fp-cuda: Hopper wgmma.b1 device-resident F₂ GEMM backend - #273
Open
JoeyBF wants to merge 18 commits into
Open
Conversation
Prototype CUDA kernel and the Phase 2 wgmma.b1 core: both operands pre-arranged on the host as K-major tiles, one m64n128k256 binary MMA. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B1XztNMycTtSC9SQtpNYjR
…zation Steps (a)-(d): hoist the A TMA load out of the column-group loop, double-buffer A in the K pipeline, load B via TMA, and split into producer/consumer warpgroups. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B1XztNMycTtSC9SQtpNYjR
Both b1 operands are K-major; move them to CU_TENSOR_MAP_SWIZZLE_128B so wgmma operand reads avoid bank conflicts. The TMA applies the swizzle on load, so the host now emits plain row-major K-major tiles (the hand-rolled cm() interleave is gone) and the wgmma matrix descriptors carry the matching layout bits (layout=1, LBO=16B, SBO=1024B), derived from CUTLASS make_gmma_desc<Major::K> / LayoutType::B128. The SMEM K-tile grows to 1024 bits (one full 128B K-major swizzle atom = 4 k256 sub-chunks) and moves to dynamic shared memory (~82KB, opt-in via CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES). Per-stage wgmmas now run behind a single commit_group/wait_group and accumulate popcounts in-hardware (scale-D=1) into one resident accumulator per column group, replacing the previous one-wgmma-per-commit/wait serialization. Compile-verified only (nvcc + rustc); not yet validated on H100. Validate with a 64x256x64 identity product first, then matmul_b1_demo, then the bench. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the four m64n64k256 wgmmas per k-step (one per output limb) with a single m64n256k256 covering all NG=4 limbs at once. Binary wgmma is k256-only, so N is the throughput lever; n256 is the max. Same accumulator register count (128 s32/thread) and same SMEM as the 4x n64 version, but 1/4 the wgmma instructions, B descriptors, and fence/commit churn. B is now arranged as one contiguous 256-column tile per CTA (host transpose_b packs 4 limbs side by side; the B TMA box is 256 rows tall) instead of four separate 64-column tiles, and is zero-padded to whole 256-column groups. A and the K tiling are unchanged. The output bit-pack splits the single acc[128] into the 4 output limbs: the m64n256 fragment is the m64n64 layout tiled along N, so register group gi in 0..32 maps to columns [gi*8, gi*8+8). Compile-verified (PTX emits m64n256k256); fragment->limb bit-pack still needs the 64x256x64 identity check on H100. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The producer warpgroup needs few registers; the consumer holds the 128-reg m64n256 accumulator. Issue setmaxnreg.dec(40) in the producer and setmaxnreg.inc(216) in the consumer so the consumer claims the producer's surplus. 128*(40+216)=32768 regs/CTA, leaving room for 2 CTAs/SM. Both are warpgroup-aligned and executed by all 128 threads of their warpgroup (the wg branch splits at warpgroup granularity). Compile-verified; PTX emits the dec/inc. Also fixes a stale m64n64k256 mention in the lib doc comment. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Bump the producer-consumer pipeline depth from 2 to 3 so the producer can run
two K-chunks ahead, hiding more TMA latency. SMEM grows to ~122 KB/CTA (one
CTA/SM vs two at STAGES=2); the host smem_bytes/cuFuncSetAttribute track it
automatically. STAGES is the latency-vs-occupancy knob to sweep on hardware.
Phase-array inits made depth-agnostic ({0}).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the per-thread global stores with a single cp.async.bulk.tensor.2d.global.shared::cta per CTA. sC is now packed row-major ([row][limb]) so the 64×NG tile is contiguous for the store; thread 0 issues it after a __syncthreads + fence.proxy.async (making the atomicXor writes visible to the async proxy), then cp.async.bulk.commit_group/wait_group, then a final __syncthreads keeps sC alive until the store lands. The kernel now takes a C tensor map instead of a raw pointer (nlim/C* params dropped). C is padded on the host to whole NG-limb column groups (n_padded_lim = n_groups*NG) so every stored tile is complete; padded columns carry zeros from the zero-padded B and are trimmed on readback. README roadmap updated to reflect Phases 4-7. Compile-verified (PTX emits the bulk store + fence.proxy.async); store path and n256 bit-pack still need the H100 identity check. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the cuda-oxide `cuda-core` dependency with `cudarc`. cuda-core pulled in nightly via an incidental `#![feature(f16)]` (a DeviceCopy impl for the unstable f16 primitive we never use); cudarc is stable Rust, mainstream/maintained, and dynamically loads the driver at runtime — so the crate builds with no CUDA present and the whole crate (lib + examples + benches) now compiles AND links locally without libcuda (previously examples failed at link with -lcuda). Host glue rewritten against cudarc: CudaContext/load_module(Ptx)/load_function, clone_htod/alloc_zeros/clone_dtoh, device_ptr for the TMA descriptor addresses, CudaFunction::set_attribute for the >48 KB dynamic-SMEM opt-in, and the typed launch builder. The three CUtensorMaps are passed by value as grid-constant args via a #[repr(transparent)] TmaArg: DeviceRepr wrapper. cuTensorMapEncodeTiled is called through cudarc::driver::sys. Kernel (.cu/PTX) and build.rs are unchanged. cudarc features: driver + nvrtc (for Ptx) + cuda-12080 + dynamic-loading; the version only selects pre-generated bindings, the real driver (12.8+/13.x) is resolved at runtime. Also delete examples/bench_breakdown.rs: it was stale (called the pre-TMA 7-pointer kernel signature with the old cm() transpose) and the only remaining cuda-core user; bench_kernel_only supersedes it. Compile + link verified locally (lib, examples, benches). Stable-toolchain build to be confirmed on the server (dev box only has Nix nightly). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The >16384 throughput cliff is L2-residency bound on B: with the old (grid_x=n_groups, grid_y=m_tiles) launch (N fastest), every B column-panel is re-touched only once per full sweep of B, so its reuse distance is the whole matrix and it spills from L2 once K*N/8 > ~50 MB, getting re-streamed from HBM per M-tile. Convert to a persistent 1-D grid of ~SM-count CTAs that sweep all output tiles in a grouped-along-M order (bi fastest within a GROUP_M-row band, bj slowest). This shortens each B-panel's reuse distance to <= GROUP_M, cutting B's HBM re-reads by a factor of GROUP_M while A still streams once. Kernel: take n_groups as an arg (no longer gridDim.x), add GROUP_M=8, wrap the per-tile body in a persistent `for (tile = blockIdx.x; ...)` loop, and re-init the mbarriers per tile (fresh phase-0 bookkeeping) so barrier parity isn't carried across iterations. setmaxnreg stays a one-time per-warpgroup action before the loop. Per-tile compute is byte-identical to Phase 7. Host: query CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, launch (num_ctas = min(SMs, total_tiles), 1, 1), and pass n_groups. TMA descriptors, padding, and readback are unchanged. GROUP_M is a tuning knob (8/16/32) to sweep on the H100. Code-only on the dev box; validate on the server. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Phase 8 shortened each B-panel's reuse distance via rasterization; this shares the remaining HBM read of a B-panel across CTAs. CLUSTER (=2) CTAs along M form a thread-block cluster: each computes a different M-tile (its own n256 accumulator) but receives the same B-panel via one multicast HBM read (cp.async.bulk.tensor.2d...multicast::cluster), cutting B's HBM traffic by a further factor of CLUSTER on top of Phase 8's GROUP_M. (A second per-CTA accumulator can't fit the 256-reg budget, so cross-M B-reuse must be cross-CTA, i.e. clusters — not single-CTA SMEM reuse.) Kernel: __cluster_dims__(CLUSTER,1,1); the schedule walks M-super-rows of CLUSTER tiles (bi = sbi*CLUSTER + rank, shared bj). A is loaded per-CTA; B is multicast by rank 0 with an all-ranks mask into every member's sB and counted against every member's full barrier. The empty barrier is cluster-wide (init count CLUSTER; each consumer arrives on every member via mapa). Pipeline barriers are initialized once and flow continuously across tiles (single qidx/p) rather than re-init per tile, which would race the cross-CTA arrivals. One barrier.cluster after init. Mirrors pranjalssh/fast.cu matmul_9. Host: pad m_tiles to a multiple of CLUSTER (every rank gets a valid M-tile), round the persistent grid down to a whole number of clusters. CLUSTER must match the kernel constant (like STAGES). GROUP_M and CLUSTER are tuning knobs to sweep on the H100. Multicast/cluster ordering is the highest-risk, hardware-unverifiable part — see HANDOFF.md. Code-only on the dev box; validate on the server. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Diagnosis on H200: the kernel was L2-refill-bandwidth bound, not compute bound. Microbenchmarks put the single-warpgroup wgmma.b1 ceiling at ~12,468 TOPS while sustained L2->SMEM refill tops out at ~8 TB/s (full tensor rate needs ~14.6). Each CTA computed a 64x256 tile, so it reloaded 8KB A + 32KB B per k-chunk -- B was 80% of the traffic because the tile was 4x wider than tall, and the tensor core starved waiting for it. Fix: register-block the output. Each CTA now computes a TM x NB block as MSTRIPS m64n128 wgmma strips (wgmma_n128, acc[MSTRIPS][64]) that all reuse one loaded B sub-tile, so a single L2->SMEM read of B feeds every strip. Winning config MSTRIPS=3 (192x128 block), NB=128, STAGES=4, GROUP_M=16, CLUSTER=2 cuts refill bytes/MAC by 33%. Also validates the previously-unvalidated Phase 8-9 batch on hardware (bit-exact) and lands the STAGES 3->4 / GROUP_M 8->16 tuning. Kernel-only, bit-exact (32768 cube): 8,585 -> 9,344 TOPS (+8.9%); 16384 7,636 -> 8,301; now ~75% of the tensor-core ceiling. Constants MSTRIPS in the kernel and TILE_M(=64*MSTRIPS)/NG in src/lib.rs must match. NB must be a multiple of 128 (NG even) or the C-store TMA box breaks 16B alignment; MSTRIPS=4 would exceed the 255-reg/thread cap. Adds examples/tune.rs, a fast no-CPU-check sweep harness. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The per-tile epilogue drained its C store inline (cp.async.bulk.wait_group 0 between two __syncthreads), stalling the whole CTA for the store latency every output tile -- a bubble that grows as a fraction of tile time at smaller K. Double-buffer sC and defer the wait: each tile packs into sC[titer&1], issues its store, then does cp.async.bulk.wait_group.read 1, which blocks only until every store but the newest has finished *reading* its SMEM source. So tile T's store drains during tile T+1's compute, and the buffer freed (two tiles back) is safe to reuse. A final wait_group 0 after the persistent loop drains the last store before the CTA exits. Kernel-only, bit-exact: 32768 9,344 -> 9,426; 16384 8,301 -> 8,420; 8192 6,425 -> 6,591; 4096 3,835 -> 3,948 (smaller K gains more, as the epilogue is a larger share of tile time there). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
wgmma.fence.sync.aligned is a warpgroup-wide (128-thread) sync; the consumer was issuing two per k-chunk (one before the wgmma group, one after the wait), i.e. 64 warpgroup syncs per output tile. But the fence is only needed to order non-wgmma writes to the accumulators before a wgmma reads them -- in steady state the accumulators are touched only by wgmma, so one fence before the whole K-loop (ordering the accumulator zeroing) suffices, matching CUTLASS's warpgroup_arrive-once pattern. The per-chunk wgmma.wait_group 0 still makes results readable for the epilogue. Kernel-only, bit-exact: 32768 9,426 -> 9,605; 16384 8,420 -> 8,618; 8192 6,591 -> 6,744. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Launch (occupancy × SM-count) CTAs instead of hard-coding SM-count, so the persistent grid exactly fills the machine for whatever the resource footprint allows. No change at the shipping config (occ=1/SM), but it's the correct, self-adapting launch and it drove the 2-CTA/SM experiment. That experiment is a documented dead end (H200, 2026-07-07): 2 CTAs/SM needs the compiled register count <=128/thread (2*256*128 = the 64K reg file), but the resident accumulator that gives the kernel its arithmetic intensity is 192 regs/thread at MSTRIPS=3. The only way to reach occ=2 is MSTRIPS=1 (64-reg acc), which collapses AI and drops 16384 from ~8,600 to ~5,500 TOPS. High AI (big accumulator) and high occupancy compete for the same register file and AI wins -- so 1 CTA/SM with the largest accumulator under the 255-reg cap is optimal. This confirms on-device that the ~10-13k TOPS bandwidth wall cannot be moved by raising occupancy. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds an opt-in Hopper CUDA backend for packed F₂ matrix multiplication, integrates GPU dispatch with CPU fallback, provides PTX compilation and CUDA environment setup, and adds correctness tests, benchmarks, examples, and performance documentation. ChangesCUDA backend integration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Matrix
participant CUDA dispatch
participant fp-cuda
participant Hopper kernel
Matrix->>CUDA dispatch: multiply F₂ matrices
CUDA dispatch->>fp-cuda: convert limbs and call matmul_b1_raw
fp-cuda->>Hopper kernel: launch TMA/wgmma kernel
Hopper kernel-->>fp-cuda: return packed output limbs
fp-cuda-->>CUDA dispatch: reconstruct result Matrix
CUDA dispatch-->>Matrix: return GPU result or CPU fallback
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
Wire the Hopper GPU backend into `fp` behind an optional `gpu` feature and keep the workspace CI green: `--workspace` builds `fp-cuda`, whose build.rs emits a stub PTX when nvcc is absent (every ubuntu-latest runner) instead of panicking. fp-cuda's library API is fp-agnostic (raw limbs), breaking the dependency cycle so `fp` can dispatch large p=2 products to the device. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B1XztNMycTtSC9SQtpNYjR
JoeyBF
marked this pull request as ready for review
July 18, 2026 02:57
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ext/crates/fp-cuda/build.rs`:
- Around line 61-68: Update the nvcc command-status error handling in the build
script so the stub PTX fallback occurs only when the error has
ErrorKind::NotFound. Propagate or fail fast for all other errors instead of
emitting the unavailable-backend warning and writing STUB_PTX.
In `@ext/crates/fp-cuda/examples/bench_kernel.rs`:
- Around line 14-15: Update the benchmark output in main to remove obsolete
claims about scalar B transposition, wgmma starvation, and pending optimization
phases. Report only measured end-to-end overhead or replace the diagnosis with
the current verified bottleneck, including the corresponding messages in the
additional affected section.
In `@ext/crates/fp-cuda/examples/bench_shapes.rs`:
- Around line 23-25: Update the benchmark setup around GpuContext::new and l2_mb
so the L2 capacity is obtained from the selected device’s attributes or supplied
through explicit configuration; do not hardcode 50 MB as a device fact. Ensure
the output labels any configured value as an assumption and uses the actual
capacity for fit/spill calculations.
In `@ext/crates/fp-cuda/README.md`:
- Around line 195-196: Use the consistent Cargo feature name “gpu” in both
documentation references: update the optional feature wording in
ext/crates/fp-cuda/README.md lines 195-196 and the fp feature wording in
ext/Cargo.toml line 79 from “cuda” to “gpu”.
- Around line 26-28: Update the runtime GPU requirement in the fp-cuda README to
remove sm_100 and state that Hopper (sm_90/sm_90a) is required. Keep the
existing explanation about PTX instructions and pre-Hopper incompatibility
unchanged.
In `@ext/crates/fp/src/blas/cuda.rs`:
- Around line 3-4: Replace the stale cuda feature name with gpu in the
documentation for the dispatch module at ext/crates/fp/src/blas/cuda.rs lines
3-4 and the integration test at ext/crates/fp/tests/cuda_dispatch.rs lines 1-2;
update only these feature references and preserve the surrounding dispatch
documentation.
In `@ext/crates/fp/tests/cuda_dispatch.rs`:
- Line 22: Add non-tile-aligned matrix shapes to the dimension list in the CUDA
dispatch test loop, such as (2049, 2051, 2053), while retaining the existing
aligned cases. Ensure the added cases exercise edge masking, partial limbs, and
raster-tail handling.
- Around line 26-29: Update the CUDA dispatch test around the dispatched and
reference computations to require a successfully initialized, usable GpuContext
before executing. Obtain the expected value through a direct raw GPU
kernel/result rather than fast_mul_concurrent, and fail or skip explicitly when
CUDA context or launch setup is unavailable so fallback execution cannot make
the test pass.
In `@ext/flake.nix`:
- Around line 47-75: Update the flake’s devShell definitions so
devShells.default uses only commonPackages and does not reference cudaPackages
or CUDA-specific shell setup. Add a separate devShells.cuda containing
commonPackages ++ cudaPackages and the existing CUDA/libclang shellHook
environment variables, preserving the current CUDA development behavior for
opt-in users.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: fdbf92f4-e29c-48f5-9f49-183cc4f6dcdb
📒 Files selected for processing (20)
ext/Cargo.tomlext/crates/fp-cuda/Cargo.tomlext/crates/fp-cuda/README.mdext/crates/fp-cuda/benches/matmul_b1.rsext/crates/fp-cuda/build.rsext/crates/fp-cuda/cuda_kernels/matmul_b1.cuext/crates/fp-cuda/examples/bench_giant.rsext/crates/fp-cuda/examples/bench_kernel.rsext/crates/fp-cuda/examples/bench_kernel_only.rsext/crates/fp-cuda/examples/bench_shapes.rsext/crates/fp-cuda/examples/common/mod.rsext/crates/fp-cuda/examples/matmul_b1_demo.rsext/crates/fp-cuda/examples/prof_sizes.rsext/crates/fp-cuda/examples/tune.rsext/crates/fp-cuda/src/lib.rsext/crates/fp/Cargo.tomlext/crates/fp/src/blas/cuda.rsext/crates/fp/src/blas/mod.rsext/crates/fp/tests/cuda_dispatch.rsext/flake.nix
Comment on lines
+23
to
+25
| let gpu = GpuContext::new(0)?; | ||
| let l2_mb = 50.0; // H100 NVL L2 | ||
| println!("GPU L2 ~= {l2_mb} MB. B in L2 (bytes = K*N/8) governs cross-M-tile reuse.\n"); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Do not report hardcoded L2 classifications as device facts.
The context can target any supported device, but every fit/spill result assumes a 50 MB H100 NVL cache. Read the capacity from device attributes or require an explicit configuration value and label it as an assumption.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ext/crates/fp-cuda/examples/bench_shapes.rs` around lines 23 - 25, Update the
benchmark setup around GpuContext::new and l2_mb so the L2 capacity is obtained
from the selected device’s attributes or supplied through explicit
configuration; do not hardcode 50 MB as a device fact. Ensure the output labels
any configured value as an assumption and uses the actual capacity for fit/spill
calculations.
Comment on lines
+47
to
+75
| # CUDA toolkit is only needed for `cargo build -p fp-cuda` (the Hopper | ||
| # wgmma.b1 backend). Kept out of `commonPackages` to avoid pulling | ||
| # multi-GB CUDA into the `apps.test` closure used by CI. | ||
| cudaPackages = [ | ||
| pkgs.cudaPackages.cudatoolkit | ||
| # cuda-oxide's `cuda-bindings` crate runs `bindgen` against cuda.h, | ||
| # which needs libclang at build time. | ||
| pkgs.llvmPackages.libclang.lib | ||
| ]; | ||
| in { | ||
| devShells.default = pkgs.mkShell { | ||
| packages = commonPackages; | ||
| packages = commonPackages ++ cudaPackages; | ||
| shellHook = '' | ||
| export RUST_LOG=info | ||
|
|
||
| # CUDA: make nvcc find headers + libs, and satisfy cuda-oxide's | ||
| # cuda-bindings build.rs (which reads CUDA_TOOLKIT_PATH, defaulting | ||
| # to /usr/local/cuda otherwise). | ||
| export CUDA_PATH=${pkgs.cudaPackages.cudatoolkit} | ||
| export CUDA_TOOLKIT_PATH=${pkgs.cudaPackages.cudatoolkit} | ||
| export CPATH="$CUDA_PATH/include''${CPATH:+:$CPATH}" | ||
| export LIBRARY_PATH="$CUDA_PATH/lib64''${LIBRARY_PATH:+:$LIBRARY_PATH}" | ||
|
|
||
| # libclang for bindgen (used by cuda-oxide's cuda-bindings crate). | ||
| # libclang loaded as a .so doesn't pick up the wrapped clang's | ||
| # auto-discovered libc/gcc include paths the way the clang binary | ||
| # does, so we feed them via BINDGEN_EXTRA_CLANG_ARGS. | ||
| export LIBCLANG_PATH=${pkgs.llvmPackages.libclang.lib}/lib | ||
| export BINDGEN_EXTRA_CLANG_ARGS="$(< ${pkgs.stdenv.cc}/nix-support/libc-crt1-cflags) $(< ${pkgs.stdenv.cc}/nix-support/libc-cflags) $(< ${pkgs.stdenv.cc}/nix-support/cc-cflags) $(< ${pkgs.stdenv.cc}/nix-support/libcxx-cxxflags 2>/dev/null || true)" |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Keep CUDA out of the default development shell.
devShells.default unconditionally includes cudaPackages, making every contributor fetch the large unfree CUDA closure even though the backend is opt-in. Keep default on commonPackages and expose these packages and environment variables through a separate devShells.cuda.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ext/flake.nix` around lines 47 - 75, Update the flake’s devShell definitions
so devShells.default uses only commonPackages and does not reference
cudaPackages or CUDA-specific shell setup. Add a separate devShells.cuda
containing commonPackages ++ cudaPackages and the existing CUDA/libclang
shellHook environment variables, preserving the current CUDA development
behavior for opt-in users.
- build.rs: fall back to the stub PTX only when nvcc is genuinely absent (io::ErrorKind::NotFound); fail fast on other spawn errors (permissions, broken exec) instead of silently building the stub. - flake.nix: move the CUDA toolkit out of the default dev shell into a dedicated `devShells.gpu` (`nix develop .#gpu`), matching the sibling GPU PR, so contributors/CI don't fetch the multi-GB unfree CUDA closure for the opt-in backend. Drop the obsolete cuda-oxide/libclang/bindgen env (the crate uses cudarc, which dlopens libcuda). - README/Cargo.toml/cuda.rs: call the Cargo feature `gpu` consistently; note the kernel is sm_90a-only (architecture-specific, not forward-compatible — no Blackwell/sm_100). - cuda_dispatch test: add non-tile-aligned shapes (2049x2051x2053, 3000x2112x4097) to exercise edge masks, partial limbs, and raster tails. - bench examples: drop obsolete pre-optimization diagnosis prose; label the bench_shapes L2 size as an assumption rather than a device fact. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B1XztNMycTtSC9SQtpNYjR
The matmul dispatch serialized every submission behind one Mutex<GpuContext> on the context's single default stream, so concurrent `&Matrix * &Matrix` from rayon workers ran one at a time. Add GpuContext::stream() (a lazily-created per-thread CudaStream) and drop the mutex: GpuContext is Send+Sync (its cudarc handles are), matmul_b1_inner allocates its device buffers per call and launches non-cooperatively, so concurrent matmuls run on independent streams — overlapping transfers with compute — with no shared state to guard. context() now shares a &'static GpuContext; stream() falls back to the context default stream if creation fails, so a poisoned context degrades to the CPU path instead of panicking. This is the foundation the row-reduce concurrency (blas3 PR, stacked on this one) builds on: it inherits stream() + the lock-free context() and only swaps its own reduce methods onto per-thread streams. Validated: gpu_matmul_concurrent (new) runs 16 threads x 6 concurrent GPU matmuls, each bit-identical to the CPU kernel; existing dispatch test unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
ext/crates/fp-cuda/README.md (1)
210-212: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the stale Nix-flake roadmap item.
ext/flake.nixnow definesdevShells.gpuwith the CUDA toolkit andnvccsetup, so this remaining item contradicts the current opt-in shell. Remove it or rewrite it to describe a still-missing parent-flake integration.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ext/crates/fp-cuda/README.md` around lines 210 - 212, Remove the stale “Extend the parent Nix flake to provide nvcc when the user opts in” roadmap item from the README, since the opt-in GPU shell already supplies the CUDA toolkit and nvcc. Leave the unrelated device-residency item unchanged.ext/crates/fp-cuda/examples/bench_kernel.rs (2)
57-69: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winMeasure the complete host packing path, or rename this metric.
This block times only
Matrix::to_bytes, whilecommon::to_limbsadditionally converts the bytes intou64limbs beforematmul_b1launches. Therefore this is not the full host-packing overhead used by the GPU path.Suggested fix
- // Estimate host overhead: time just serialization + padding + // Measure host packing used by matmul_b1 let t_host = Instant::now(); - let _a_ser = { - let mut v = Vec::new(); - a.to_bytes(&mut v).unwrap(); - v - }; - let _b_ser = { - let mut v = Vec::new(); - b.to_bytes(&mut v).unwrap(); - v - }; - let host_ser_ms = t_host.elapsed().as_secs_f64() * 1e3; + let _a_limbs = common::to_limbs(&a); + let _b_limbs = common::to_limbs(&b); + let host_pack_ms = t_host.elapsed().as_secs_f64() * 1e3;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ext/crates/fp-cuda/examples/bench_kernel.rs` around lines 57 - 69, Update the host-overhead measurement around t_host and host_ser_ms to time the complete packing path used before matmul_b1, including common::to_limbs for both serialized matrices. Alternatively, rename the metric and comments to clearly identify that it measures serialization only, not full host packing.
45-55: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winLabel the result as best-of-three.
bestuses the minimum of three trials, but the output presents the value without identifying that policy. This can overstate throughput relative to an average or median measurement.Suggested wording
- " End-to-end: {:.1} ms → {:.1} binary TOPS", + " End-to-end (best of {trials}): {:.1} ms → {:.1} binary TOPS",Also applies to: 71-79
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ext/crates/fp-cuda/examples/bench_kernel.rs` around lines 45 - 55, Update the benchmark output around the throughput calculation in main to explicitly label the reported timing or throughput as best-of-three, matching the existing trials count and minimum-selection behavior. Apply the same label to the corresponding output covered by the comment’s additional location.ext/crates/fp-cuda/examples/bench_shapes.rs (1)
76-89: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDeclare the timing scope of these TOPS values.
bench_shapes.rscallsmatmul_b1_timed, which the siblingbench_kernel_only.rspresents as excluding host setup and H2D/D2H transfers. This table only saysTOPS, so users may incorrectly compare it with the end-to-end TOPS frombench_kernel.rs.- "{:>7} {:>7} {:>7} | {:>9} {:>6} | {:>9} | note", + "{:>7} {:>7} {:>7} | {:>9} {:>6} | {:>9} | note", ... - "M", "K", "N", "B (MB)", "fits", "TOPS" + "M", "K", "N", "B (MB)", "fits", "kernel TOPS"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ext/crates/fp-cuda/examples/bench_shapes.rs` around lines 76 - 89, Clarify the timing scope of the TOPS values in the table printed by the shapes benchmark around matmul_b1_timed, indicating that they represent kernel-only timing and exclude host setup and H2D/D2H transfers. Update the TOPS column label or adjacent note so it is clearly distinct from the end-to-end metrics reported by bench_kernel.rs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ext/crates/fp-cuda/src/lib.rs`:
- Around line 70-92: Update GpuContext::stream so its cached TLS stream is keyed
by the owning GpuContext/device instead of using one unscoped slot per thread.
Ensure calls on different contexts, such as devices 0 and 1, create and reuse
distinct streams, and add a regression test covering both contexts on the same
thread.
---
Outside diff comments:
In `@ext/crates/fp-cuda/examples/bench_kernel.rs`:
- Around line 57-69: Update the host-overhead measurement around t_host and
host_ser_ms to time the complete packing path used before matmul_b1, including
common::to_limbs for both serialized matrices. Alternatively, rename the metric
and comments to clearly identify that it measures serialization only, not full
host packing.
- Around line 45-55: Update the benchmark output around the throughput
calculation in main to explicitly label the reported timing or throughput as
best-of-three, matching the existing trials count and minimum-selection
behavior. Apply the same label to the corresponding output covered by the
comment’s additional location.
In `@ext/crates/fp-cuda/examples/bench_shapes.rs`:
- Around line 76-89: Clarify the timing scope of the TOPS values in the table
printed by the shapes benchmark around matmul_b1_timed, indicating that they
represent kernel-only timing and exclude host setup and H2D/D2H transfers.
Update the TOPS column label or adjacent note so it is clearly distinct from the
end-to-end metrics reported by bench_kernel.rs.
In `@ext/crates/fp-cuda/README.md`:
- Around line 210-212: Remove the stale “Extend the parent Nix flake to provide
nvcc when the user opts in” roadmap item from the README, since the opt-in GPU
shell already supplies the CUDA toolkit and nvcc. Leave the unrelated
device-residency item unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 8836bf7e-5de7-47f4-9da2-d8c9765d2a03
📒 Files selected for processing (9)
ext/Cargo.tomlext/crates/fp-cuda/README.mdext/crates/fp-cuda/build.rsext/crates/fp-cuda/examples/bench_kernel.rsext/crates/fp-cuda/examples/bench_shapes.rsext/crates/fp-cuda/src/lib.rsext/crates/fp/src/blas/cuda.rsext/crates/fp/tests/cuda_dispatch.rsext/flake.nix
GpuContext::stream() cached the per-thread stream in a thread-keyed slot with no context key, so a thread that built two contexts (e.g. one per device) would reuse the first context's stream for the second — wrong context/device. Move the cache into the GpuContext itself (Mutex<HashMap<ThreadId, Arc<CudaStream>>>): streams are now owned by and scoped to the context. The mutex guards only the map lookup, never a GPU submission, so it does not serialize device work (unlike the whole-op lock this stream design replaced). Adds stream_is_scoped_per_context: two GpuContexts on one thread get distinct streams; each context reuses its own. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ext/crates/fp-cuda/src/lib.rs`:
- Around line 45-49: The GpuContext::stream implementation retains one
CudaStream per ThreadId indefinitely through the process-wide context. Replace
the unbounded streams map with thread-lifetime cleanup or a bounded per-context
stream pool, while preserving distinct stream usage and per-context isolation;
ensure streams are released or reused when threads exit or capacity is reached.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 0d127abd-6b19-4655-acc3-fa3e7641be30
📒 Files selected for processing (1)
ext/crates/fp-cuda/src/lib.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
fp-cuda, a CUDA backend for F₂ (GF(2)) matrix multiplication on Hopper(H200/H100), and wires it into
fpbehind an optionalgpufeature so largep = 2matrix products dispatch to binary tensor cores.What's here
wgmma.b1(binary AND+popcount tensor-core) GEMM kernel, both operandspre-arranged on the host as K-major tiles and loaded via TMA with 128B
swizzle into the layout the swizzled wgmma descriptors expect.
(TMA + double-buffering + producer/consumer warpgroup specialization),
128B swizzle,
m64n256k256, per-warpgroup register reallocation, deeperK-pipeline, TMA bulk output store, persistent kernel + grouped tile
rasterization, thread-block clusters + TMA multicast of B, output-tile
register-blocking, store/compute overlap, and fence hoisting.
cudarc(stable Rust).fpintegration (gpufeature):fp-cuda's library API is fp-agnostic(raw limbs), breaking the dependency cycle;
fpdispatches large p=2products to the device. The
fp-cudabuild.rsemits a stub PTX when nvccis absent, so
--workspaceCI stays green on runners without a GPU.Testing
cargo test -p fp --features gpu --test cuda_dispatch—gpu_dispatch_matches_cpupasses on an H200 (bit-exact vs the CPU path).
cargo build -p fp(default, no GPU) and--features gpuboth compile.Base for the follow-up device-resident row-reduction PR, which rebases on top
of this.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
gpufeature.Documentation
Tests