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

crypto/math-cuda/build.rs

@@ -191,6 +191,10 @@ fn main() {
     compile_kernel("deep.cu", "deep.cubin", have_nvcc, &[]);
     compile_kernel("fri.cu", "fri.cubin", have_nvcc, &[]);
     compile_kernel("inverse.cu", "inverse.cubin", have_nvcc, &[]);
+    // RPX256 (XHash12) leaves and parents — the algebraic hash's device
+    // kernels. Pinned on the host by `tests/host_kat/rpx_host_kat.cpp`; the
+    // cubin needs no `-D`: RPX has no compile-time knob.
+    compile_kernel("rpx.cu", "rpx.cubin", have_nvcc, &[]);
     compile_kernel("logup.cu", "logup.cubin", have_nvcc, &[]);
     compile_kernel(
         "constraint_interp.cu",

crypto/math-cuda/src/device.rs

@@ -136,6 +136,7 @@ const LOGUP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/logup.cubin
 const CONSTRAINT_INTERP_CUBIN: &[u8] =
     include_bytes!(concat!(env!("OUT_DIR"), "/constraint_interp.cubin"));
 const BLAKE3_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/blake3.cubin"));
+const RPX_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rpx.cubin"));
 
 /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel
 /// callers overlap on the GPU without serializing on stream ownership. The
@@ -238,6 +239,22 @@ pub struct Backend {
     pub blake3_blocks_of_felts_probe: CudaFunction,
     pub blake3_chain_probe: CudaFunction,
 
+    // rpx.cubin — the RPX256 (XHash12) leaf kernels, Merkle level/tail
+    // compressors and the permutation probe (see `kernels/rpx.cu`). Twin for
+    // twin with the blake3 set above and in the same order; the probe is the
+    // only host-visible handle on the bare device permutation, which the parity
+    // tests check against the host `Rpx256`.
+    pub rpx_leaves_base_row_major_row_pair: CudaFunction,
+    pub rpx_leaves_base_row_major_row_pair_range: CudaFunction,
+    pub rpx_leaves_base_batched: CudaFunction,
+    pub rpx_leaves_base_row_pair_batched: CudaFunction,
+    pub rpx_leaves_ext3_batched: CudaFunction,
+    pub rpx_comp_poly_leaves_ext3: CudaFunction,
+    pub rpx_fri_leaves_ext3: CudaFunction,
+    pub rpx_merkle_level: CudaFunction,
+    pub rpx_merkle_tail: CudaFunction,
+    pub rpx_permute_probe: CudaFunction,
+
     // barycentric.cubin
     pub barycentric_base_batched: CudaFunction,
     pub barycentric_ext3_batched: CudaFunction,
@@ -448,6 +465,7 @@ impl Backend {
         let constraint_interp =
             ctx.load_module(Ptx::from_binary(CONSTRAINT_INTERP_CUBIN.to_vec()))?;
         let blake3 = ctx.load_module(Ptx::from_binary(BLAKE3_CUBIN.to_vec()))?;
+        let rpx = ctx.load_module(Ptx::from_binary(RPX_CUBIN.to_vec()))?;
 
         let mut streams = Vec::with_capacity(STREAM_POOL_SIZE);
         for _ in 0..STREAM_POOL_SIZE {
@@ -560,6 +578,20 @@ impl Backend {
             blake3_serialize_felts_probe: blake3.load_function("blake3_serialize_felts_probe")?,
             blake3_blocks_of_felts_probe: blake3.load_function("blake3_blocks_of_felts_probe")?,
             blake3_chain_probe: blake3.load_function("blake3_chain_probe")?,
+
+            rpx_leaves_base_row_major_row_pair: rpx
+                .load_function("rpx_leaves_base_row_major_row_pair")?,
+            rpx_leaves_base_row_major_row_pair_range: rpx
+                .load_function("rpx_leaves_base_row_major_row_pair_range")?,
+            rpx_leaves_base_batched: rpx.load_function("rpx_leaves_base_batched")?,
+            rpx_leaves_base_row_pair_batched: rpx
+                .load_function("rpx_leaves_base_row_pair_batched")?,
+            rpx_leaves_ext3_batched: rpx.load_function("rpx_leaves_ext3_batched")?,
+            rpx_comp_poly_leaves_ext3: rpx.load_function("rpx_comp_poly_leaves_ext3")?,
+            rpx_fri_leaves_ext3: rpx.load_function("rpx_fri_leaves_ext3")?,
+            rpx_merkle_level: rpx.load_function("rpx_merkle_level")?,
+            rpx_merkle_tail: rpx.load_function("rpx_merkle_tail")?,
+            rpx_permute_probe: rpx.load_function("rpx_permute_probe")?,
             barycentric_base_batched: bary.load_function("barycentric_base_batched")?,
             barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?,
             barycentric_base_batched_strided: bary

crypto/math-cuda/src/fri.rs

@@ -205,7 +205,14 @@ impl FriCommitState {
                     num_leaves_u64,
                     &mut leaves_view,
                 )?,
-                DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!(
+                DeviceHash::Rpx256 => crate::rpx::launch_fri_leaves_ext3(
+                    self.stream.as_ref(),
+                    be,
+                    &out,
+                    num_leaves_u64,
+                    &mut leaves_view,
+                )?,
+                DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!(
                     "{:?} device commit not yet ported (FRI layer ext3 leaves)",
                     self.hash
                 ),
@@ -224,7 +231,13 @@ impl FriCommitState {
                 &mut nodes_dev,
                 num_leaves,
             )?,
-            DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!(
+            DeviceHash::Rpx256 => crate::rpx::build_inner_tree_levels(
+                self.stream.as_ref(),
+                be,
+                &mut nodes_dev,
+                num_leaves,
+            )?,
+            DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!(
                 "{:?} device commit not yet ported (FRI layer inner tree levels)",
                 self.hash
             ),

crypto/math-cuda/src/lde.rs

@@ -1054,7 +1054,10 @@ fn build_inner_tree_levels_for(
         DeviceHash::Blake3 => {
             crate::blake3::build_inner_tree_levels(stream, be, nodes_dev, leaves_len)
         }
-        DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => {
+        DeviceHash::Rpx256 => {
+            crate::rpx::build_inner_tree_levels(stream, be, nodes_dev, leaves_len)
+        }
+        DeviceHash::Rpo256 | DeviceHash::Poseidon => {
             unimplemented!("{hash:?} device commit not yet ported (inner tree levels)")
         }
     }
@@ -1141,7 +1144,16 @@ fn coset_lde_row_major_inner(
                 log_lde,
                 &mut leaves_view,
             )?,
-            DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => {
+            DeviceHash::Rpx256 => crate::rpx::launch_leaves_base_row_major_row_pair(
+                stream.as_ref(),
+                be,
+                &buf,
+                cols_u64,
+                lde_u64,
+                log_lde,
+                &mut leaves_view,
+            )?,
+            DeviceHash::Rpo256 | DeviceHash::Poseidon => {
                 unimplemented!("{hash:?} device commit not yet ported (row-major row-pair leaves)")
             }
         }
@@ -1339,7 +1351,18 @@ pub fn coset_lde_row_major_split_trees(
                     log_lde,
                     &mut leaves_view,
                 )?,
-                DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!(
+                DeviceHash::Rpx256 => crate::rpx::launch_leaves_base_row_major_row_pair_range(
+                    stream.as_ref(),
+                    be,
+                    &buf,
+                    cols_u64,
+                    col_start,
+                    col_end,
+                    lde_u64,
+                    log_lde,
+                    &mut leaves_view,
+                )?,
+                DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!(
                     "{hash:?} device commit not yet ported (row-major row-pair leaves, column range)"
                 ),
             }
@@ -2154,7 +2177,23 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner(
                 lde_u64,
                 &mut leaves_view,
             )?,
-            (DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon, _) => {
+            (DeviceHash::Rpx256, true) => crate::rpx::launch_leaves_base_row_pair(
+                stream.as_ref(),
+                &buf,
+                col_stride_u64,
+                m as u64,
+                lde_u64,
+                &mut leaves_view,
+            )?,
+            (DeviceHash::Rpx256, false) => crate::rpx::launch_leaves_base(
+                stream.as_ref(),
+                &buf,
+                col_stride_u64,
+                m as u64,
+                lde_u64,
+                &mut leaves_view,
+            )?,
+            (DeviceHash::Rpo256 | DeviceHash::Poseidon, _) => {
                 unimplemented!("{hash:?} device commit not yet ported (column-major base leaves)")
             }
         }
@@ -2399,7 +2438,17 @@ fn evaluate_poly_coset_batch_ext3_into_inner(
                     log_num_rows,
                     &mut leaves_view,
                 )?,
-                DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => {
+                DeviceHash::Rpx256 => crate::rpx::launch_comp_poly_leaves_ext3(
+                    stream.as_ref(),
+                    be,
+                    &buf,
+                    col_stride_u64,
+                    num_parts_u64,
+                    lde_u64,
+                    log_num_rows,
+                    &mut leaves_view,
+                )?,
+                DeviceHash::Rpo256 | DeviceHash::Poseidon => {
                     unimplemented!("{hash:?} device commit not yet ported (comp-poly ext3 leaves)")
                 }
             }

crypto/math-cuda/src/lib.rs

@@ -21,6 +21,7 @@ pub mod merkle;
 pub mod mmcs;
 pub mod ntt;
 pub mod nvtx;
+pub mod rpx;
 
 // Re-exported for downstream crates so they can refer to CUDA primitive
 // types without depending on cudarc directly.
@@ -42,8 +43,9 @@ pub type Result<T> = std::result::Result<T, cudarc::driver::DriverError>;
 /// (keccak-256, or `Blake3Chain` at the compiled round count), exactly as on
 /// the host.
 ///
-/// ★ The three ALGEBRAIC keys name hashes whose device kernels are not yet
-/// ported. Every dispatch site in this crate carries an arm for them that
+/// ★ Of the three ALGEBRAIC keys, [`DeviceHash::Rpx256`] is ported
+/// ([`rpx`]); RPO256 and Poseidon name hashes whose device kernels are not.
+/// Every dispatch site in this crate carries an arm for the unported keys that
 /// aborts with `unimplemented!` naming the hash — never an arm that launches a
 /// byte-hash kernel in its place. The keys exist ahead of their kernels so the
 /// host side (`stark::config::DeviceTreeBackend`) can name every commitment
@@ -59,8 +61,8 @@ pub enum DeviceHash {
     /// RPO256 leaves and parents. No device kernels yet: every dispatch site
     /// aborts loudly on this key.
     Rpo256,
-    /// RPX256 (XHash12) leaves and parents. No device kernels yet: every
-    /// dispatch site aborts loudly on this key.
+    /// RPX256 (XHash12) leaves and parents — [`rpx`]'s kernels, the
+    /// algebraic family's first device port.
     Rpx256,
     /// ⚠ Poseidon-original — UNSHIPPABLE on the host side too; present so the
     /// key set mirrors `CommitmentHash` one-to-one. No device kernels.