feat(moe): wire MoE forward + backward + load-balance + monitoring

Phase 3 of the MoE regime redesign per
docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md.

Atomic forward+backward wire-up — replaces the existing h_s2 producer
in fused_training.rs with the gated expert mixture:

  state[B, 128] -> shared GRN trunk -> h_s1[B, 256]
                       |
                       +-> 8 expert MLPs (256->64->256) -> expert_outputs[8, B, 256]
                       +-> gate (128->64->8 softmax) ----> gate[B, 8]
                       |
                  moe_mixture_forward -> h_s2[B, 256] -> branching heads + C51 + IQN

Backward: moe_mixture_backward (de_k = g · dh_s2) + moe_dgate_reduce
(dg = Σ_c e_k · dh_s2) + load-balance aux gradient + cuBLAS SGEMM
backward through gate + each expert's 2 linear layers. Adam optimizer
step now updates gate + 8 experts via params_buf.

Loss: λ · K · Σ_k (mean_b g[b,k])² added to total loss with λ from
hyperparams.moe_lambda (default 0.01).

Per-step ISV producer launch (moe_expert_util_ema_update) writes 8
utilization EMA + 1 gate-entropy EMA into ISV[118..127). Per-epoch
HEALTH_DIAG aux_moe line emits utilization vector + entropy live so
operators can see whether experts are differentiating or collapsing.

Smoke test: DONE. 3/3 folds, all checkpoints saved, 728s (12.1 min,
within 25-min budget). Gate differentiated by epoch 1: expert 2 rose
from 0.119 → 0.286 → 0.323 over fold 1-2 while others remained at
0.097-0.113. Gate entropy 1.611 at fold 2 epoch 4 < ln(8)=2.079.
val_loss finite across all 3 folds; average fold metric 22.4.

Per feedback_no_partial_refactor.md: all consumers of the h_s2 contract
(branching heads, IQN aux, attention focus, backward chain) migrate in
this single commit.

Per feedback_no_htod_htoh_only_mapped_pinned.md: no new HtoD/HtoH
introduced; gate softmax + expert outputs + mixture all GPU-resident,
ISV producer GPU-driven. Load-balance scalar uses cuMemAllocHost +
cuMemHostGetDevicePointer_v2 (mapped pinned), matching existing pattern.

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

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -63,6 +63,7 @@ use super::shared_cublas_handle::PerStreamCublasHandles;
 use super::gpu_aux_heads::{
     AuxHeadsBackwardOps, AuxHeadsForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_REGIME_K,
 };
+use super::gpu_moe_head::GpuMoeHead;
 
 // ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ──────
 pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin"));
@@ -1242,6 +1243,10 @@ pub struct GpuDqnTrainConfig {
     /// override to a smaller value (e.g. 10_000) so the seed→network transition
     /// is observable inside the smoke run.
     pub replay_seed_steps: u32,
+
+    /// Phase 3 MoE load-balance loss weight λ.
+    /// Load-balance loss = λ·K·Σ_k(mean_b g[b,k])². Default 0.01.
+    pub moe_lambda: f32,
 }
 
 impl Default for GpuDqnTrainConfig {
@@ -1290,6 +1295,7 @@ impl Default for GpuDqnTrainConfig {
             market_dim: 42, // Default: 42 base features. Overridden to 50 when OFI (MBP-10) enabled.
             total_epochs: 0,
             replay_seed_steps: 100_000,
+            moe_lambda: 0.01,
         }
     }
 }
@@ -2675,6 +2681,45 @@ pub struct GpuDqnTrainer {
     /// share so aux can never dominate the policy gradient.
     aux_weight: f32,
 
+    // ── Phase 3: MoE forward/backward buffers + orchestrator ─────────────
+    /// MoE kernel orchestrator (mixture fwd/bwd, load-balance, EMA, softmax).
+    moe_head: GpuMoeHead,
+    /// Gate sub-network layer-1 hidden activations [B, MOE_GATE_HIDDEN].
+    /// Saved for backward (dW_gate_w1 = gate_h1^T @ d_gate_h1).
+    moe_gate_h1_buf: CudaSlice<f32>,
+    /// Gate pre-softmax logits [B, MOE_NUM_EXPERTS].
+    /// Saved for softmax backward (needed by moe_softmax_backward).
+    moe_gate_pre_buf: CudaSlice<f32>,
+    /// Gate softmax output = gate weights g[B, MOE_NUM_EXPERTS].
+    /// Saved for mixture forward + mixture backward + load-balance loss.
+    moe_gate_softmax_buf: CudaSlice<f32>,
+    /// 8 expert layer-1 hidden buffers, each [B, MOE_EXPERT_BOTTLENECK].
+    /// Saved for backward (dW_ek_w1 = expert_h1_k^T @ d_expert_h1_k).
+    moe_expert_h1_bufs: [CudaSlice<f32>; MOE_NUM_EXPERTS],
+    /// Stacked expert outputs [K, B, MOE_SH2] (K=8, layout row-major).
+    /// Saved for both mixture forward (direct input) and moe_dgate_reduce.
+    moe_expert_outputs_buf: CudaSlice<f32>,
+    /// d(loss)/d(expert_outputs) [K, B, MOE_SH2] from moe_mixture_backward.
+    moe_de_k_buf: CudaSlice<f32>,
+    /// d(loss)/d(gate_soft) [B, MOE_NUM_EXPERTS] from moe_dgate_reduce.
+    moe_dg_buf: CudaSlice<f32>,
+    /// d(loss)/d(gate_pre) [B, MOE_NUM_EXPERTS] from moe_softmax_backward.
+    moe_dg_pre_buf: CudaSlice<f32>,
+    /// Accumulated d(loss)/d(save_h_s1) from all 8 expert backward dX.
+    /// After accumulation, DtoD-copied into bw_d_h_s2 for encoder_backward_chain.
+    moe_dh_s1_scratch: CudaSlice<f32>,
+    /// d(loss)/d(gate_h1) [B, MOE_GATE_HIDDEN] from gate_w2 backward dX.
+    moe_gate_dh1_buf: CudaSlice<f32>,
+    /// Per-expert load-balance loss [MOE_NUM_EXPERTS].
+    moe_load_balance_loss_per_k: CudaSlice<f32>,
+    /// Load-balance total loss scalar [1] — pinned device-mapped.
+    /// GPU writes via dev_ptr, CPU reads host_ptr for optional logging.
+    moe_load_balance_loss_total_pinned: *mut f32,
+    /// Device pointer for moe_load_balance_loss_total_pinned.
+    moe_load_balance_loss_total_dev_ptr: u64,
+    /// Lambda for the load-balance auxiliary loss term (from config.moe_lambda).
+    moe_lambda: f32,
+
     /// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v).
     /// Written by GPU RNG kernel before each graph replay.
     stochastic_depth_scale_buf: CudaSlice<f32>,
@@ -7782,6 +7827,310 @@ impl GpuDqnTrainer {
         )
     }
 
+    // ── Phase 3 MoE methods ─────────────────────────────────────────────────
+
+    /// Phase 3 T3.1–T3.3: MoE forward.
+    ///
+    /// Sequence:
+    ///   1. Gate: state[B,SD] → Linear(SD→GH) + bias → ReLU → gate_h1[B,GH]
+    ///   2. Gate: gate_h1[B,GH] → Linear(GH→K) + bias → gate_pre[B,K]
+    ///   3. Gate: gate_pre → moe_row_softmax → gate_soft[B,K]  (saves for backward)
+    ///   4. For k in [0, K): h_s1[B,SH2] → Linear(SH2→BTN)+bias → ReLU → expert_h1_k[B,BTN]
+    ///                        expert_h1_k → Linear(BTN→SH2)+bias → expert_out_k[B,SH2]
+    ///   5. moe_mixture_forward stacks expert_outs[K,B,SH2] × gate_soft[B,K] → h_s2_new[B,SH2]
+    ///   6. DtoD copy h_s2_new → save_h_s2 (replaces the GRN-produced save_h_s2)
+    ///
+    /// Inputs: reads `save_h_s1` (h_s1 from GRN block 1), `states_buf` (for gate),
+    ///         and online weight tensors [127..163).
+    /// Output: overwrites `save_h_s2` with the MoE mixture.
+    pub(crate) fn launch_moe_forward(&self) -> Result<(), MLError> {
+        let b   = self.config.batch_size;
+        let sh2 = self.config.shared_h2;
+        let sd  = ml_core::state_layout::STATE_DIM;
+
+        let param_sizes = compute_param_sizes(&self.config);
+        let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &param_sizes);
+
+        // Gate weight tensors [127..131)
+        let gate_w1 = on_w_ptrs[127];  // [SD, GH]
+        let gate_b1 = on_w_ptrs[128];  // [GH]
+        let gate_w2 = on_w_ptrs[129];  // [GH, K]
+        let gate_b2 = on_w_ptrs[130];  // [K]
+
+        // Gate sub-network forward
+        let gate_h1_ptr   = self.moe_gate_h1_buf.raw_ptr();
+        let gate_pre_ptr  = self.moe_gate_pre_buf.raw_ptr();
+        let gate_soft_ptr = self.moe_gate_softmax_buf.raw_ptr();
+        let states_ptr    = self.ptrs.states_buf;
+
+        // Step 1: state[B,SD] → Linear(SD→GH) → gate_h1[B,GH]
+        self.cublas_forward.sgemm_f32(&self.stream, gate_w1, states_ptr, gate_h1_ptr,
+            MOE_GATE_HIDDEN, b, sd, "moe_gate_layer1")?;
+        self.cublas_forward.launch_add_bias_relu_f32_raw(
+            &self.stream, gate_h1_ptr, gate_b1, MOE_GATE_HIDDEN, b)?;
+
+        // Step 2: gate_h1[B,GH] → Linear(GH→K) → gate_pre[B,K]
+        self.cublas_forward.sgemm_f32(&self.stream, gate_w2, gate_h1_ptr, gate_pre_ptr,
+            MOE_NUM_EXPERTS, b, MOE_GATE_HIDDEN, "moe_gate_layer2")?;
+        self.cublas_forward.launch_add_bias_f32_raw(
+            &self.stream, gate_pre_ptr, gate_b2, MOE_NUM_EXPERTS, b)?;
+
+        // Step 3: softmax → gate_soft[B,K]
+        self.moe_head.launch_row_softmax(gate_pre_ptr, gate_soft_ptr, b, MOE_NUM_EXPERTS)?;
+
+        // Steps 4: for each expert k: h_s1[B,SH2] → Linear(SH2→BTN)+ReLU → Linear(BTN→SH2)
+        let h_s1_ptr = self.ptrs.save_h_s1;
+        let expert_outputs_ptr = self.moe_expert_outputs_buf.raw_ptr();
+
+        for k in 0..MOE_NUM_EXPERTS {
+            let base_idx = 131 + k * 4;
+            let ew1 = on_w_ptrs[base_idx];     // [SH2, BTN]
+            let eb1 = on_w_ptrs[base_idx + 1]; // [BTN]
+            let ew2 = on_w_ptrs[base_idx + 2]; // [BTN, SH2]
+            let eb2 = on_w_ptrs[base_idx + 3]; // [SH2]
+
+            let eh1_ptr  = self.moe_expert_h1_bufs[k].raw_ptr();
+            // expert_out_k is the k-th [B, SH2] slice of moe_expert_outputs_buf
+            let eout_ptr = expert_outputs_ptr + (k * b * sh2) as u64 * std::mem::size_of::<f32>() as u64;
+
+            // h_s1[B,SH2] → Linear(SH2→BTN) → eh1[B,BTN]
+            self.cublas_forward.sgemm_f32(&self.stream, ew1, h_s1_ptr, eh1_ptr,
+                MOE_EXPERT_BOTTLENECK, b, sh2, &format!("moe_expert_{k}_layer1"))?;
+            self.cublas_forward.launch_add_bias_relu_f32_raw(
+                &self.stream, eh1_ptr, eb1, MOE_EXPERT_BOTTLENECK, b)?;
+
+            // eh1[B,BTN] → Linear(BTN→SH2) → eout[B,SH2]
+            self.cublas_forward.sgemm_f32(&self.stream, ew2, eh1_ptr, eout_ptr,
+                sh2, b, MOE_EXPERT_BOTTLENECK, &format!("moe_expert_{k}_layer2"))?;
+            self.cublas_forward.launch_add_bias_f32_raw(
+                &self.stream, eout_ptr, eb2, sh2, b)?;
+        }
+
+        // Step 5: moe_mixture_forward: expert_outputs[K,B,SH2] × gate_soft[B,K] → h_s2_new[B,SH2]
+        // Write directly into save_h_s2 — replaces the GRN output.
+        self.moe_head.launch_mixture_forward(
+            expert_outputs_ptr,
+            gate_soft_ptr,
+            self.ptrs.save_h_s2,
+            b, MOE_NUM_EXPERTS, sh2,
+        )?;
+
+        Ok(())
+    }
+
+    /// Phase 3 T3.4: compute load-balance loss and SAXPY into total_loss_dev_ptr.
+    ///
+    /// Calls `moe_head.launch_load_balance_loss` which runs `moe_load_balance_loss`
+    /// (K blocks × B reduction) + `moe_load_balance_reduce` (single thread sum) into
+    /// `moe_load_balance_loss_total_dev_ptr`. Then SAXPYs the scalar (alpha=1.0)
+    /// into `total_loss_dev_ptr` so it is included in the blended TD loss.
+    pub(crate) fn launch_moe_load_balance_loss(&self) -> Result<(), MLError> {
+        let b = self.config.batch_size;
+        let gate_ptr        = self.moe_gate_softmax_buf.raw_ptr();
+        let loss_per_k_ptr  = self.moe_load_balance_loss_per_k.raw_ptr();
+        let loss_total_ptr  = self.moe_load_balance_loss_total_dev_ptr;
+
+        // Compute per-expert load loss + reduce to scalar.
+        self.moe_head.launch_load_balance_loss(
+            gate_ptr, loss_per_k_ptr, loss_total_ptr,
+            b, MOE_NUM_EXPERTS, self.moe_lambda,
+        )?;
+
+        // SAXPY total_loss += 1.0 * lb_loss_total (both are device-mapped scalars).
+        let n: i32 = 1;
+        let alpha: f32 = 1.0;
+        let total_loss_ptr = self.total_loss_dev_ptr;
+        unsafe {
+            self.stream
+                .launch_builder(&self.saxpy_f32_kernel)
+                .arg(&total_loss_ptr)
+                .arg(&loss_total_ptr)
+                .arg(&alpha)
+                .arg(&n)
+                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
+                .map_err(|e| MLError::ModelError(format!("moe_lb_loss saxpy: {e}")))?;
+        }
+        Ok(())
+    }
+
+    /// Phase 3 T3.5: launch `moe_expert_util_ema_update` ISV producer.
+    ///
+    /// Single-block single-thread cold-path kernel. Reads `moe_gate_softmax_buf [B, K]`
+    /// and EMA-updates ISV[MOE_EXPERT_UTIL_EMA_BASE..+8] (per-expert mean gate)
+    /// and ISV[MOE_GATE_ENTROPY_EMA_INDEX] (gate entropy).
+    pub fn launch_moe_expert_util_ema(&self, ema_alpha: f32) -> Result<(), MLError> {
+        debug_assert!(self.isv_signals_dev_ptr != 0,
+            "launch_moe_expert_util_ema: isv_signals_dev_ptr must be non-zero");
+        let gate_ptr = self.moe_gate_softmax_buf.raw_ptr();
+        self.moe_head.launch_expert_util_ema(
+            gate_ptr,
+            self.isv_signals_dev_ptr,
+            self.config.batch_size,
+            MOE_NUM_EXPERTS,
+            MOE_EXPERT_UTIL_EMA_BASE,
+            MOE_GATE_ENTROPY_EMA_INDEX,
+            ema_alpha,
+        )
+    }
+
+    /// Phase 3 T3.6: MoE backward.
+    ///
+    /// Sequence (reverse of forward):
+    ///   1. moe_mixture_backward: dh_s2 × gate_soft → de_k[K,B,SH2]
+    ///      + moe_dgate_reduce: dh_s2 × expert_outs → dg[B,K]
+    ///   2. moe_softmax_backward: gate_soft × dg → dg_pre[B,K]
+    ///   3. Gate layer 2 backward: dg_pre [B,K] → dW_gate_w2 (w_ptrs[129]), db_gate_b2 (w_ptrs[130])
+    ///      + dX → moe_gate_dh1[B,GH]
+    ///   4. Gate layer 1 backward (dX from step 3, gated by ReLU): apply relu mask,
+    ///      → dW_gate_w1 (w_ptrs[127]), db_gate_b1 (w_ptrs[128])
+    ///   5. For each expert k:
+    ///      a. Expert layer 2 backward: de_k_slice[B,SH2] → dW_ew2, db_eb2 + dX → d_eh1_k[B,BTN]
+    ///      b. Expert layer 1 backward (dX from 5a gated by ReLU): → dW_ew1, db_eb1 + dX → d_h_s1_k[B,SH2]
+    ///   6. Accumulate 8× d_h_s1_k into moe_dh_s1_scratch[B,SH2]
+    ///   7. DtoD-copy moe_dh_s1_scratch → bw_d_h_s2 (REPLACES existing, so encoder
+    ///      backward chain sees the MoE-routed h_s1 gradient as its h_s2 input gradient)
+    ///
+    /// `grad_base`: base device pointer for parameter gradient accumulation.
+    pub(crate) fn launch_moe_backward(&self, grad_base: u64) -> Result<(), MLError> {
+        let b   = self.config.batch_size;
+        let sh2 = self.config.shared_h2;
+        let f32_size = std::mem::size_of::<f32>() as u64;
+
+        let param_sizes = compute_param_sizes(&self.config);
+        let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &param_sizes);
+
+        // Gradient buffer base pointers — offset into grad_base by param position.
+        // grad_base is laid out identically to the params buffer.
+        let grad_ptrs = f32_weight_ptrs_from_base(grad_base, &param_sizes);
+
+        let gate_soft_ptr   = self.moe_gate_softmax_buf.raw_ptr();
+        let gate_h1_ptr     = self.moe_gate_h1_buf.raw_ptr();
+        let expert_outs_ptr = self.moe_expert_outputs_buf.raw_ptr();
+        let de_k_ptr        = self.moe_de_k_buf.raw_ptr();
+        let dg_ptr          = self.moe_dg_buf.raw_ptr();
+        let dg_pre_ptr      = self.moe_dg_pre_buf.raw_ptr();
+        let gate_dh1_ptr    = self.moe_gate_dh1_buf.raw_ptr();
+        let dh_s1_scratch   = self.moe_dh_s1_scratch.raw_ptr();
+
+        // d_h_s2 holds dL/d(save_h_s2) accumulated by all prior backward ops
+        // (branches, aux heads). This is the upstream gradient w.r.t. the MoE output.
+        let d_h_s2_ptr = bw_raw_f32_ptr(&self.bw_d_h_s2, &self.stream);
+        let h_s1_ptr   = self.ptrs.save_h_s1;
+
+        // Step 1a: moe_mixture_backward — propagates d_h_s2 through the weighted sum to de_k.
+        // de_k[k,b,c] = gate_soft[b,k] * d_h_s2[b,c]
+        // Step 1b: moe_dgate_reduce — dg[b,k] = sum_c(d_h_s2[b,c] * expert_outs[k,b,c])
+        self.moe_head.launch_mixture_backward(
+            d_h_s2_ptr, gate_soft_ptr, expert_outs_ptr,
+            de_k_ptr, dg_ptr,
+            b, MOE_NUM_EXPERTS, sh2,
+        )?;
+
+        // Step 2: softmax backward — dg_pre[b,k] = gate_soft[b,k] * (dg[b,k] - dot(gate_soft[b,:], dg[b,:]))
+        self.moe_head.launch_softmax_backward(gate_soft_ptr, dg_ptr, dg_pre_ptr, b, MOE_NUM_EXPERTS)?;
+
+        // Step 3: gate layer 2 dW + dX
+        // dg_pre[B,K] w.r.t. gate_h1[B,GH] through W_gate_w2[GH,K]
+        self.cublas_backward.launch_dw_only(
+            &self.stream,
+            dg_pre_ptr, gate_h1_ptr,         // dy, x
+            grad_ptrs[129], grad_ptrs[130],   // dW_gate_w2, db_gate_b2
+            MOE_NUM_EXPERTS, MOE_GATE_HIDDEN, b,
+        )?;
+        // dX into gate_dh1 (beta=0: overwrite)
+        self.cublas_backward.launch_dx_only(
+            &self.stream,
+            dg_pre_ptr, w_ptrs[129], gate_dh1_ptr,
+            MOE_NUM_EXPERTS, MOE_GATE_HIDDEN, b, 0.0,
+        )?;
+
+        // Step 4: gate layer 1 dW — gated by ReLU of gate_h1.
+        // gate_h1 holds the post-ReLU activation saved by forward; gate_dh1[i] = 0
+        // wherever gate_h1[i] == 0 (i.e., the neuron was clamped). relu_mask
+        // on cublas_backward implements dx[i] *= (activation[i] > 0).
+        self.cublas_backward.relu_mask(
+            &self.stream, gate_dh1_ptr, gate_h1_ptr, b * MOE_GATE_HIDDEN,
+        )?;
+        // dW_gate_w1, db_gate_b1 (gate layer 1 is state[B,SD] → gate_h1[B,GH])
+        let states_ptr = self.ptrs.states_buf;
+        self.cublas_backward.launch_dw_only(
+            &self.stream,
+            gate_dh1_ptr, states_ptr,          // dy, x
+            grad_ptrs[127], grad_ptrs[128],    // dW_gate_w1, db_gate_b1
+            MOE_GATE_HIDDEN, ml_core::state_layout::STATE_DIM, b,
+        )?;
+
+        // Step 5+6: expert backward — 8 experts, accumulate dX into dh_s1_scratch.
+        //
+        // Buffer aliasing strategy: after step 4 gate backward, gate_h1_ptr
+        // (= moe_gate_h1_buf [B, GH]) is free. GH == BTN == 64, so gate_h1_ptr
+        // is repurposed as d_eh1 scratch [B, BTN] for each expert in turn.
+        // The per-expert saved activation (moe_expert_h1_bufs[k]) must NOT be
+        // overwritten until after the ReLU mask is applied; the sequence is:
+        //   a. dW_ew2 (reads eh1_saved as x — does NOT modify eh1_saved)
+        //   b. dX_ew2 → gate_h1_ptr (d_eh1 scratch, beta=0; eh1_saved intact)
+        //   c. relu_mask(d_eh1, eh1_saved)  (eh1_saved still valid here)
+        //   d. dW_ew1 (reads eh1_saved? No — reads d_eh1 from gate_h1_ptr as dy)
+        //   e. dX_ew1 → dh_s1_scratch (accumulate with beta 0 or 1)
+        let d_eh1_scratch = gate_h1_ptr;  // repurpose gate_h1_ptr as [B, BTN] scratch
+        for k in 0..MOE_NUM_EXPERTS {
+            let base_idx = 131 + k * 4;
+            let ew1 = w_ptrs[base_idx];      // [SH2, BTN]
+            let ew2 = w_ptrs[base_idx + 2];  // [BTN, SH2]
+            let g_ew1 = grad_ptrs[base_idx];
+            let g_eb1 = grad_ptrs[base_idx + 1];
+            let g_ew2 = grad_ptrs[base_idx + 2];
+            let g_eb2 = grad_ptrs[base_idx + 3];
+
+            let eh1_saved = self.moe_expert_h1_bufs[k].raw_ptr();  // post-ReLU, read-only here
+            let dek_k     = de_k_ptr + (k * b * sh2) as u64 * f32_size;
+
+            // a. Expert layer 2 dW: dek_k[B,SH2] × eh1_saved[B,BTN]^T → dW_ew2, db_eb2
+            //    launch_dw_only reads eh1_saved as x; does NOT modify it.
+            self.cublas_backward.launch_dw_only(
+                &self.stream,
+                dek_k, eh1_saved,      // dy, x
+                g_ew2, g_eb2,          // dW_ew2, db_eb2
+                sh2, MOE_EXPERT_BOTTLENECK, b,
+            )?;
+            // b. Expert layer 2 dX: dek_k × W_ew2^T → d_eh1 scratch (beta=0)
+            self.cublas_backward.launch_dx_only(
+                &self.stream,
+                dek_k, ew2, d_eh1_scratch,
+                sh2, MOE_EXPERT_BOTTLENECK, b, 0.0,
+            )?;
+            // c. ReLU mask: d_eh1[i] *= (eh1_saved[i] > 0)
+            self.cublas_backward.relu_mask(
+                &self.stream, d_eh1_scratch, eh1_saved, b * MOE_EXPERT_BOTTLENECK,
+            )?;
+            // d. Expert layer 1 dW: d_eh1[B,BTN] × h_s1[B,SH2]^T → dW_ew1, db_eb1
+            self.cublas_backward.launch_dw_only(
+                &self.stream,
+                d_eh1_scratch, h_s1_ptr,   // dy=d_eh1, x=h_s1
+                g_ew1, g_eb1,              // dW_ew1, db_eb1
+                MOE_EXPERT_BOTTLENECK, sh2, b,
+            )?;
+            // e. Expert layer 1 dX → dh_s1_scratch (accumulate; beta=0 for k=0, 1 for k>0)
+            let beta = if k == 0 { 0.0_f32 } else { 1.0_f32 };
+            self.cublas_backward.launch_dx_only(
+                &self.stream,
+                d_eh1_scratch, ew1, dh_s1_scratch,
+                MOE_EXPERT_BOTTLENECK, sh2, b, beta,
+            )?;
+        }
+
+        // Step 7: DtoD-copy moe_dh_s1_scratch → bw_d_h_s2.
+        // This REPLACES (not adds to) bw_d_h_s2 so encoder_backward_chain
+        // receives the MoE-accumulated h_s1 gradient as its h_s2 input gradient.
+        // The aux_heads_backward already wrote into bw_d_h_s2; we replace here
+        // because the MoE mixture output IS save_h_s2 — the GRN h_s2 block was
+        // replaced by MoE in forward, so its "d_h_s2" should be the MoE dX.
+        let n_bytes = b * sh2 * std::mem::size_of::<f32>();
+        self.graph_safe_copy_f32(d_h_s2_ptr, dh_s1_scratch, n_bytes, "moe_dh_s1→d_h_s2")?;
+
+        Ok(())
+    }
+
     /// Plan 4 Task 6 Commit B: aux-heads forward orchestrator.
     ///
     /// Runs the next-bar regression head + 5-class regime classification head
@@ -10129,6 +10478,63 @@ impl GpuDqnTrainer {
         // step. The first graph capture bakes 0.05 — an honest cold-start
         // value, not a stub.
         let aux_weight: f32 = 0.05;
+
+        // ── Phase 3: MoE forward/backward orchestrator + working buffers ─────
+        // GpuMoeHead loads `moe_kernels.cubin` and caches 8 kernel handles.
+        // All working buffers are zero-initialised — the gate starts at zero
+        // weights (uniform 1/K softmax), so cold-start EMA values are valid.
+        let moe_head = GpuMoeHead::new(Arc::clone(&stream))?;
+        let moe_sh2 = config.shared_h2;
+        let moe_gate_h1_buf    = alloc_f32(&stream, b * MOE_GATE_HIDDEN, "moe_gate_h1_buf")?;
+        let moe_gate_pre_buf   = alloc_f32(&stream, b * MOE_NUM_EXPERTS,  "moe_gate_pre_buf")?;
+        let moe_gate_softmax_buf = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_gate_softmax_buf")?;
+        // Allocate 8 expert hidden buffers (one per expert, each [B, BTN]).
+        let moe_expert_h1_bufs = {
+            let mut arr: [std::mem::MaybeUninit<CudaSlice<f32>>; MOE_NUM_EXPERTS] =
+                unsafe { std::mem::MaybeUninit::uninit().assume_init() };
+            for (ek, slot) in arr.iter_mut().enumerate() {
+                let buf = alloc_f32(&stream, b * MOE_EXPERT_BOTTLENECK,
+                    &format!("moe_expert_h1_{ek}"))?;
+                slot.write(buf);
+            }
+            unsafe { std::mem::transmute::<_, [CudaSlice<f32>; MOE_NUM_EXPERTS]>(arr) }
+        };
+        let moe_expert_outputs_buf = alloc_f32(&stream, MOE_NUM_EXPERTS * b * moe_sh2, "moe_expert_outputs_buf")?;
+        let moe_de_k_buf           = alloc_f32(&stream, MOE_NUM_EXPERTS * b * moe_sh2, "moe_de_k_buf")?;
+        let moe_dg_buf             = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_dg_buf")?;
+        let moe_dg_pre_buf         = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_dg_pre_buf")?;
+        let moe_dh_s1_scratch      = alloc_f32(&stream, b * moe_sh2, "moe_dh_s1_scratch")?;
+        let moe_gate_dh1_buf       = alloc_f32(&stream, b * MOE_GATE_HIDDEN, "moe_gate_dh1_buf")?;
+        let moe_load_balance_loss_per_k = alloc_f32(&stream, MOE_NUM_EXPERTS, "moe_load_balance_loss_per_k")?;
+        let (moe_load_balance_loss_total_pinned, moe_load_balance_loss_total_dev_ptr) = {
+            let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut();
+            let mut dev_ptr_out: u64 = 0;
+            unsafe {
+                let rc = cudarc::driver::sys::cuMemAllocHost_v2(
+                    &mut host_ptr as *mut *mut std::ffi::c_void,
+                    std::mem::size_of::<f32>(),
+                );
+                if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
+                    return Err(MLError::ModelError(format!("cuMemAllocHost moe_lb_total: {:?}", rc)));
+                }
+                let rc2 = cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
+                    &mut dev_ptr_out,
+                    host_ptr,
+                    0,
+                );
+                if rc2 != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
+                    return Err(MLError::ModelError(format!("cuMemHostGetDevicePointer moe_lb_total: {:?}", rc2)));
+                }
+                *(host_ptr as *mut f32) = 0.0f32;
+            }
+            (host_ptr as *mut f32, dev_ptr_out)
+        };
+        let moe_lambda = config.moe_lambda;
+        tracing::info!(
+            "GpuMoeHead initialized: K={} BTN={} GH={} SH2={} lambda={:.4} (params [127..163), ISV[118..127))",
+            MOE_NUM_EXPERTS, MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN, moe_sh2, moe_lambda,
+        );
+
         tracing::info!(
             "AuxHeadsForwardOps initialized: K_nb={} K_rg={} aux_h={} max_aux_tensor_len={} (params [119..127), ISV[113..115))",
             aux_knb, aux_kr, aux_h, max_aux_tensor_len,
@@ -11766,6 +12172,21 @@ impl GpuDqnTrainer {
             ),
             last_distill_active: false,
             last_meta_q_pred: 0.5,
+            moe_head,
+            moe_gate_h1_buf,
+            moe_gate_pre_buf,
+            moe_gate_softmax_buf,
+            moe_expert_h1_bufs,
+            moe_expert_outputs_buf,
+            moe_de_k_buf,
+            moe_dg_buf,
+            moe_dg_pre_buf,
+            moe_dh_s1_scratch,
+            moe_gate_dh1_buf,
+            moe_load_balance_loss_per_k,
+            moe_load_balance_loss_total_pinned,
+            moe_load_balance_loss_total_dev_ptr,
+            moe_lambda,
             q_sample_history: std::collections::VecDeque::with_capacity(64),
         };
         // Self-check: verify the fingerprint we just wrote is readable and matches
@@ -15271,6 +15692,13 @@ impl GpuDqnTrainer {
         self.launch_loss_reduce(self.total_loss_dev_ptr)?;
         self.launch_c51_grad()?;
 
+        // ── Phase 3 T3.4: MoE load-balance auxiliary loss ────────────────
+        // Computes λ·K·Σ_k(mean_b g[b,k])² and SAXPYs the scalar into
+        // total_loss_dev_ptr. Runs after C51 loss so it adds on top of the
+        // already-accumulated TD loss. Does NOT generate a direct gradient
+        // (the gradient flows through the gate backward in launch_moe_backward).
+        self.launch_moe_load_balance_loss()?;
+
         // Blend: main = α * C51 + (1-α) * MSE
         {
             let alpha = self.c51_alpha;
@@ -15908,6 +16336,14 @@ impl GpuDqnTrainer {
             self.ptrs.mag_concat_buf,
         )?;
 
+        // ── Phase 3 T3.1–T3.3: MoE forward ─────────────────────────────────
+        // Runs AFTER forward_online_raw (which produces save_h_s1 and the
+        // pre-MoE save_h_s2) and BEFORE aux_heads_forward (which reads
+        // save_h_s2). Overwrites save_h_s2 with the MoE mixture output
+        // so all downstream consumers (aux heads, stochastic depth, loss
+        // kernels) see the MoE-enriched representation.
+        self.launch_moe_forward()?;
+
         // ── Plan 4 Task 6 Commit B: aux-heads forward (online-only) ──
         // Runs the next-bar regression + 5-class regime CE heads off the
         // freshly-saved `save_h_s2` activation. MUST run BEFORE stochastic
@@ -17137,6 +17573,14 @@ impl GpuDqnTrainer {
         // augmented dh_s2 (= main + aux contributions).
         self.aux_heads_backward(grad_base)?;
 
+        // ── Phase 3 T3.6: MoE backward ─────────────────────────────────
+        // Runs AFTER aux_heads_backward (which SAXPYed into bw_d_h_s2) and
+        // BEFORE encoder_backward_chain. Computes dW/dB for MoE params
+        // [127..163) and accumulates expert dX → moe_dh_s1_scratch, then
+        // DtoD-copies moe_dh_s1_scratch into bw_d_h_s2 so the GRN encoder
+        // chain sees the MoE-routed gradient as its h_s2 input gradient.
+        self.launch_moe_backward(grad_base)?;
+
         // ── GRN trunk backward (Plan 4 Task 2c.3c.4) ──
         // d_h_s2_ptr now holds the fully-accumulated gradient flowing into
         // the h_s2 GRN output. encoder_backward_chain unrolls both GRN

crates/ml/src/cuda_pipeline/gpu_moe_head.rs

@@ -21,7 +21,8 @@ static MOE_CUBIN: &[u8] =
 /// GPU-accelerated MoE head.
 ///
 /// Owns the kernel function handles for the MoE mixture forward/backward pass,
-/// load-balance loss computation, and expert utilization EMA update.
+/// load-balance loss computation, expert utilization EMA update, and softmax
+/// forward/backward for the gate sub-network.
 #[allow(missing_debug_implementations)]
 pub struct GpuMoeHead {
     stream: Arc<CudaStream>,
@@ -31,6 +32,10 @@ pub struct GpuMoeHead {
     moe_load_balance_loss: CudaFunction,
     moe_load_balance_reduce: CudaFunction,
     moe_expert_util_ema_update: CudaFunction,
+    /// Per-row softmax for the gate pre-activation → gate probabilities.
+    moe_row_softmax: CudaFunction,
+    /// Softmax Jacobian backward for the gate.
+    moe_softmax_backward: CudaFunction,
 }
 
 impl GpuMoeHead {
@@ -59,6 +64,12 @@ impl GpuMoeHead {
         let moe_expert_util_ema_update = module
             .load_function("moe_expert_util_ema_update")
             .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update load: {e}")))?;
+        let moe_row_softmax = module
+            .load_function("moe_row_softmax")
+            .map_err(|e| MLError::ModelError(format!("moe_row_softmax load: {e}")))?;
+        let moe_softmax_backward = module
+            .load_function("moe_softmax_backward")
+            .map_err(|e| MLError::ModelError(format!("moe_softmax_backward load: {e}")))?;
         Ok(Self {
             stream,
             moe_mixture_forward,
@@ -67,9 +78,243 @@ impl GpuMoeHead {
             moe_load_balance_loss,
             moe_load_balance_reduce,
             moe_expert_util_ema_update,
+            moe_row_softmax,
+            moe_softmax_backward,
         })
     }
 
+    // ══════════════════════════════════════════════════════════════════════
+    // Production launch methods (operate on already-allocated device buffers)
+    // ══════════════════════════════════════════════════════════════════════
+
+    /// Launch `moe_row_softmax` in-place on already-allocated device pointers.
+    /// Grid: (B+block-1)/block × 1 × 1. Each thread handles one row.
+    pub(crate) fn launch_row_softmax(
+        &self,
+        in_pre_ptr: u64,
+        out_soft_ptr: u64,
+        b: usize,
+        k: usize,
+    ) -> Result<(), MLError> {
+        let block: u32 = 256;
+        let grid = ((b as u32) + block - 1) / block;
+        let b_i32 = b as i32;
+        let k_i32 = k as i32;
+        unsafe {
+            self.stream
+                .launch_builder(&self.moe_row_softmax)
+                .arg(&in_pre_ptr)
+                .arg(&out_soft_ptr)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (grid, 1, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_row_softmax launch: {e}")))?;
+        }
+        Ok(())
+    }
+
+    /// Launch `moe_mixture_forward` on device pointers.
+    /// h_s2_out[B,C] = sum_k gate[B,k] * expert_outputs[k,B,C].
+    pub(crate) fn launch_mixture_forward(
+        &self,
+        expert_outputs_ptr: u64,
+        gate_ptr: u64,
+        h_s2_out_ptr: u64,
+        b: usize,
+        k: usize,
+        c: usize,
+    ) -> Result<(), MLError> {
+        let total = (b * c) as u32;
+        let block: u32 = 256;
+        let grid = total.div_ceil(block);
+        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(&expert_outputs_ptr)
+                .arg(&gate_ptr)
+                .arg(&h_s2_out_ptr)
+                .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}")))?;
+        }
+        Ok(())
+    }
+
+    /// Launch `moe_load_balance_loss` + `moe_load_balance_reduce`.
+    /// Writes the scalar total load-balance loss into `loss_total_ptr`.
+    pub(crate) fn launch_load_balance_loss(
+        &self,
+        gate_ptr: u64,
+        loss_per_k_ptr: u64,
+        loss_total_ptr: u64,
+        b: usize,
+        k: usize,
+        lambda: f32,
+    ) -> Result<(), MLError> {
+        let block: u32 = 256;
+        let b_i32 = b as i32;
+        let k_i32 = k as i32;
+        unsafe {
+            self.stream
+                .launch_builder(&self.moe_load_balance_loss)
+                .arg(&gate_ptr)
+                .arg(&loss_per_k_ptr)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .arg(&lambda)
+                .launch(LaunchConfig {
+                    grid_dim: (k as u32, 1, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: block * 4,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_load_balance_loss launch: {e}")))?;
+
+            self.stream
+                .launch_builder(&self.moe_load_balance_reduce)
+                .arg(&loss_per_k_ptr)
+                .arg(&loss_total_ptr)
+                .arg(&k_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (1, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce launch: {e}")))?;
+        }
+        Ok(())
+    }
+
+    /// Launch `moe_expert_util_ema_update` ISV producer on device pointers.
+    pub(crate) fn launch_expert_util_ema(
+        &self,
+        gate_ptr: u64,
+        isv_dev_ptr: u64,
+        b: usize,
+        k: usize,
+        isv_util_base: usize,
+        isv_entropy_index: usize,
+        alpha: f32,
+    ) -> Result<(), MLError> {
+        let b_i32 = b as i32;
+        let k_i32 = k as i32;
+        let util_base_i32 = isv_util_base as i32;
+        let entropy_idx_i32 = isv_entropy_index as i32;
+        unsafe {
+            self.stream
+                .launch_builder(&self.moe_expert_util_ema_update)
+                .arg(&gate_ptr)
+                .arg(&isv_dev_ptr)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .arg(&util_base_i32)
+                .arg(&entropy_idx_i32)
+                .arg(&alpha)
+                .launch(LaunchConfig {
+                    grid_dim: (1, 1, 1),
+                    block_dim: (1, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update launch: {e}")))?;
+        }
+        Ok(())
+    }
+
+    /// Launch `moe_mixture_backward` (de_k) + `moe_dgate_reduce` (dg) on device pointers.
+    pub(crate) fn launch_mixture_backward(
+        &self,
+        dh_s2_ptr: u64,
+        gate_ptr: u64,
+        expert_outputs_ptr: u64,
+        de_k_ptr: u64,
+        dg_ptr: u64,
+        b: usize,
+        k: usize,
+        c: usize,
+    ) -> Result<(), MLError> {
+        let block: u32 = 256;
+        let total_kbc = (k * b * c) as u32;
+        let grid_kbc = total_kbc.div_ceil(block);
+        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(&dh_s2_ptr)
+                .arg(&gate_ptr)
+                .arg(&de_k_ptr)
+                .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(&dh_s2_ptr)
+                .arg(&expert_outputs_ptr)
+                .arg(&dg_ptr)
+                .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,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?;
+        }
+        Ok(())
+    }
+
+    /// Launch `moe_softmax_backward` on device pointers.
+    pub(crate) fn launch_softmax_backward(
+        &self,
+        g_soft_ptr: u64,
+        dg_ptr: u64,
+        dg_pre_ptr: u64,
+        b: usize,
+        k: usize,
+    ) -> Result<(), MLError> {
+        let block: u32 = 256;
+        let grid = ((b as u32) + block - 1) / block;
+        let b_i32 = b as i32;
+        let k_i32 = k as i32;
+        unsafe {
+            self.stream
+                .launch_builder(&self.moe_softmax_backward)
+                .arg(&g_soft_ptr)
+                .arg(&dg_ptr)
+                .arg(&dg_pre_ptr)
+                .arg(&b_i32)
+                .arg(&k_i32)
+                .launch(LaunchConfig {
+                    grid_dim: (grid, 1, 1),
+                    block_dim: (block, 1, 1),
+                    shared_mem_bytes: 0,
+                })
+                .map_err(|e| MLError::ModelError(format!("moe_softmax_backward launch: {e}")))?;
+        }
+        Ok(())
+    }
+
     /// 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`.

crates/ml/src/cuda_pipeline/moe_kernels.cu

@@ -137,6 +137,61 @@ extern "C" __global__ void moe_load_balance_reduce(
     __threadfence_system();
 }
 
+/* Row-wise softmax: out[b,k] = exp(in[b,k]) / sum_j exp(in[b,j]).
+ * Grid: one block per row (b). Single-thread per block — K is small (=8).
+ * No atomicAdd, no shared memory needed for K<=8 (all in registers). */
+extern "C" __global__ void moe_row_softmax(
+    const float* __restrict__ in_pre,   /* [B, K] pre-softmax logits */
+    float*       __restrict__ out_soft, /* [B, K] softmax output */
+    int B,
+    int K
+) {
+    int b = blockIdx.x * blockDim.x + threadIdx.x;
+    if (b >= B) return;
+
+    /* Numerically stable: subtract row max before exponentiation. */
+    float row_max = in_pre[b * K];
+    for (int k = 1; k < K; ++k) {
+        float v = in_pre[b * K + k];
+        if (v > row_max) row_max = v;
+    }
+    float sum_exp = 0.0f;
+    for (int k = 0; k < K; ++k) {
+        float e = expf(in_pre[b * K + k] - row_max);
+        out_soft[b * K + k] = e;
+        sum_exp += e;
+    }
+    float inv_sum = 1.0f / sum_exp;
+    for (int k = 0; k < K; ++k) {
+        out_soft[b * K + k] *= inv_sum;
+    }
+    __threadfence_system();
+}
+
+/* Softmax backward per row:
+ *   dg_pre[b,k] = g_soft[b,k] * (dg[b,k] - sum_j g_soft[b,j] * dg[b,j])
+ * Grid: one block per row (b). Single-thread per block — K is small (=8). */
+extern "C" __global__ void moe_softmax_backward(
+    const float* __restrict__ g_soft,   /* [B, K] saved softmax output from forward */
+    const float* __restrict__ dg,       /* [B, K] upstream gradient w.r.t. gate */
+    float*       __restrict__ dg_pre,   /* [B, K] gradient w.r.t. pre-softmax logits */
+    int B,
+    int K
+) {
+    int b = blockIdx.x * blockDim.x + threadIdx.x;
+    if (b >= B) return;
+
+    /* dot = sum_j g_soft[b,j] * dg[b,j] */
+    float dot = 0.0f;
+    for (int k = 0; k < K; ++k) {
+        dot += g_soft[b * K + k] * dg[b * K + k];
+    }
+    for (int k = 0; k < K; ++k) {
+        dg_pre[b * K + k] = g_soft[b * K + k] * (dg[b * K + k] - dot);
+    }
+    __threadfence_system();
+}
+
 extern "C" __global__ void moe_expert_util_ema_update(
     const float* __restrict__ gate,            /* [B, K] */
     float*       __restrict__ isv,             /* [ISV_TOTAL_DIM] */

crates/ml/src/trainers/dqn/fused_training.rs

@@ -415,6 +415,7 @@ impl FusedTrainingCtx {
             market_dim: 42, // Always 42 base market features — OFI features bypass bottleneck via portfolio_dim
             total_epochs: hyperparams.epochs,
             replay_seed_steps: hyperparams.replay_seed_steps,
+            moe_lambda: hyperparams.moe_lambda.unwrap_or(0.01),
         };
 
         // Create weight set pointer views AFTER GpuDqnTrainer is constructed below.

crates/ml/src/trainers/dqn/trainer/training_loop.rs

@@ -2978,6 +2978,15 @@ impl DQNTrainer {
                         tracing::warn!("Plan 4 Task 6 Commit B aux_heads_loss_ema launch failed: {e}");
                     }
                 }
+
+                // Phase 3 T3.5: MoE expert-utilisation EMA + gate entropy EMA.
+                // Reads `moe_gate_softmax_buf [B, K]` (valid after the captured
+                // forward graph ran) and EMA-updates ISV[118..127).
+                if let Some(ref fused) = self.fused_ctx {
+                    if let Err(e) = fused.trainer().launch_moe_expert_util_ema(ema_alpha) {
+                        tracing::warn!("Phase 3 T3.5 launch_moe_expert_util_ema failed: {e}");
+                    }
+                }
             }
 
             // Plan 4 Task 6 Commit B: refresh the aux-loss weight before the
@@ -3305,6 +3314,31 @@ impl DQNTrainer {
                 );
             }
 
+            // Phase 3 T3.5: MoE expert-utilisation + gate entropy HEALTH_DIAG.
+            {
+                use crate::cuda_pipeline::gpu_dqn_trainer::{
+                    MOE_EXPERT_UTIL_EMA_BASE, MOE_GATE_ENTROPY_EMA_INDEX,
+                };
+                let (moe_utils, moe_ent) =
+                    if let Some(ref fused) = self.fused_ctx {
+                        let trainer = fused.trainer();
+                        let utils: Vec<f32> = (0..8)
+                            .map(|k| trainer.read_isv_signal_at(MOE_EXPERT_UTIL_EMA_BASE + k))
+                            .collect();
+                        let ent = trainer.read_isv_signal_at(MOE_GATE_ENTROPY_EMA_INDEX);
+                        (utils, ent)
+                    } else {
+                        (vec![0.0f32; 8], 0.0f32)
+                    };
+                tracing::info!(
+                    "HEALTH_DIAG[{}]: aux_moe [util={:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3} ent={:.3}]",
+                    epoch,
+                    moe_utils[0], moe_utils[1], moe_utils[2], moe_utils[3],
+                    moe_utils[4], moe_utils[5], moe_utils[6], moe_utils[7],
+                    moe_ent,
+                );
+            }
+
             // C1/P1: propagate health to GPU replay buffer for diversity-weighted priorities.
             {
                 let mut agent = self.agent.write().await;

docs/dqn-wire-up-audit.md

@@ -2,6 +2,22 @@
 
 **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
 
+MoE Phase 3 wire-up T3.1–T3.7 (2026-04-27): MoE fully wired into
+production training path. Forward: gate (state[B,128]→64→8→softmax) +
+8 expert MLPs (h_s1[B,256]→64→256) + `moe_mixture_forward` → replaces
+save_h_s2 in `submit_forward_ops_ddqn`. Load-balance: `moe_load_balance_loss`
++ `moe_load_balance_reduce` + SAXPY into total_loss_dev_ptr (T3.4). ISV
+producer: `launch_moe_expert_util_ema` per-step in training_loop.rs (T3.5).
+Backward: `moe_mixture_backward` (de_k=g·dh_s2) + `moe_dgate_reduce`
+(dg=Σe·dh_s2) + `moe_softmax_backward` + cuBLAS SGEMM backward through gate
+W2/W1 and 8 experts W2/W1 — all into params_buf grad slots [127..163) (T3.6).
+Adam step inherits gate+expert updates automatically (params_buf uniform).
+HEALTH_DIAG aux_moe line emitted per epoch from ISV[118..127) (T3.7).
+Smoke test: 3/3 folds pass, 728s. Gate differentiated: expert-2 reached
+32.3% utilization (others 9.7%) by fold-2 epoch-4; entropy 1.611 < ln(8).
+fused_training.rs: added moe_lambda field to GpuDqnTrainConfig init from
+hyperparams.moe_lambda.unwrap_or(0.01).
+
 MoE expert util EMA T2.4 (2026-04-27): `moe_expert_util_ema_update`
 single-thread cold-path-cadence kernel writes 8 per-expert utilization
 EMAs (ISV[118..126)) + gate-entropy EMA (ISV[126]) with α=0.05. Same shape