Follow-up from #273, where this is listed under "Not yet done" in the crate README but was never tracked.
transpose_b (ext/crates/fp-cuda/src/lib.rs) rearranges B into the K-major tile layout the TMA descriptors expect. It does so one bit at a time, single-threaded:
for jj in 0..64usize {
let j = lg * 64 + jj;
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;
}
}Every bit of B passes through that inner loop, so the work is ~k * n single-bit operations — roughly 1.07e9 at 32768³. The kernel itself runs in ~7.35 ms at that size, so the host-side transpose is the dominant cost of an end-to-end multiply and is the main reason cargo bench reports ≤30 TOPS against the kernel's ~9,600.
Two independent improvements, either of which helps on its own:
- Blocked 64×64 bit transpose. The delta-swap / recursive-halving trick does a 64×64 bit-matrix transpose in ~6 masked shift-and-XOR rounds over 64 words, instead of 4096 bit extractions. This is the bulk of the win.
- Parallelise over
(kk, cg). The outer loops write to disjoint base offsets of out, so the iteration space is already embarrassingly parallel and rayon is a workspace dependency.
Correctness is easy to hold onto here: the transpose is a pure function, so a new implementation can be checked against the current one directly, and the existing CPU-vs-GPU bit-exactness checks in matmul_b1_demo and the kernel-only bench cover the integrated path.
Related: keeping operands resident on the device across step_resolution's successive multiplications (the other README "Not yet done" item) would avoid re-running this per product, and the two interact.
Follow-up from #273, where this is listed under "Not yet done" in the crate README but was never tracked.
transpose_b(ext/crates/fp-cuda/src/lib.rs) rearranges B into the K-major tile layout the TMA descriptors expect. It does so one bit at a time, single-threaded:Every bit of B passes through that inner loop, so the work is ~
k * nsingle-bit operations — roughly 1.07e9 at 32768³. The kernel itself runs in ~7.35 ms at that size, so the host-side transpose is the dominant cost of an end-to-end multiply and is the main reasoncargo benchreports ≤30 TOPS against the kernel's ~9,600.Two independent improvements, either of which helps on its own:
(kk, cg). The outer loops write to disjointbaseoffsets ofout, so the iteration space is already embarrassingly parallel and rayon is a workspace dependency.Correctness is easy to hold onto here: the transpose is a pure function, so a new implementation can be checked against the current one directly, and the existing CPU-vs-GPU bit-exactness checks in
matmul_b1_demoand the kernel-only bench cover the integrated path.Related: keeping operands resident on the device across
step_resolution's successive multiplications (the other README "Not yet done" item) would avoid re-running this per product, and the two interact.