fp-cuda: transpose B on the device - #286
Conversation
`transpose_b` rearranged B into the kernel's K-major layout one bit at a time on the host, and at large sizes that cost more than the multiply it fed. It is replaced by `transpose_tile_b1_kernel`: B is uploaded exactly as it stands and rearranged on the device, so the host does no bit-level work on either operand. The kernel supplies the K padding too, so that padding costs no host copy. End to end on an H200 NVL this is 1.24x at 4096 rising to 1.70x at 32768 against transposing on the host, even when the host uses a blocked transpose. See EXPERIMENTS.md, which also records why the kernel is left uncoalesced: it runs 113 us against a 1.36 ms matmul, so making it free would buy ~0.15%. `matmul_b1_kernel` is untouched and the public signatures are unchanged, so `fp`'s dispatch needs no edit. The demo built matrices from random limbs, leaving the bits past the last column set; `Matrix` requires them zero and compares limb-wise. It happened not to matter before and there is no reason to keep relying on that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EFfXYt9zvAvHWhNHnULe6r
|
Important Review skippedNo new commits to review since the last review. ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
Included review availability: Your plan includes up to 1 review per rolling hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe fp-cuda matmul path now uploads B unchanged and transposes it on the GPU. A new CUDA kernel produces the layout consumed by matmul. Host-side B transposition and padding were removed. Demo data generation now clears unused bits. ChangesCUDA B transpose
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to The device-side transpose is validated across ragged shapes and concurrent GPU dispatches, with reported kernel throughput preserved; no actionable merge-blocking risk remains beyond normal checks and review. Sequence Diagram(s)sequenceDiagram
participant GpuContext
participant CUDA_Module
participant transpose_tile_b1_kernel
participant matmul_kernel
GpuContext->>CUDA_Module: load transpose_tile_b1_kernel
GpuContext->>CUDA_Module: upload B and allocate transposed buffer
GpuContext->>transpose_tile_b1_kernel: launch B transpose
transpose_tile_b1_kernel->>matmul_kernel: provide transposed B tile layout
GpuContext->>matmul_kernel: launch matrix multiplication
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
1 participant
Closes #285.
transpose_brearranged B into the kernel's K-major layout one bit at a time on the host, and atlarge sizes that cost more than the multiply it fed.
transpose_tile_b1_kerneldoes it on thedevice instead: B is uploaded exactly as it stands, and the host does no bit-level work on either
operand. The kernel also supplies the K padding, so that padding costs no host copy either.
matmul_b1_kernelis untouched, and the public signatures are unchanged —fp's dispatch alreadypasses B row-major, so this PR touches one crate and needs no companion change.
Why on the device
The obvious alternative is to keep the rearrangement on the host but make it fast, which is cheap to
do: with Bᵀ in hand each output tile row is a contiguous limb run, so the tiling becomes a gather,
and #175 gives
fpa blockedp = 2transpose to produce Bᵀ with. That version was built andbenchmarked first. It still loses.
End to end on an H200 NVL, arms interleaved, 5 iterations, medians, each arm paying its own operand
serialization:
Reproduced on a second run (1.24 / 1.33 / 1.68 / 1.73×). The margin exceeds the host transpose alone
because the host arm also runs
pad_2dand a tiling pass, so the kernel displaces two pieces of hostwork rather than one. The host transpose is not slow in isolation — it reaches 7.4–9.4 GB/s while its
working set fits in cache — but at 32768² that set is 128 MiB and it falls to roughly 1.4 GB/s.
Since the device wins at every size the GPU path is used at, there is one entry point rather than
two. Both this and the kernel-optimization decision below are recorded in
EXPERIMENTS.md.Validation, on an H200 NVL
matmul_b1_demopasses every shape, including65×65 * 65×65.across eight shapes — ragged on each axis independently, spanning multiple K chunks (
k = 1025, 2049) and column groups (n = 129, 333).cargo test -p fp --features gpupasses, sogpu_dispatch_matches_cpuandgpu_matmul_concurrentcover the dispatched path, including 16 threads multiplying at once.4096 / 8192 / 16384 / 32768, within ~2% of before, and the PTX still emits exactly 12
wgmma.mma_async.Deliberately not done
bitreads a column of B, so consecutive threads aren_limlimbs apart;ncuat 16384² shows 9.94% DRAM throughput against 92.4% on the memory pipes.Staging a row-major tile in shared memory would fix it and
__ballot_syncwould replace the64-iteration gather. Neither is worth doing: the kernel runs 113 µs against
matmul_b1_kernel's1.36 ms inside a 77 ms end-to-end call, so making it free would buy about 0.15%.
transpose_bdominates end-to-end time at large sizes #285's second suggestion).pad_2dandinterleave_aare still serial; they are now a smaller share of the total than the transfers.
serialization, ~1.5 ms is the two kernels, and the remaining ~55 ms is H2D/D2H plus per-call device
allocation — though
pad_2dandinterleave_aalso rewrite ~33 MB on the host insidematmul_b1_inner, so not all of it is PCIe. It wants annsystimeline before anyone optimizesit, and it is the largest single item left.
Relationship to #175
Independent — they share no file, and either can merge first. #175 adds
Matrix::transpose, whichwas written for the host-transposing version of this change and remains useful on its own merits;
this PR does not use it.
Summary by CodeRabbit
Performance
Correctness
Documentation