feat(moe): moe_mixture_forward kernel + Rust wrapper + unit test

Single-thread-per-(b,c) kernel, no atomicAdd, capture-friendly.
CPU-reference test verifies correctness for B=4, K=8, C=256 within 1e-6
tolerance.

Test wrapper uses mapped pinned memory exclusively
(feedback_no_htod_htoh_only_mapped_pinned.md): allocate MappedF32Buffer,
write CPU-side via host_ptr, kernel reads via dev_ptr, output via host_ptr
read after stream sync. NO memcpy_stod / memcpy_dtov anywhere.

GpuMoeHead struct in crates/ml/src/cuda_pipeline/gpu_moe_head.rs follows
existing cuda_pipeline head pattern (cubin loaded once, kernel handles
cached). Subsequent kernels (moe_mixture_backward,
moe_load_balance_loss, moe_expert_util_ema_update) extend the same
struct.

Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.2.
Plan: docs/superpowers/plans/2026-04-27-moe-regime-redesign.md Task 2.1.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/ml/build.rs

@@ -187,6 +187,12 @@ fn main() {
         // integrate into `experience_action_select`. Test-only; no production
         // callers in this commit.
         "thompson_test_kernel.cu",
+        // MoE Regime Redesign Phase 2 Task 2.1: mixture-forward kernel.
+        // Single-thread-per-(b,c), no atomicAdd, capture-friendly.
+        // h_s2[b,c] = Σ_k g[b,k] · expert_outputs[k,b,c].
+        // Subsequent tasks extend this cubin with moe_mixture_backward,
+        // moe_load_balance_loss, moe_expert_util_ema_update.
+        "moe_kernels.cu",
     ];
 
     // ALL kernels get common header (BF16 types + wrappers)

crates/ml/src/cuda_pipeline/gpu_moe_head.rs

@@ -0,0 +1,130 @@
+#![allow(unsafe_code)] // Required for CUDA kernel launches.
+
+//! Rust wrapper for MoE GPU kernels.
+//! Loaded from cubin at runtime, like other cuda_pipeline heads (gpu_iqn_head, etc.).
+//!
+//! Per `feedback_no_htod_htoh_only_mapped_pinned.md`: all CPU<->GPU
+//! communication uses mapped pinned memory. The test entry points
+//! allocate MappedF32Buffers, write CPU data via host_ptr, launch the
+//! kernel against dev_ptr, sync, and read GPU-written output via host_ptr.
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
+
+use crate::MLError;
+use super::mapped_pinned::MappedF32Buffer;
+
+static MOE_CUBIN: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));
+
+/// GPU-accelerated MoE head.
+///
+/// Owns the kernel function handles for the MoE forward mixture pass.
+/// Subsequent tasks will add: moe_mixture_backward, moe_dgate_reduce,
+/// moe_load_balance_loss, moe_load_balance_reduce, moe_expert_util_ema_update.
+#[allow(missing_debug_implementations)]
+pub struct GpuMoeHead {
+    stream: Arc<CudaStream>,
+    moe_mixture_forward: CudaFunction,
+}
+
+impl GpuMoeHead {
+    /// Load MoE cubin and cache all kernel function handles.
+    /// Mirrors the loader pattern in `GpuAuxHeads::new` and `GrnBlock::new`.
+    pub fn new(stream: Arc<CudaStream>) -> Result<Self, MLError> {
+        let context = stream.context();
+        let module = context
+            .load_cubin(MOE_CUBIN.to_vec())
+            .map_err(|e| MLError::ModelError(format!("MoE cubin load: {e}")))?;
+        let moe_mixture_forward = module
+            .load_function("moe_mixture_forward")
+            .map_err(|e| MLError::ModelError(format!("moe_mixture_forward load: {e}")))?;
+        Ok(Self {
+            stream,
+            moe_mixture_forward,
+        })
+    }
+
+    /// Test-only entry — stages CPU data via mapped pinned memory, runs the
+    /// kernel, reads mapped pinned output. NO HtoD/HtoH per
+    /// `feedback_no_htod_htoh_only_mapped_pinned.md`.
+    ///
+    /// Arguments:
+    ///   `expert_outputs` — `[K, B, C]` row-major flat slice.
+    ///   `gate`           — `[B, K]` row-major flat slice (rows must sum to 1).
+    ///   `b`              — batch size.
+    ///   `k`              — number of experts.
+    ///   `c`              — hidden/channel dimension.
+    ///
+    /// Returns `h_s2` as a `[B, C]` flat `Vec<f32>`.
+    pub fn test_mixture_forward(
+        &self,
+        expert_outputs: &[f32],
+        gate: &[f32],
+        b: usize,
+        k: usize,
+        c: usize,
+    ) -> Result<Vec<f32>, MLError> {
+        assert_eq!(
+            expert_outputs.len(),
+            k * b * c,
+            "expert_outputs length mismatch: expected {} ({k}×{b}×{c}), got {}",
+            k * b * c,
+            expert_outputs.len()
+        );
+        assert_eq!(
+            gate.len(),
+            b * k,
+            "gate length mismatch: expected {} ({b}×{k}), got {}",
+            b * k,
+            gate.len()
+        );
+
+        // Allocate mapped pinned buffers (CPU-and-GPU coherent).
+        let expert_buf = unsafe { MappedF32Buffer::new(k * b * c) }
+            .map_err(|e| MLError::ModelError(format!("expert buf alloc: {e}")))?;
+        let gate_buf = unsafe { MappedF32Buffer::new(b * k) }
+            .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?;
+        let h_s2_buf = unsafe { MappedF32Buffer::new(b * c) }
+            .map_err(|e| MLError::ModelError(format!("h_s2 buf alloc: {e}")))?;
+
+        // Direct host_ptr writes, no memcpy.
+        expert_buf.write_from_slice(expert_outputs);
+        gate_buf.write_from_slice(gate);
+
+        let total = (b * c) as u32;
+        let block: u32 = 256;
+        let grid = total.div_ceil(block);
+
+        let dev_expert = expert_buf.dev_ptr;
+        let dev_gate = gate_buf.dev_ptr;
+        let dev_h_s2 = h_s2_buf.dev_ptr;
+        let b_i32 = b as i32;
+        let k_i32 = k as i32;
+        let c_i32 = c as i32;
+
+        unsafe {
+            self.stream
+                .launch_builder(&self.moe_mixture_forward)
+                .arg(&dev_expert)
+                .arg(&dev_gate)
+                .arg(&dev_h_s2)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .arg(&c_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (grid, 1, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?;
+        }
+        self.stream
+            .synchronize()
+            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;
+
+        // Read GPU-written result via host_ptr (mapped pinned coherence).
+        Ok(h_s2_buf.read_all())
+    }
+}

crates/ml/src/cuda_pipeline/mod.rs

@@ -43,6 +43,7 @@ pub mod gpu_attention;
 pub mod gpu_tlob;
 pub mod gpu_grn;
 pub mod gpu_aux_heads;
+pub mod gpu_moe_head;
 pub mod decision_transformer;
 pub mod learning_health;
 pub mod q_snapshot;

crates/ml/src/cuda_pipeline/moe_kernels.cu

@@ -0,0 +1,33 @@
+/* MoE mixture / load-balance / ISV-producer kernels for the Mixture-of-
+ * Experts redesign. Spec:
+ * docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.2.
+ *
+ * GPU-resident, capture-friendly, no atomicAdd
+ * (per `feedback_no_atomicadd.md`).
+ */
+
+extern "C" __global__ void moe_mixture_forward(
+    const float* __restrict__ expert_outputs,  /* [K, B, C] */
+    const float* __restrict__ gate,            /* [B, K] */
+    float*       __restrict__ h_s2,            /* [B, C] */
+    int B,
+    int K,
+    int C
+) {
+    /* One thread per (b, c) pair. */
+    int total = B * C;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= total) return;
+
+    int b = tid / C;
+    int c = tid % C;
+
+    float acc = 0.0f;
+    for (int k = 0; k < K; ++k) {
+        float g  = gate[b * K + k];
+        float ek = expert_outputs[k * (B * C) + b * C + c];
+        acc += g * ek;
+    }
+    h_s2[tid] = acc;
+    __threadfence_system();  /* make write PCIe-visible to mapped pinned host_ptr */
+}

crates/ml/tests/moe_kernels_test.rs

@@ -0,0 +1,62 @@
+//! Layer 1 unit tests for MoE kernels.
+//! Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §8.
+
+#![allow(unsafe_code)] // Required for CUDA + mapped pinned memory.
+
+use std::sync::Arc;
+use cudarc::driver::CudaContext;
+use ml::cuda_pipeline::gpu_moe_head::GpuMoeHead;
+
+#[test]
+#[ignore]  // Requires CUDA; run with --ignored.
+fn moe_mixture_forward_matches_cpu_reference() {
+    let ctx = CudaContext::new(0).expect("CUDA context — GPU available?");
+    let stream = ctx.default_stream();
+
+    const B: usize = 4;
+    const C: usize = 256;
+    const K: usize = 8;
+
+    // Random expert outputs [K, B, C] and gate [B, K]
+    let expert_outputs_host: Vec<f32> = (0..K * B * C)
+        .map(|i| ((i * 7 + 3) % 100) as f32 / 100.0)
+        .collect();
+    let mut gate_host: Vec<f32> = (0..B * K)
+        .map(|i| ((i * 11 + 5) % 100) as f32 / 100.0)
+        .collect();
+    // Normalize gate rows to sum to 1
+    for b in 0..B {
+        let row_sum: f32 = (0..K).map(|k| gate_host[b * K + k]).sum();
+        for k in 0..K {
+            gate_host[b * K + k] /= row_sum;
+        }
+    }
+
+    // CPU reference: h_s2[b, c] = sum_k g[b, k] * e_k[b, c]
+    let mut expected: Vec<f32> = vec![0.0; B * C];
+    for b in 0..B {
+        for c in 0..C {
+            for k in 0..K {
+                expected[b * C + c] +=
+                    gate_host[b * K + k] * expert_outputs_host[k * B * C + b * C + c];
+            }
+        }
+    }
+
+    // GPU: stage via mapped pinned memory, run kernel, read mapped pinned output.
+    let head = GpuMoeHead::new(Arc::clone(&stream)).expect("MoE head init");
+    let actual = head
+        .test_mixture_forward(&expert_outputs_host, &gate_host, B, K, C)
+        .expect("kernel run");
+
+    // Compare with tolerance 1e-6
+    for i in 0..B * C {
+        assert!(
+            (expected[i] - actual[i]).abs() < 1e-6,
+            "Mismatch at index {}: expected {:.6}, got {:.6}",
+            i,
+            expected[i],
+            actual[i]
+        );
+    }
+}

docs/dqn-wire-up-audit.md

@@ -787,3 +787,12 @@ all kernel test wrappers (Test 0.F + upcoming MoE kernel tests) share
 one implementation. Adds `write_from_slice` helper for direct host_ptr
 writes (no memcpy). Test 0.F bit-identical post-move.
 Per `feedback_no_htod_htoh_only_mapped_pinned.md`.
+
+MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE
+CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] =
+Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust
+wrapper `GpuMoeHead` in `crates/ml/src/cuda_pipeline/gpu_moe_head.rs`
+loads cubin, exposes test_mixture_forward using mapped pinned buffers
+exclusively (no HtoD/HtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`).
+Unit test in `crates/ml/tests/moe_kernels_test.rs` verifies CPU
+reference match within 1e-6 for B=4, K=8, C=256.