diff --git a/ext/crates/fp-cuda/EXPERIMENTS.md b/ext/crates/fp-cuda/EXPERIMENTS.md index e4f5c0a565..272ceb5bd9 100644 --- a/ext/crates/fp-cuda/EXPERIMENTS.md +++ b/ext/crates/fp-cuda/EXPERIMENTS.md @@ -152,3 +152,41 @@ on this hardware, but not guaranteed by the model. Output stays bit-exact either way. The reorder costs ~0.96% at 4096 and ~0.73% at 8192 (every post-change run below the lowest pre-change run, so the loss is real), and is at noise level at 16384 and 32768, the sizes the kernel is actually used at. Paid. + +## B is transposed on the device, not the host (2026-08-16, H200 NVL) + +**Chosen:** B is uploaded untransposed and rearranged by `transpose_tile_b1_kernel`. There is no +host-side path. + +The alternative was to take Bᵀ from the caller and gather the tiles on the host — cheap, since with +Bᵀ in hand each tile row is a contiguous limb run, and `fp` can produce Bᵀ quickly with a blocked +`p = 2` transpose. It still loses, end to end, arms interleaved, 5 iterations, medians, each arm +paying its own operand serialization: + +| n | host total | device total | speedup | +|---|------------|--------------|---------| +| 4096 | 2.59 ms | 2.09 ms | 1.24x | +| 8192 | 16.32 ms | 11.18 ms | 1.46x | +| 16384 | 127.24 ms | 76.93 ms | 1.65x | +| 32768 | 459.54 ms | 271.18 ms | 1.70x | + +Reproduced on a second run (1.24 / 1.33 / 1.68 / 1.73x). The margin exceeds the host transpose alone +because the host arm also runs `pad_2d` and a tiling pass, so the kernel displaces two pieces of +host work rather than one. + +The host transpose is not slow in isolation — it reaches 7.4-9.4 GB/s while the working set fits in +cache — but at 32768 that set is 128 MiB and it falls to roughly 1.4 GB/s. Measurements taken only +up to 8192 make the host arm look far better than it is. + +## The transpose kernel is not worth optimizing (2026-08-16, H200 NVL) + +**Chosen:** leave `transpose_tile_b1_kernel` as written. + +Its loads are uncoalesced — thread `bit` reads a column of B, so consecutive threads are `n_lim` +limbs apart — and `ncu` at 16384 confirms it: 9.94% DRAM throughput against 92.4% on the memory +pipes. Staging a row-major tile in shared memory would fix that, and `__ballot_sync` would replace +the 64-iteration gather. + +Neither is worth doing. The kernel runs 113 us against `matmul_b1_kernel`'s 1.36 ms, inside a 77 ms +end-to-end call. Making it infinitely fast would buy about 0.15%. Of that 77 ms, roughly 55 ms is +H2D/D2H and device allocation, which is where the time actually is. diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index e0e3174941..5d75bb0398 100644 --- a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -406,3 +406,44 @@ extern "C" __global__ void matmul_b1_kernel( // Drain the last outstanding output store before the CTA exits. if (t == 0) tma_store_wait(); } + +// Build the K-major tiles the matmul consumes, reading B in its natural row-major layout. +// +// Tile (k_chunk, column group) row `lg*64 + jj`, limb `kl`, bit `bit` is bit `jj` of +// B[k_chunk*TK + kl*64 + bit][cg*NG + lg]. One block owns one (k_chunk, cg, lg, kl) quadruple and +// so one 64-square bit block: thread `bit` loads that block's row into shared memory, then thread +// `jj` gathers bit `jj` out of all 64 rows. The gather reads one shared slot at a time across the +// whole block, so every read is a broadcast rather than a bank conflict. +// +// Both the loads (a column of B, stride n_lim) and the stores (stride KL) are strided. +extern "C" __global__ void transpose_tile_b1_kernel( + const unsigned long long* __restrict__ b, // k_padded x n_lim, row-major + unsigned long long* __restrict__ out, // k_chunks x n_groups x (NB*KL) + int n_lim, // limbs per row of B + int k_rows, // rows of B actually uploaded (the unpadded k) + int n_groups) // column groups of NG limbs +{ + __shared__ unsigned long long sB[64]; + + const int kl = blockIdx.x % KL; + const int lg = (blockIdx.x / KL) % (NB / 64); + const int cg = blockIdx.y; + const int kk = blockIdx.z; + const int t = threadIdx.x; // 0..63 + + const int limb = cg * (NB / 64) + lg; + const int row = kk * TK + kl * 64 + t; + // Column groups past the operand, and K rows past the end of B, contribute zeros — so B is + // uploaded unpadded and the K padding costs no host copy. + sB[t] = (limb < n_lim && row < k_rows) ? b[(long long)row * n_lim + limb] : 0ULL; + __syncthreads(); + + unsigned long long val = 0; + for (int bit = 0; bit < 64; ++bit) { + val |= ((sB[bit] >> t) & 1ULL) << bit; + } + + const long long tile = (long long)NB * KL; + const long long base = ((long long)kk * n_groups + cg) * tile; + out[base + (long long)(lg * 64 + t) * KL + kl] = val; +} diff --git a/ext/crates/fp-cuda/examples/matmul_b1_demo.rs b/ext/crates/fp-cuda/examples/matmul_b1_demo.rs index ee553715d3..f744f8c350 100644 --- a/ext/crates/fp-cuda/examples/matmul_b1_demo.rs +++ b/ext/crates/fp-cuda/examples/matmul_b1_demo.rs @@ -18,9 +18,17 @@ fn main() -> anyhow::Result<()> { let mut rng = rand::rng(); let mut make = |rows: usize, cols: usize| { - let data: Vec = (0..rows * cols.div_ceil(64)) - .map(|_| rng.random()) - .collect(); + let stride = cols.div_ceil(64); + let mut data: Vec = (0..rows * stride).map(|_| rng.random()).collect(); + // `Matrix` requires the bits past the last column to be zero, and random limbs do not + // respect that. Comparisons are limb-wise, so leaving them set makes two matrices that + // agree on every entry compare unequal. + if !cols.is_multiple_of(64) { + let mask = (1u64 << (cols % 64)) - 1; + for row in 0..rows { + data[row * stride + stride - 1] &= mask; + } + } Matrix::from_data(TWO, rows, cols, data) }; diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 22db5523e2..bb9a1a1e47 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -57,6 +57,7 @@ pub struct GpuContext { #[allow(dead_code)] module: Arc, kernel: CudaFunction, + transpose_kernel: CudaFunction, } impl GpuContext { @@ -69,11 +70,13 @@ impl GpuContext { let ptx = Ptx::from_src(String::from_utf8(PTX_IMAGE.to_vec())?); let module = ctx.load_module(ptx)?; let kernel = module.load_function("matmul_b1_kernel")?; + let transpose_kernel = module.load_function("transpose_tile_b1_kernel")?; Ok(Self { ctx, streams: Mutex::new(HashMap::new()), module, kernel, + transpose_kernel, }) } @@ -129,6 +132,9 @@ impl GpuContext { /// - `a`: the `m`×`k` left operand, `m * k.div_ceil(64)` limbs. /// - `b`: the `k`×`n` right operand, `k * n.div_ceil(64)` limbs. /// +/// The kernel wants B K-major; `transpose_tile_b1_kernel` rearranges it on the device, so B is +/// uploaded exactly as it stands and the host does no bit-level work on either operand. +/// /// Returns C = A·B as `m * n.div_ceil(64)` limbs in the same layout, ready to hand to /// `fp::Matrix::from_data`. pub fn matmul_b1_raw( @@ -173,7 +179,8 @@ fn matmul_b1_inner( time_iters: usize, ) -> anyhow::Result<(Vec, f64)> { let n_lim = n.div_ceil(64); - assert_eq!(a.len(), m * k.div_ceil(64), "A limb count mismatch"); + let k_lim = k.div_ceil(64); + assert_eq!(a.len(), m * k_lim, "A limb count mismatch"); assert_eq!(b.len(), k * n_lim, "B limb count mismatch"); let k_padded = k.next_multiple_of(TILE_K); @@ -189,14 +196,34 @@ fn matmul_b1_inner( let stream = gpu.stream(); - let a_padded = pad_2d(a, m, k.div_ceil(64), m_padded, k_padded / 64); - let b_padded = pad_2d(b, k, n_lim, k_padded, n_lim); - + let a_padded = pad_2d(a, m, k_lim, m_padded, k_padded / 64); let a_interleaved = interleave_a(&a_padded, m_padded, k_padded); - let bt = transpose_b(&b_padded, k_padded, n_lim); - let a_dev = stream.clone_htod(&a_interleaved)?; - let bt_dev = stream.clone_htod(&bt)?; + + // B goes up unpadded; the kernel reads rows past `k` as zeros, so the K padding costs no host + // copy. + let b_dev = stream.clone_htod(b)?; + let bt_dev = stream.alloc_zeros::(k_chunks * n_groups * (NG as usize * 64) * KL)?; + { + let cfg = LaunchConfig { + grid_dim: ((KL * NG as usize) as u32, n_groups as u32, k_chunks as u32), + block_dim: (64, 1, 1), + shared_mem_bytes: 0, + }; + let mut lb = stream.launch_builder(&gpu.transpose_kernel); + let (n_lim_i, k_i, n_groups_i) = (n_lim as i32, k as i32, n_groups as i32); + lb.arg(&b_dev) + .arg(&bt_dev) + .arg(&n_lim_i) + .arg(&k_i) + .arg(&n_groups_i); + // SAFETY: the five pushed arguments match `transpose_tile_b1_kernel`'s parameter list in + // order and type; `b_dev` holds `k * n_lim` limbs and the kernel indexes it only where + // `row < k` and `limb < n_lim`; `bt_dev` is exactly the tile count the grid covers; both + // buffers outlive the launch, their guards being held until the final synchronize. + unsafe { lb.launch(cfg) }?; + } + let c_dev = stream.alloc_zeros::(m_padded * n_padded_lim)?; // Raw device addresses for the TMA descriptors. The returned guards keep the @@ -379,51 +406,6 @@ fn interleave_a(a: &[u64], m: usize, k: usize) -> Vec { out } -/// Pre-transpose B into plain row-major K-major tiles for TMA 128B swizzle. -/// -/// Each (k_chunk, column group) tile is NB = NG*64 rows (= the NG*64 output columns of the group) × -/// KL u64s (= TILE_K K bits); the consumer feeds it to MSTRIPS m64n128 wgmmas that share it. -/// Operand row `lg*64 + jj` is output column `cg*NG*64 + lg*64 + jj`; element `[..][kl] bit` is bit -/// `jj` of `B[k_chunk*TILE_K + kl*64 + bit][cg*NG + lg]`. -/// -/// Groups whose limb runs past `n_lim` are left zero-padded. Output is row-major; the TMA applies -/// the swizzle on load. -fn transpose_b(b: &[u64], k: usize, n_lim: usize) -> Vec { - let k_chunks = k / TILE_K; - let ng = NG as usize; - let n_groups = n_lim.div_ceil(ng); - let tile = ng * 64 * KL; // NB rows × KL u64 - let mut out = vec![0u64; k_chunks * n_groups * tile]; - let mut buf = [0u64; TILE_K]; - - for kk in 0..k_chunks { - for cg in 0..n_groups { - let base = (kk * n_groups + cg) * tile; - for lg in 0..ng { - let limb = cg * ng + lg; - if limb >= n_lim { - continue; // padded column group → leave zeros - } - for (i, slot) in buf.iter_mut().enumerate() { - let br = kk * TILE_K + i; - *slot = if br < k { b[br * n_lim + limb] } else { 0 }; - } - for jj in 0..64usize { - let j = lg * 64 + jj; // operand row within the NB-col tile - for kl in 0..KL { - let mut val: u64 = 0; - for bit in 0..64usize { - val |= ((buf[kl * 64 + bit] >> jj) & 1) << bit; - } - out[base + j * KL + kl] = val; - } - } - } - } - } - out -} - fn pad_2d(src: &[u64], rows: usize, stride: usize, nr: usize, ns: usize) -> Vec { if rows == nr && stride == ns { return src.to_vec();