feat(math-cuda): RPX256 leaf and tree kernels, launch module and device dispatch (lane K, phase 2)

[MauroToscano]

Lane K, phase 2 of the per-table GPU redo: the RPX256 (XHash12) device surface — leaf and tree kernels, the launch module, the cubin wiring, and the nine DeviceHash::Rpx256 dispatch arms turned from loud aborts into launches. Follows #952 (permutation + host-KAT) and #959 (canonicalisation pin); RPO256 and Poseidon keep their loud arms. Commits: b78d0d5 (the port), 1b22e66 (the code-shape fix the first cubin build forced — see below), 856091b (lane M's test-faults sticky comp-tree hook mirrored into the two RPX builders, per review), and f8d4147 (the literal-grouping lint fix: probe_states(0x0052_5058, 256); laptop make lint all six passes exit 0 at that head).

What changes

• crypto/math-cuda/kernels/rpx.cu — ten extern "C" kernels, kernel for kernel the twin of blake3.cu:338-620: rpx_leaves_base_batched, rpx_leaves_base_row_pair_batched, rpx_leaves_ext3_batched, rpx_comp_poly_leaves_ext3, rpx_fri_leaves_ext3, rpx_leaves_base_row_major_row_pair, …_range, rpx_merkle_level, rpx_merkle_tail, rpx_permute_probe. A node is four canonical felts as big-endian bytes (digest_to_commitment), so device trees are the host's bytes; a parent decodes its children as commitment_to_digest and runs compress.
• crypto/math-cuda/src/rpx.rs — the launch module, launcher for launcher the twin of blake3.rs (RPX_BLOCK_DIM = 128).
• build.rs compiles rpx.cu; lib.rs exports rpx and re-documents DeviceHash::Rpx256; device.rs loads rpx.cubin and its ten functions.
• The nine arms: lde.rs (inner tree levels; row-major row-pair leaves, plain and column-ranged; column-major base leaves per row and per pair; comp-poly ext3 leaves), fri.rs (FRI ext3 leaves; FRI inner levels), stark/src/gpu_lde.rs (comp-poly tree from ext3 evals; from resident slabs — inside match device_hash_of::<B>(), after lane M's admission checks).
• tests/host_kat/rpx_host_kat.cpp — layer 7: every leaf kernel, both compressors and the probe replayed thread by thread through the shim against the CPU leaf spec and the host parent (read patterns + node encoding, no GPU).
• prover/tests/rpx_device_parity.rs (cuda-gated) — the device gate; lives in the prover crate because RpxStarkHash and the algebraic backends do.

The hash_bytes == hash_data leaf contract

A device leaf absorbs exactly the felt sequence the host leaf hashes — bit-reversed rows, each column by column, an ext3 element as its three components (leaves_bit_reversed_grouped) — which is the sequence felts_from_bytes rebuilds from the leaf bytes, so the contract holds on device by construction. Raw [0, 2^64) storage is absorbed as is: the permutation is representation-independent and the host canonicalises before serialising; same field value, same digest. The felt count is known before the first absorb (the overwrite duplex's padding flag, A1).

Code shape — why permute is __noinline__ and every loop is rolled

The first cubin build (f546de7, box B) had every device function force-inlined and every loop fully unrolled. One permutation compiled to ~49k straight-line PTX lines (the 72-step inverse S-box chain over twelve lanes, three rounds) and every kernel carried one copy per permute call site — seven in the comp-poly kernel. The build was aborted after 41 minutes. 1b22e66 makes permute a real __noinline__ device function (empty under the host shim), gives the rounds a runtime index in a rolled round loop, rolls the lane, squaring, triple and constant loops (#pragma unroll 1), and rolls the MDS output-lane loop over a doubled circulant row (no modulo); the 24-MAC inner loop stays unrolled. Nothing arithmetic moved: the host-KAT passes unchanged with the same operation counts and the witness pin intact.

	f546de7 (aborted)	d3f235d / 1b22e66
rpx.ptx	55,808,145 B / 1,330,061 lines	191,959 B / 6,241 lines
.func bodies	0 (all inlined)	1 (permute, 1,426 lines)
rpx_comp_poly_leaves_ext3	345,057 lines	981
rpx_leaves_ext3_batched	197,091	573
rpx_leaves_base_row_pair_batched	147,895	786
rpx_leaves_base_row_major_row_pair / …_range	147,892 / 147,890	757 / 769
rpx_leaves_base_batched	98,507	482
rpx_fri_leaves_ext3	98,453	99
rpx_permute_probe / rpx_merkle_level / rpx_merkle_tail	49,385 / 48,955 / 48,915	78 / 127 / 137
ld.local / st.local	333 / 1,490	270 / 706
nvcc	cicc ~39 min, ptxas killed at 4.6 min / 2.1 GB	nvcc -ptx 0.99 s; full cargo build -p math-cuda --release 27.84 s
rpx.cubin	—	554,320 B (keccak.cubin 488,784 B)

ptxas -v at sm_120, every entry 0 bytes spill stores / 0 bytes spill loads: probe 40 regs · merkle_level 40 · merkle_tail 53 (1 barrier) · fri 40 · base_batched 56 · comp_poly 62 · row_pair_batched 64 · ext3_batched 64 · row_major_row_pair 64 · …_range 64. The price of the called shape is a 192–208 B per-thread stack frame: the twelve-lane state crosses the permute call by pointer and is indexed dynamically inside it (that is the ld.local/st.local traffic — by design, not spills), plus loop overhead of order 10% of the permutation's instructions. The unroll factors of square_n and the lane loops are the tuning knobs if profiling says so.

Gate

Box A @ d3f235d (rpx.cu, src/rpx.rs, the host-KAT files and the parity test are byte-identical at 1b22e66):

cargo build -p math-cuda --release: BUILD_EXIT=0, BUILD_WALL=27.84 s
cargo test -p lambda-vm-prover --release --features cuda --test rpx_device_parity -- --nocapture --test-threads=1
  rpx_comp_poly_tree_root_matches_cpu ok · rpx_device_permutation_matches_the_host_oracle ok
  rpx_fri_layer_tree_medium ok · rpx_fri_layer_tree_small ok
  rpx_fused_base_root_matches_cpu ok · rpx_fused_ext3_root_matches_cpu ok
  rpx_fused_tamper_diverges ok · rpx_row_major_range_leaves_match_cpu ok
  test result: ok. 8 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.54s

Laptop @ 1b22e66, stub cubins: cargo fmt --check, cargo check -p math-cuda --tests, cargo check -p stark --features cuda, cargo check -p lambda-vm-prover --features cuda --test rpx_device_parity — exit 0. make test-rpx-host-kat:

★ EXTERNAL: seven fb_round(s, r) = RPO256 vs miden-crypto hash_elements: 19/19 matched
★ ORACLE: rpx::permute vs Rust Rpx256::permute: 11/11 vectors matched
★ ORACLE: rpx::sponge_leaf vs Rust sponge_leaf(Rpx): 7/7 lengths matched
★ ORACLE: rpx::compress vs Rust HasherKind::Rpx.compress: 2/2 parents matched
★ canonicalisation pin: witness leaves 1 raw lane(s) >= p before the loop; permute() returns them canonical
-- layer 7: leaf kernels, Merkle compressors and the probe, replayed thread by thread --
base leaf kernels: read pattern + node encoding match the CPU leaf spec
ext3 + comp-poly leaf kernels: read pattern + node encoding match the CPU leaf spec
FRI leaf kernel: read pattern + node encoding match the Pair backend's leaf
row-major leaf kernels: read pattern + node encoding match the CPU leaf spec, all column ranges
Merkle compressors: level and tail kernels reproduce the host parent over a 16-leaf tree
permute probe: 11 oracle states reproduced through the kernel entry point
ALL HOST KAT CHECKS PASS

Box A @ d3f235d, the remaining two steps (relayed verbatim): make test-rpx-host-kat → KAT_EXIT=0, ALL HOST KAT CHECKS PASS (19/19 Miden · 11/11 permutation · 7/7 leaf · 2/2 parent · canonicalisation pin). make lint → LINT_EXIT=2: error: digits of hex, binary or octal literal not in groups of equal size at prover/tests/rpx_device_parity.rs:106 (probe_states(0x5250_58, 256), the cuda-arm clippy deny unusual_byte_groupings, invisible to cargo check) — fixed in the follow-up commit below.

Merge gate @ 1b22e66 on box A (relayed verbatim; the four RPX paths byte-identical to d3f235d):

cargo test -p lambda-vm-prover --release --features cuda --test rpx_device_parity: test result: ok. 8 passed; 0 failed; ... finished in 0.55s
make test-cuda-integration: test result: ok. 7 passed; 0 failed; ... finished in 14.09s   (ABORT 0, TEST-ONLY 0)
make test-cuda-d1:          test result: ok. 1 passed; 0 failed; ... finished in 2.61s
make test-cuda-fallback:    test result: ok. 5 passed; 0 failed; ... finished in 26.44s
                            test result: ok. 1 passed; 0 failed; ... finished in 12.32s
make lint: exit 2 — the rpx_device_parity.rs:106 literal (fixed by f8d41476)

Final gate (make test-cuda-fallback + parity + make lint) at f8d4147, appended here when relayed.

Open items

• The stack frame and the rolled loops above are the first shape that builds; profiling under a real prove decides whether square_n / the lane loops get a small unroll factor back.
• F6 (a run-permutation gather kernel for lane V's transpose run pass) is not in this PR.