fp-cuda: transpose B on the device

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.

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;
+}

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<u64> = (0..rows * cols.div_ceil(64))
-            .map(|_| rng.random())
-            .collect();
+        let stride = cols.div_ceil(64);
+        let mut data: Vec<u64> = (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)
     };
 

ext/crates/fp-cuda/src/lib.rs

@@ -57,6 +57,7 @@ pub struct GpuContext {
     #[allow(dead_code)]
     module: Arc<CudaModule>,
     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<u64>, 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::<u64>(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::<u64>(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<u64> {
     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<u64> {
-    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<u64> {
     if rows == nr && stride == ns {
         return src.to_vec();