GPU offload for Nassau's Milnor multiply (p = 2)

[JoeyBF]

Reorganized successor to #264: the GPU Milnor-multiply offload, with the benchmark suite (#269, merged) and the seqno index (#270) split out.

Stacked on #270. This PR is based on the seqno branch, so its diff currently shows the seqno commits too. The net-new commit here is the top one, "Add CubeCL GPU offload…". Once #270 merges I'll rebase and the diff will reduce to the GPU commit alone. Review against #270, or just the head commit.

What

Offloads the admissible-matrix Milnor multiply at the heart of Nassau's get_partial_matrix to the GPU via CubeCL, for the p = 2 / trivial-profile regime. Gated behind the gpu feature.

• milnor_algebra.rs: an admissible-matrix multiply (multiply_basis_element_by_element_2 + AdmissibleMatrix) as the CPU reference model for the kernel, plus the #[cfg(feature = "gpu")] upload helpers. Covered by admissible_multiply_agrees_with_reference.
• milnor_gpu.rs: the batched CubeCL kernel (multiply_batch_kernel / multiply_batch_on_gpu), fusing every (R, s) product of one matrix build into a single launch, with dedup-by-R, device-resident admissible buffers, and per-launch memory release pinned to one CUDA stream. Single-kernel building blocks used only to validate it live under #[cfg(test)], so the production surface is just multiply_batch_on_gpu + GpuProduct.
• ext::nassau_gpu + the build_partial_matrix hook in nassau.rs: route the differential-matrix build through the GPU path when compiled in, opted into (NASSAU_GPU), applicable, and large enough to amortise launch overhead; otherwise the CPU path is unchanged.
• free_module.rs: skip zero generator blocks in act (a general CPU win).
• Supporting nassau_e2e example, GPU/CPU consistency tests, and the milnor_gpu_ab A/B bench.

Rebased onto current master

Master has moved a long way since this was opened, in two ways that touch this PR directly.

#280 packed the Milnor p-part into a u64. The host-side extraction is ported to the PPart API throughout. One place this is a straight improvement: the GPU resident cache keyed its per-R admissible-matrix table on Vec<PPartEntry>, and that key is now the packed PPart itself — Copy, Hash, no allocation per lookup. The admissible sweep in multiply_basis_element_by_element_2 was rewritten to assemble its result in a plain u64 and store it once, the same approach #280 took in PPartMultiplier; the trailing-zero trim it used to need is gone, because the packed form does not represent trailing zeros.

fp-cuda arrived, with ext's own gpu feature meaning "fp's Hopper F₂ matmul". Rather than introduce a second GPU feature name, ext's gpu now enables both backends: they sit at different layers (F₂ linear algebra vs the Milnor product) and are independent at runtime. Note the consequence — fp-cuda compiles its kernel with nvcc in a build script, so --features gpu now needs the toolkit at build time, which is why the workspace excludes that crate. The Milnor half alone has no such requirement (cubecl JITs through NVRTC at runtime), so cargo check -p algebra --features gpu type-checks it with no CUDA installed.

The flake.nix changes this PR used to carry are dropped entirely — master's devShells.gpu already provides the CUDA toolkit and loader paths they were adding.

Notes vs #264

Dead code and dev-only profiling scaffolding removed; the validation-only kernels moved behind #[cfg(test)]. The env-var dispatch and the process-global GPU context are unchanged (deferred as larger follow-ups).

Test plan

• cargo test -p algebra — 89 pass, incl. admissible_multiply_agrees_with_reference, the CPU oracle for the kernel
• cargo test -p ext --features nassau --test milnor_vs_nassau — CPU Nassau correctness
• cargo clippy -p algebra --features gpu --all-targets — type-checks the kernel and the cfg(test) validation kernels without CUDA
• just lint (nightly fmt + cargo hack clippy feature-powerset, -D warnings)
• GPU-execution tests (tests/nassau_gpu*.rs, and the in-crate milnor_gpu device tests) require a live CUDA device — not run here or in CI

🤖 Generated with Claude Code

---

Generated by Claude Code

[coderabbitai]

Warning

Review limit reached

Next included review available in 19 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Team

Run ID: a22276c3-82fc-4b5e-a869-49f0eb9b6dfe

📥 Commits

Reviewing files that changed from the base of the PR and between e1e0f6f and 0512d8a.

📒 Files selected for processing (15)

• ext/Cargo.toml
• ext/benches/milnor_gpu_ab.rs
• ext/crates/algebra/Cargo.toml
• ext/crates/algebra/benches/seqno.rs
• ext/crates/algebra/src/algebra/milnor_algebra.rs
• ext/crates/algebra/src/algebra/milnor_gpu.rs
• ext/crates/algebra/src/algebra/mod.rs
• ext/crates/algebra/src/module/free_module.rs
• ext/examples/nassau_e2e.rs
• ext/src/lib.rs
• ext/src/nassau.rs
• ext/src/nassau_gpu.rs
• ext/tests/nassau_gpu.rs
• ext/tests/nassau_gpu_reuse.rs
• ext/tests/nassau_gpu_timing.rs

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}