feat(moe): moe_mixture_backward + moe_dgate_reduce kernels + FD test

Two-stage backward: moe_mixture_backward computes de_k = g · dh_s2 in
one kernel; moe_dgate_reduce computes dg via shmem reduction over c.

FD test verifies analytic backward matches numerical for B=2, K=4, C=32
within 1e-3 tolerance (perturbation 1e-3).

All test data flows via mapped pinned buffers (no HtoD/HtoH).

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

crates/ml/src/cuda_pipeline/gpu_moe_head.rs

@@ -20,13 +20,15 @@ static MOE_CUBIN: &[u8] =
 
 /// 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.
+/// Owns the kernel function handles for the MoE mixture forward/backward pass.
+/// Subsequent tasks will add: 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,
+    moe_mixture_backward: CudaFunction,
+    moe_dgate_reduce: CudaFunction,
 }
 
 impl GpuMoeHead {
@@ -40,9 +42,17 @@ impl GpuMoeHead {
         let moe_mixture_forward = module
             .load_function("moe_mixture_forward")
             .map_err(|e| MLError::ModelError(format!("moe_mixture_forward load: {e}")))?;
+        let moe_mixture_backward = module
+            .load_function("moe_mixture_backward")
+            .map_err(|e| MLError::ModelError(format!("moe_mixture_backward load: {e}")))?;
+        let moe_dgate_reduce = module
+            .load_function("moe_dgate_reduce")
+            .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce load: {e}")))?;
         Ok(Self {
             stream,
             moe_mixture_forward,
+            moe_mixture_backward,
+            moe_dgate_reduce,
         })
     }
 
@@ -127,4 +137,85 @@ impl GpuMoeHead {
         // Read GPU-written result via host_ptr (mapped pinned coherence).
         Ok(h_s2_buf.read_all())
     }
+
+    /// Test-only entry for the backward pass.
+    ///
+    /// Launches `moe_mixture_backward` (de_k = g · dh_s2) and
+    /// `moe_dgate_reduce` (dg = Σ_c e_k · dh_s2, shmem reduction over c).
+    ///
+    /// Returns `(de_k [K*B*C], dg [B*K])`.
+    pub fn test_mixture_backward(
+        &self,
+        expert_outputs: &[f32],
+        gate: &[f32],
+        dh_s2: &[f32],
+        b: usize,
+        k: usize,
+        c: usize,
+    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
+        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 dh_s2_buf = unsafe { MappedF32Buffer::new(b * c) }
+            .map_err(|e| MLError::ModelError(format!("dh_s2 buf alloc: {e}")))?;
+        let de_k_buf = unsafe { MappedF32Buffer::new(k * b * c) }
+            .map_err(|e| MLError::ModelError(format!("de_k buf alloc: {e}")))?;
+        let dg_buf = unsafe { MappedF32Buffer::new(b * k) }
+            .map_err(|e| MLError::ModelError(format!("dg buf alloc: {e}")))?;
+
+        expert_buf.write_from_slice(expert_outputs);
+        gate_buf.write_from_slice(gate);
+        dh_s2_buf.write_from_slice(dh_s2);
+
+        let block: u32 = 256;
+        let total_kbc = (k * b * c) as u32;
+        let grid_kbc = total_kbc.div_ceil(block);
+
+        let dev_expert = expert_buf.dev_ptr;
+        let dev_gate = gate_buf.dev_ptr;
+        let dev_dh_s2 = dh_s2_buf.dev_ptr;
+        let dev_de_k = de_k_buf.dev_ptr;
+        let dev_dg = dg_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_backward)
+                .arg(&dev_dh_s2)
+                .arg(&dev_gate)
+                .arg(&dev_de_k)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .arg(&c_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (grid_kbc, 1, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_mixture_backward launch: {e}")))?;
+
+            self.stream
+                .launch_builder(&self.moe_dgate_reduce)
+                .arg(&dev_dh_s2)
+                .arg(&dev_expert)
+                .arg(&dev_dg)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .arg(&c_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (b as u32, k as u32, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: block * 4, // block * sizeof(f32)
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?;
+        }
+        self.stream
+            .synchronize()
+            .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;
+
+        Ok((de_k_buf.read_all(), dg_buf.read_all()))
+    }
 }

crates/ml/src/cuda_pipeline/moe_kernels.cu

@@ -31,3 +31,60 @@ extern "C" __global__ void moe_mixture_forward(
     h_s2[tid] = acc;
     __threadfence_system();  /* make write PCIe-visible to mapped pinned host_ptr */
 }
+
+extern "C" __global__ void moe_mixture_backward(
+    const float* __restrict__ dh_s2,           /* [B, C] */
+    const float* __restrict__ gate,            /* [B, K] */
+    float*       __restrict__ de_k,            /* [K, B, C] */
+    int B,
+    int K,
+    int C
+) {
+    /* One thread per (k, b, c). de_k[k,b,c] = g[b,k] * dh_s2[b,c]. */
+    int total = K * B * C;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= total) return;
+
+    int k = tid / (B * C);
+    int rem = tid - k * (B * C);
+    int b = rem / C;
+    int c = rem - b * C;
+
+    de_k[tid] = gate[b * K + k] * dh_s2[b * C + c];
+    __threadfence_system();
+}
+
+extern "C" __global__ void moe_dgate_reduce(
+    const float* __restrict__ dh_s2,           /* [B, C] */
+    const float* __restrict__ expert_outputs,  /* [K, B, C] */
+    float*       __restrict__ dg,              /* [B, K] */
+    int B,
+    int K,
+    int C
+) {
+    /* Grid: (B, K, 1); block: 256 threads (or chosen by host). */
+    int b = blockIdx.x;
+    int k = blockIdx.y;
+    if (b >= B || k >= K) return;
+
+    extern __shared__ float smem[];
+    int tid = threadIdx.x;
+    int block = blockDim.x;
+
+    float local = 0.0f;
+    for (int c = tid; c < C; c += block) {
+        local += dh_s2[b * C + c] * expert_outputs[k * (B * C) + b * C + c];
+    }
+    smem[tid] = local;
+    __syncthreads();
+
+    for (int s = block / 2; s > 0; s >>= 1) {
+        if (tid < s) smem[tid] += smem[tid + s];
+        __syncthreads();
+    }
+
+    if (tid == 0) {
+        dg[b * K + k] = smem[0];
+        __threadfence_system();
+    }
+}

crates/ml/tests/moe_kernels_test.rs

@@ -60,3 +60,100 @@ fn moe_mixture_forward_matches_cpu_reference() {
         );
     }
 }
+
+#[test]
+#[ignore]
+fn moe_mixture_backward_finite_differences() {
+    let ctx = CudaContext::new(0).expect("CUDA context");
+    let stream = ctx.default_stream();
+
+    const B: usize = 2;
+    const C: usize = 32;
+    const K: usize = 4;
+    const EPS: f32 = 1e-3;
+    const TOL: f32 = 1e-3;
+
+    let expert_outputs: Vec<f32> = (0..K * B * C)
+        .map(|i| ((i * 13 + 1) % 200) as f32 / 200.0 - 0.25)
+        .collect();
+    let mut gate: Vec<f32> = (0..B * K)
+        .map(|i| ((i * 17 + 7) % 100) as f32 / 100.0)
+        .collect();
+    for b in 0..B {
+        let s: f32 = (0..K).map(|k| gate[b * K + k]).sum();
+        for k in 0..K {
+            gate[b * K + k] /= s;
+        }
+    }
+
+    let dh_s2: Vec<f32> = (0..B * C)
+        .map(|i| ((i * 19 + 11) % 100) as f32 / 100.0 - 0.5)
+        .collect();
+
+    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
+    let (de_analytic, dg_analytic) = head
+        .test_mixture_backward(&expert_outputs, &gate, &dh_s2, B, K, C)
+        .unwrap();
+
+    // Numerical de_k via finite differences over h_s2 = sum_k g_k * e_k.
+    // d(loss)/d(e_k0[b0,c0]) = dh_s2[b0,c0] * gate[b0,k0]  (analytic).
+    // Numerically: reuse the cpu_mixture_forward helper.
+    for k0 in 0..K {
+        for b0 in 0..B {
+            for c0 in 0..C {
+                let idx = k0 * B * C + b0 * C + c0;
+                let mut e_plus = expert_outputs.clone();
+                e_plus[idx] += EPS;
+                let h_plus = cpu_mixture_forward(&e_plus, &gate, B, K, C);
+                let mut e_minus = expert_outputs.clone();
+                e_minus[idx] -= EPS;
+                let h_minus = cpu_mixture_forward(&e_minus, &gate, B, K, C);
+                let mut numerical: f32 = 0.0;
+                for i in 0..B * C {
+                    numerical += dh_s2[i] * (h_plus[i] - h_minus[i]) / (2.0 * EPS);
+                }
+                assert!(
+                    (de_analytic[idx] - numerical).abs() < TOL,
+                    "de_k[{},{},{}]: analytic {:.6}, numerical {:.6}",
+                    k0, b0, c0, de_analytic[idx], numerical,
+                );
+            }
+        }
+    }
+
+    // Numerical dg over gate (note: perturbing g without renormalizing is OK
+    // for the local gradient check since we're testing the kernel's stage
+    // before any softmax/gate-softmax backward — that lives in cuBLAS).
+    for b0 in 0..B {
+        for k0 in 0..K {
+            let idx = b0 * K + k0;
+            let mut g_plus = gate.clone();
+            g_plus[idx] += EPS;
+            let h_plus = cpu_mixture_forward(&expert_outputs, &g_plus, B, K, C);
+            let mut g_minus = gate.clone();
+            g_minus[idx] -= EPS;
+            let h_minus = cpu_mixture_forward(&expert_outputs, &g_minus, B, K, C);
+            let mut numerical: f32 = 0.0;
+            for i in 0..B * C {
+                numerical += dh_s2[i] * (h_plus[i] - h_minus[i]) / (2.0 * EPS);
+            }
+            assert!(
+                (dg_analytic[idx] - numerical).abs() < TOL,
+                "dg[{},{}]: analytic {:.6}, numerical {:.6}",
+                b0, k0, dg_analytic[idx], numerical,
+            );
+        }
+    }
+}
+
+fn cpu_mixture_forward(expert_outputs: &[f32], gate: &[f32], b: usize, k: usize, c: usize) -> Vec<f32> {
+    let mut out = vec![0.0_f32; b * c];
+    for bi in 0..b {
+        for ci in 0..c {
+            for ki in 0..k {
+                out[bi * c + ci] += gate[bi * k + ki] * expert_outputs[ki * b * c + bi * c + ci];
+            }
+        }
+    }
+    out
+}

docs/dqn-wire-up-audit.md

@@ -2,6 +2,13 @@
 
 **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
 
+MoE backward kernels T2.2 (2026-04-27): `moe_mixture_backward` computes
+`de_k[k,b,c] = g[b,k] · dh_s2[b,c]` (one thread per (k,b,c));
+`moe_dgate_reduce` computes `dg[b,k] = Σ_c e_k[k,b,c] · dh_s2[b,c]` via
+shmem-tree reduction over c (one block per (b,k)). Finite-differences test
+verifies analytic backward matches numerical for B=2, K=4, C=32 within 1e-3.
+All data flows via mapped pinned buffers. Consumers land in Phase 3 wire-up.
+
 MoE ISV slot reset registration (2026-04-27): registered fold-boundary
 reset entries for the 8 `MOE_EXPERT_UTIL_EMA` slots (reset to 1/K=0.125)
 and the `MOE_GATE_ENTROPY_EMA_INDEX` slot (reset to ln(8)). Producers