# MoE Regime Redesign Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Replace vestigial `RegimeConditionalDQN` (3 hardcoded heads, only `trending_head` trains) with a Mixture-of-Experts policy network (K=8 experts, learned gate) integrated into the existing single-DQN production training path. Atomic deletion of all regime-conditional infrastructure with no fallback paths.

**Architecture:** Shared GRN trunk → 8 small expert MLPs (256→64→256 bottleneck) → learned gating network (state[STATE_DIM=128]→64→8 softmax) → mixed `h_s2` → existing 4 branching heads + C51 + IQN dual head. Soft full mixture (no top-k hardcoding); anti-collapse load-balancing aux loss `λ·K·Σ_k(mean_b g[b,k])²` with `λ=0.01` configurable.

**Tech Stack:** Rust 1.85+, CUDA 12.4, cuBLAS, cudarc, safetensors, GPU PER replay, CUDA Graph capture, Argo Workflows on Scaleway L40S/H100 K8s.

**Spec:** [`docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`](../specs/2026-04-27-moe-regime-redesign-design.md). Read before starting.

**Pearl:** [`pearl_learned_gate_subsumes_handcoded.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_learned_gate_subsumes_handcoded.md). Load-bearing rationale.

**Memory-transfer constraint (`feedback_no_htod_htoh_only_mapped_pinned.md`):** Never use HtoD (`memcpy_stod`) or HtoH (defensive Vec→Vec copies) anywhere in this plan — including test wrappers. The only allowed CPU↔GPU path is **mapped pinned memory** via `cuMemHostAlloc(DEVICEMAP | PORTABLE)`. Reference implementation: `MappedF32Buffer` / `MappedI32Buffer` in `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (used by Test 0.F). Every kernel input that originates CPU-side allocates a mapped buffer, the CPU writes to `host_ptr`, the kernel reads via `dev_ptr` — no copy step. Every kernel output that the CPU needs to read is similarly written by the kernel via `dev_ptr` and read CPU-side via `host_ptr` (with `__threadfence_system()` on the kernel side and `read_volatile` on the CPU side).

---

## File Structure

### New files

| Path | Responsibility |
|---|---|
| `crates/ml/src/cuda_pipeline/moe_kernels.cu` | 4 new CUDA kernels: `moe_mixture_forward`, `moe_mixture_backward`, `moe_load_balance_loss`, `moe_expert_util_ema_update` |
| `crates/ml/src/cuda_pipeline/gpu_moe_head.rs` | Rust wrapper around the MoE kernels; analogous to `gpu_iqn_head.rs` shape |
| `crates/ml/tests/moe_kernels_test.rs` | Layer 1 unit tests for all 4 new kernels |
| `crates/ml/src/trainers/dqn/smoke_tests/moe_gate_differentiation.rs` | Layer 3 gate-differentiation validation |

### Files modified (substantial)

| Path | Change |
|---|---|
| `crates/ml-dqn/src/count_bonus.rs` | Add `bonuses_branched_f32` write-into API (Phase 0) |
| `crates/ml-dqn/src/dqn.rs` | Strip 6 `use_*` flags + tighten count_bonus API (Phase 0); add `expert_*`, `gate_*` weight tensor fields (Phase 1) |
| `crates/ml-dqn/src/lib.rs` | Drop `RegimeConditional*` exports (Phase 4) |
| `crates/ml/src/trainers/dqn/config.rs` | `DQNAgentType` rewrite — wrap `DQN` directly (Phase 4); add `moe_lambda` to `DQNHyperparameters` (Phase 1) |
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | MoE producer launches + `aux_moe` HEALTH_DIAG line (Phase 3) |
| `crates/ml/src/trainers/dqn/trainer/action.rs` | Caller migration from regime-routed to direct (Phase 0 + Phase 4) |
| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | Wire MoE config into `GpuDqnTrainConfig` (Phase 1) |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Construct `GpuMoeHead`; wire into forward/backward graphs (Phase 3) |
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Extend `params_buf` layout for gate + 8 experts; new ISV slot consts; integrate `GpuMoeHead` in forward/backward (Phase 1 + Phase 3) |
| `crates/ml/src/trainers/dqn/state_reset_registry.rs` | Register 9 new ISV slots with reset values (Phase 1) |

### Files deleted (atomic with Phase 4)

| Path | Why |
|---|---|
| `crates/ml-dqn/src/regime_conditional.rs` | Vestigial 3-head wrapper; MoE subsumes purpose |

### Files removed/modified for cleanup (Phase 0)

| Path | Change |
|---|---|
| `services/trading_service/src/services/dqn_model.rs` | Drop `config.use_count_bonus = false` line |
| `crates/ml/tests/dqn_action_collapse_fix_test.rs` | Drop `config.use_count_bonus = true` line + `!use_iqn` assertion |
| `crates/ml/tests/dqn_inference_test.rs` | Drop `config.use_iqn = false` |
| `crates/ml/src/trainers/dqn/distributional_q_tests.rs` | Drop `use_iqn=*` references in docstrings + Test 0.F output line |
| `crates/ml/src/training_profile.rs` | Drop comment referencing `use_branching` |
| `crates/ml/examples/evaluate_baseline.rs` | Drop `iqn_num_quantiles` references |
| `docs/dqn-wire-up-audit.md` | Append Phase 0 + Phase 4 entries per Invariant 7 |

---

## Phase 0 — Precondition: use_* flag cleanup + count_bonus refactor

Phase 0 is **partially complete in `stash@{0}`**: the bulk of the edits to `count_bonus.rs`, `dqn.rs`, `ml/src/trainers/dqn/config.rs`, and `action.rs` are stashed. Phase 0 finishes that work and commits it.

### Task 0.1 — Restore stashed working tree and verify state

**Files:**
- Modify: `crates/ml-dqn/src/count_bonus.rs`
- Modify: `crates/ml-dqn/src/dqn.rs`
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/action.rs`

- [ ] **Step 1: Verify stash exists**

```bash
git stash list | head -3
```
Expected: first line is `stash@{0}: On main: WIP: use_* + count_bonus refactor (parking before spec commit)`.

- [ ] **Step 2: Pop the stash**

```bash
git stash pop stash@{0}
```
Expected output: `Changes not staged for commit:` listing the 4 files (count_bonus.rs, dqn.rs, ml/src/trainers/dqn/config.rs, action.rs). On any merge conflict, abort: `git stash apply --reset` then re-do edits manually following Tasks 0.2–0.10 from scratch.

- [ ] **Step 3: cargo check the ml-dqn crate**

```bash
SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -10
```
Expected: `Finished 'dev' profile [unoptimized + debuginfo] target(s)`. If errors, fix per the error message before proceeding.

- [ ] **Step 4: cargo check the workspace**

```bash
SQLX_OFFLINE=true cargo check --workspace --tests 2>&1 | tail -15
```
Expected: 0 errors. Warnings about `unused` from pre-existing dead code are acceptable. If errors, identify the call site (likely in `services/`, `tests/`, or `examples/`) and proceed to Task 0.2.

### Task 0.2 — Audit remaining use_* references workspace-wide

**Files:**
- Read: workspace-wide grep results

- [ ] **Step 1: Grep for any remaining `use_dueling`, `use_distributional`, `use_noisy_nets`, `use_branching`, `use_count_bonus`, `use_cvar_action_selection` references**

```bash
grep -rn "use_dueling\|use_distributional\|use_noisy_nets\|use_branching\|use_count_bonus\|use_cvar_action_selection" --include="*.rs" --include="*.md" 2>/dev/null
```

- [ ] **Step 2: Classify findings**

For each match, determine:
- **In a doc comment / Markdown reference / proto definition (e.g., `services/ml_training_service/src/proto/`)**: leave alone — proto-generated `use_dueling` is a separate gRPC contract, not the DQNConfig field.
- **In a test, service code, example, or non-proto Rust source**: must be cleaned up.

Build a TODO list of files needing edits. Expected at most:
- `crates/ml/tests/dqn_action_collapse_fix_test.rs` (line 300)
- `crates/ml/tests/dqn_inference_test.rs` (line 136)
- `services/trading_service/src/services/dqn_model.rs` (line 59)
- `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (3 lines)
- `crates/ml/src/trainers/dqn/config.rs` (line ~1031 doc comment)
- `crates/ml/src/training_profile.rs` (line ~870)

If the stash already cleaned these (since they were already in the WIP commit drafted earlier), `grep` returns nothing in non-proto files and you skip to Task 0.7.

### Task 0.3 — Strip `use_count_bonus = true` from collapse_fix test

**Files:**
- Modify: `crates/ml/tests/dqn_action_collapse_fix_test.rs:300`

- [ ] **Step 1: Read the test setup block**

```bash
sed -n '295,305p' crates/ml/tests/dqn_action_collapse_fix_test.rs
```

- [ ] **Step 2: Delete the use_count_bonus line**

Open `crates/ml/tests/dqn_action_collapse_fix_test.rs` and remove the line:
```rust
config.use_count_bonus = true;
```
(field no longer exists; production records actions unconditionally now).

- [ ] **Step 3: Verify test still compiles**

```bash
SQLX_OFFLINE=true cargo check -p ml --tests 2>&1 | grep -E "error|dqn_action_collapse" | head -5
```
Expected: no errors. (Test isn't run, just compiled.)

### Task 0.4 — Strip `use_iqn = false` from inference test

**Files:**
- Modify: `crates/ml/tests/dqn_inference_test.rs:136`

- [ ] **Step 1: Read the setup block**

```bash
sed -n '130,140p' crates/ml/tests/dqn_inference_test.rs
```

- [ ] **Step 2: Delete the line**

Remove `config.use_iqn = false;` (field no longer exists; production runs IQN unconditionally via cuda_pipeline).

- [ ] **Step 3: Verify**

```bash
SQLX_OFFLINE=true cargo check -p ml --tests 2>&1 | tail -5
```

### Task 0.5 — Strip `use_count_bonus` from trading_service

**Files:**
- Modify: `services/trading_service/src/services/dqn_model.rs:59`

- [ ] **Step 1: Read context**

```bash
sed -n '55,65p' services/trading_service/src/services/dqn_model.rs
```

- [ ] **Step 2: Delete the line**

Remove the line `config.use_count_bonus = false;` and the trailing comment `// No UCB exploration bonus on Q-values`. Replace with a single line `config.count_bonus_coefficient = 0.0;` since the production kill-switch is the numeric coefficient (per `feedback_no_feature_flags.md`).

- [ ] **Step 3: Verify**

```bash
SQLX_OFFLINE=true cargo check -p trading_service 2>&1 | tail -5
```

### Task 0.6 — Strip stale doc comments

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs` (doc comment)
- Modify: `crates/ml/src/training_profile.rs` (informational comment)
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (Test 0.F docstrings + output line)

- [ ] **Step 1: Find all stale doc references**

```bash
grep -n "use_branching\|use_count_bonus\|use_iqn" crates/ml/src/trainers/dqn/config.rs crates/ml/src/training_profile.rs crates/ml/src/trainers/dqn/distributional_q_tests.rs 2>/dev/null
```

- [ ] **Step 2: For each match, rewrite the comment to drop the flag reference**

Pattern: `Only used when use_branching is true` → delete that sentence (branching is unconditional).

Pattern (from distributional_q_tests.rs Test 0.F docstring): `with use_iqn=false, so it has no IQN tensors` → already updated in earlier work to drop the legacy framing. If still present, rewrite per the Test 0.F's current "C51-only scope" docstring (already applied in commit `da632446c`).

- [ ] **Step 3: Verify no remaining references**

```bash
grep -rn "use_branching\|use_count_bonus" crates/ml/src/ crates/ml-dqn/src/ services/ 2>/dev/null | grep -v "/proto/"
```
Expected: empty.

### Task 0.7 — Final cargo check + Test 0.F behavioral parity

**Files:**
- Verify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (Test 0.F)

- [ ] **Step 1: cargo check workspace**

```bash
SQLX_OFFLINE=true cargo check --workspace --tests 2>&1 | tail -10
```
Expected: 0 errors. Pre-existing warnings OK.

- [ ] **Step 2: Run Test 0.F to confirm bit-identical kernel outputs**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib distributional_q_tests::test_0f -- --ignored --nocapture 2>&1 | tail -20
```

Expected output exactly matches pre-cleanup outputs from commit `da632446c`:
- `σ_C51 per direction: [141.74062, 142.49536, 140.4468, 140.72096]`
- `argmax pick: 2`
- `Thompson counts (S/H/L/F): [92, 298, 341, 269], P(active)=0.4330`
- Test panics with the same bias-hypothesis assertion message as before.

If outputs differ even slightly, the use_* removal accidentally changed behavior somewhere — `git diff HEAD` and investigate.

### Task 0.8 — Update audit doc per Invariant 7

**Files:**
- Modify: `docs/dqn-wire-up-audit.md` (append entry at the top)

- [ ] **Step 1: Read the existing audit doc header structure**

```bash
sed -n '1,10p' docs/dqn-wire-up-audit.md
```

- [ ] **Step 2: Append a new dated entry at the top describing Phase 0**

Edit `docs/dqn-wire-up-audit.md`. After the existing first dated entry, prepend:

```markdown
use_* flag family removal + count_bonus API tightening (2026-04-27): atomic
deletion of 6 dead `use_*` boolean flags from `DQNConfig` per
`feedback_no_feature_flags.md`. (1) `use_dueling`, `use_distributional`,
`use_noisy_nets`, `use_branching` were pure-cosmetic always-on flags; only
metadata strings remained. Their `if true /* always on */` patterns at 4 call
sites in `dqn.rs::select_action`, `select_action_with_confidence`,
`select_action_inference`, `q_values_for_batch` collapsed to the live
branching-only arm; the dead else arms (legacy non-branching code paths
referencing 5-flat `ExposureLevel` layout) deleted. (2) `use_count_bonus` was
redundant with the live `count_bonus_coefficient` numeric kill-switch; gating
removed, recording made unconditional, single coefficient drives behavior.
(3) `use_cvar_action_selection` was dead post-`use_iqn` cleanup (only
consumers were inside the IQN arms removed in commit `da632446c`); field gone,
`cvar_alpha` retained for cuda_pipeline CVaR loss usage and as the future dial
for CVaR action selection on the live C51 distribution. (4) `count_bonus.rs`
gained `bonuses_branched_f32(&self, exp_out: &mut [f32], ord_out: &mut [f32],
urg_out: &mut [f32])` write-into API alongside the existing `Vec<f64>`; the
production `DQN::get_count_bonuses_branched()` returns fixed `([f32; 4], [f32;
3], [f32; 3])` arrays — allocation-free, f32 throughout, fed directly to the
GPU action selector. Bit-identical Test 0.F outputs vs pre-cleanup confirm
zero behavioral change (legacy `use_*` arms were unreachable, as expected).
This commit is the precondition for the MoE regime redesign per
`docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`.
```

- [ ] **Step 3: Verify the audit doc passes the pre-commit hook**

```bash
git add docs/dqn-wire-up-audit.md
git diff --cached --stat docs/dqn-wire-up-audit.md
```
Expected: 1 file changed, ~25 lines inserted.

### Task 0.9 — Commit Phase 0

**Files:**
- All Phase 0 modified files staged

- [ ] **Step 1: Stage all Phase 0 changes**

```bash
git status --short
git add crates/ml-dqn/src/count_bonus.rs crates/ml-dqn/src/dqn.rs \
        crates/ml/src/trainers/dqn/config.rs \
        crates/ml/src/trainers/dqn/trainer/action.rs \
        crates/ml/tests/dqn_action_collapse_fix_test.rs \
        crates/ml/tests/dqn_inference_test.rs \
        services/trading_service/src/services/dqn_model.rs \
        crates/ml/src/trainers/dqn/distributional_q_tests.rs \
        crates/ml/src/training_profile.rs \
        crates/ml/src/trainers/dqn/distributional_q_tests.rs \
        docs/dqn-wire-up-audit.md
git status --short
```

- [ ] **Step 2: Commit with descriptive message**

```bash
git commit -m "$(cat <<'EOF'
refactor(dqn): strip dead use_* flags + tighten count_bonus API

Atomic removal of 6 boolean feature flags from DQNConfig per
`feedback_no_feature_flags.md`, all of which gated dead or redundant code:

- `use_dueling`, `use_distributional`, `use_noisy_nets`, `use_branching`:
  pure-cosmetic always-on flags. Only metadata strings remained.
  4 `if true /* always on */` blocks in dqn.rs collapsed to live arms;
  dead else arms (legacy non-branching code paths, 5-flat ExposureLevel)
  deleted.
- `use_count_bonus`: redundant with `count_bonus_coefficient` (the live
  numeric kill-switch). Recording made unconditional; coefficient is the
  single dial.
- `use_cvar_action_selection`: dead post-use_iqn cleanup (only consumers
  were inside the deleted IQN arms). cvar_alpha retained for cuda_pipeline
  CVaR loss usage.

count_bonus.rs API tightened: `bonuses_branched_f32(&self,
exp_out: &mut [f32], ord_out: &mut [f32], urg_out: &mut [f32])` write-into
API alongside the existing Vec<f64> form (kept for tests). Production
DQN::get_count_bonuses_branched() returns fixed `([f32; 4], [f32; 3],
[f32; 3])` arrays — allocation-free, f32 throughout, fed directly to GPU
action selector via stack-allocated buffers.

Test 0.F outputs bit-identical vs pre-cleanup (sigma_C51, argmax,
Thompson counts) — confirms legacy use_* arms were unreachable as
expected.

`docs/dqn-wire-up-audit.md` updated per Invariant 7.

Precondition for the MoE regime redesign per
`docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`.

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

- [ ] **Step 3: Verify commit landed**

```bash
git log --oneline -3
```
Expected: top commit is `refactor(dqn): strip dead use_* flags + tighten count_bonus API`.

---

## Phase 1 — MoE infrastructure (additive, no behavior change)

Phase 1 adds the gate + expert weight tensors to the params layout, the new ISV slots, and the `moe_lambda` hyperparameter — but does NOT yet wire the gate into the forward graph. After Phase 1, the workspace builds cleanly with extra-but-unused parameters; behavior is unchanged.

### Task 1.1 — Add `moe_lambda` hyperparameter

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs` (DQNHyperparameters struct + Default impl)

- [ ] **Step 1: Locate DQNHyperparameters struct**

```bash
grep -n "pub struct DQNHyperparameters\|pub struct DQNHyperparams" crates/ml/src/trainers/dqn/config.rs
```

- [ ] **Step 2: Add the `moe_lambda` field**

Add to the struct (alongside `cql_alpha`, `iqn_lambda`, `count_bonus_coefficient`):

```rust
/// Mixture-of-Experts load-balancing aux loss weight.
/// Default: 0.01 (anti-collapse only — prevents init-noise-dominated
/// single-expert lock-in without forcing uniform utilization).
/// Hyperopt search range: [0.001, 0.1].
pub moe_lambda: Option<f32>,
```

- [ ] **Step 3: Update Default**

```rust
moe_lambda: Some(0.01),
```

- [ ] **Step 4: Verify cargo check**

```bash
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/config.rs
git commit -m "feat(dqn): add moe_lambda hyperparameter (default 0.01)

Anti-collapse load-balancing aux loss weight for the upcoming Mixture-of-
Experts redesign. Configurable via hyperopt; not a kernel constant.

Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §4.3.

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

### Task 1.2 — Add 9 new ISV slot constants

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (existing pattern: `pub const ISV_*_INDEX`)

- [ ] **Step 1: Locate existing ISV slot constant block**

```bash
grep -n "pub const AUX_LABEL_SCALE_EMA_INDEX\|pub const ISV_TOTAL_DIM" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
```

- [ ] **Step 2: Add 9 new constants after AUX_LABEL_SCALE_EMA_INDEX**

```rust
/// Mixture-of-Experts per-expert utilization EMA, α=0.05.
/// 8 contiguous slots [118..126).
/// Producer: `moe_expert_util_ema_update` GPU kernel (cold-path cadence).
/// Reset value at fold boundary: 1.0 / 8.0 (uniform-gate initial).
pub const MOE_EXPERT_UTIL_EMA_BASE: usize = 118;
pub const MOE_EXPERT_UTIL_EMA_COUNT: usize = 8;

/// Entropy of batch-averaged gate distribution, EMA α=0.05.
/// Reset value at fold boundary: ln(8) ≈ 2.0794.
pub const MOE_GATE_ENTROPY_EMA_INDEX: usize = 126;
```

- [ ] **Step 3: Bump ISV_TOTAL_DIM**

Find the existing `pub const ISV_TOTAL_DIM` line. Change from `118` to `127`.

- [ ] **Step 4: Verify cargo check (no consumers yet, just constants)**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat(dqn): reserve ISV slots 118-126 for MoE monitoring

8 slots for per-expert utilization EMA + 1 slot for gate-entropy EMA.
ISV_TOTAL_DIM 118 -> 127. Not yet consumed; producers + consumers land
in subsequent commits.

Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §4.6.

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

### Task 1.3 — Register fold-boundary reset for MoE ISV slots

**Files:**
- Modify: `crates/ml/src/trainers/dqn/state_reset_registry.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (reset dispatch arm)

- [ ] **Step 1: Locate state_reset_registry pattern**

```bash
grep -n "isv_aux_label_scale_ema\|FoldReset" crates/ml/src/trainers/dqn/state_reset_registry.rs | head -5
```

- [ ] **Step 2: Add 9 entries to the registry**

Following the pattern of `isv_aux_label_scale_ema` entry, append:

```rust
// Mixture-of-Experts per-expert utilization (slots 118..126).
// Reset to 1/K = 1/8 = 0.125 (uniform-gate initial).
StateResetEntry {
    name: "isv_moe_expert_util_ema",
    timing: ResetTiming::FoldBoundary,
    description: "ISV[MOE_EXPERT_UTIL_EMA_BASE..MOE_EXPERT_UTIL_EMA_BASE+8) — \
                  per-expert gate-weight EMA (α=0.05); GPU moe_expert_util_ema_update \
                  kernel fills; reset to uniform 1/K=0.125 at fold boundary",
},
StateResetEntry {
    name: "isv_moe_gate_entropy_ema",
    timing: ResetTiming::FoldBoundary,
    description: "ISV[MOE_GATE_ENTROPY_EMA_INDEX=126] — entropy of batch-mean gate \
                  distribution, EMA α=0.05; reset to ln(8) ≈ 2.0794 at fold boundary \
                  (max entropy at uniform initial gate)",
},
```

- [ ] **Step 3: Add reset dispatch arms in training_loop.rs**

Find the existing `match name` arm for `"isv_aux_label_scale_ema"` (around line 4813). Add:

```rust
"isv_moe_expert_util_ema" => {
    if let Some(ref fused) = self.fused_ctx {
        for k in 0..crate::cuda_pipeline::gpu_dqn_trainer::MOE_EXPERT_UTIL_EMA_COUNT {
            fused.write_isv_signal_at(
                crate::cuda_pipeline::gpu_dqn_trainer::MOE_EXPERT_UTIL_EMA_BASE + k,
                1.0_f32 / crate::cuda_pipeline::gpu_dqn_trainer::MOE_EXPERT_UTIL_EMA_COUNT as f32,
            );
        }
    }
}
"isv_moe_gate_entropy_ema" => {
    if let Some(ref fused) = self.fused_ctx {
        fused.write_isv_signal_at(
            crate::cuda_pipeline::gpu_dqn_trainer::MOE_GATE_ENTROPY_EMA_INDEX,
            (crate::cuda_pipeline::gpu_dqn_trainer::MOE_EXPERT_UTIL_EMA_COUNT as f32).ln(),
        );
    }
}
```

- [ ] **Step 4: Verify cargo check**

```bash
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/state_reset_registry.rs \
        crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat(dqn): register fold-boundary reset for MoE ISV slots

Slots 118..126 reset to 1/K=0.125 (uniform gate initial); slot 126 resets
to ln(8)=2.0794 (max entropy). Producers land in Phase 2; consumers in
Phase 3.

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

### Task 1.4 — Add gate + expert weight tensors to DQN struct

**Files:**
- Modify: `crates/ml-dqn/src/dqn.rs` (DQN struct fields + initialization)

- [ ] **Step 1: Locate DQN struct definition**

```bash
grep -n "pub struct DQN " crates/ml-dqn/src/dqn.rs
```

- [ ] **Step 2: Add MoE field group**

After the existing `branching_q_network` field, add:

```rust
/// Mixture-of-Experts gate subnetwork: state[42] -> 64 -> 8 logits -> softmax.
/// Initialized to zero so initial g(s) = uniform 1/K for every state.
pub moe_gate: Option<super::moe::MoeGate>,

/// 8 expert MLPs: each is h_s1[256] -> 64 -> h_s2[256] bottleneck.
/// Independent Xavier init per expert for initial differentiation.
pub moe_experts: Option<Vec<super::moe::MoeExpert>>,
```

(`super::moe` module created in Task 1.5.)

- [ ] **Step 3: Initialize fields to None in DQN::new_with_stream**

In the existing `Ok(Self { ... })` block, add:
```rust
moe_gate: None,        // Phase 3 wire-up replaces with Some(MoeGate::new(...))
moe_experts: None,     // Phase 3 wire-up replaces with Some(vec![MoeExpert; 8])
```

(For Phase 1, fields exist but are unused.)

- [ ] **Step 4: Verify cargo check**

```bash
SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -5
```

The `super::moe::MoeGate` / `MoeExpert` types don't exist yet — expected error. Proceed to Task 1.5.

### Task 1.5 — Create MoE module skeleton in ml-dqn

**Files:**
- Create: `crates/ml-dqn/src/moe.rs`
- Modify: `crates/ml-dqn/src/lib.rs` (add `pub mod moe;`)

- [ ] **Step 1: Create the module file**

Create `crates/ml-dqn/src/moe.rs` with the type signatures:

```rust
//! Mixture-of-Experts components for the DQN policy network.
//!
//! Replaces the vestigial RegimeConditionalDQN with a learned-gate
//! K=8 expert mixture per the spec at
//! docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md.

use ml_core::cuda_autograd::GpuTensor;
use ml_core::MLError;
use std::sync::Arc;
use cudarc::driver::CudaStream;

/// Gate subnetwork: state[42] -> 64 -> 8 -> softmax over experts.
/// Parameters initialized to zero so g(s) = uniform 1/K initially.
pub struct MoeGate {
    /// Linear 1 weights [42, 64], zero-init.
    pub w1: GpuTensor,
    /// Linear 1 bias [64], zero-init.
    pub b1: GpuTensor,
    /// Linear 2 weights [64, 8], zero-init.
    pub w2: GpuTensor,
    /// Linear 2 bias [8], zero-init.
    pub b2: GpuTensor,
}

impl MoeGate {
    /// Construct with zero-initialized weights so initial gate = uniform 1/K.
    pub fn new(stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let w1 = GpuTensor::zeros(vec![42, 64], &stream)?;
        let b1 = GpuTensor::zeros(vec![64], &stream)?;
        let w2 = GpuTensor::zeros(vec![64, 8], &stream)?;
        let b2 = GpuTensor::zeros(vec![8], &stream)?;
        Ok(Self { w1, b1, w2, b2 })
    }
}

/// Single expert MLP: h_s1[256] -> 64 -> h_s2[256] with bottleneck.
/// Each expert has independent Xavier init for initial differentiation.
pub struct MoeExpert {
    /// Linear 1 weights [256, 64], Xavier init.
    pub w1: GpuTensor,
    /// Linear 1 bias [64], zero-init.
    pub b1: GpuTensor,
    /// Linear 2 weights [64, 256], Xavier init.
    pub w2: GpuTensor,
    /// Linear 2 bias [256], zero-init.
    pub b2: GpuTensor,
}

impl MoeExpert {
    /// Construct with Xavier-initialized weights.
    pub fn new(stream: Arc<CudaStream>, seed: u64) -> Result<Self, MLError> {
        // Xavier init: stddev = sqrt(2 / (fan_in + fan_out))
        let w1 = GpuTensor::xavier(vec![256, 64], seed, &stream)?;
        let b1 = GpuTensor::zeros(vec![64], &stream)?;
        let w2 = GpuTensor::xavier(vec![64, 256], seed.wrapping_add(1), &stream)?;
        let b2 = GpuTensor::zeros(vec![256], &stream)?;
        Ok(Self { w1, b1, w2, b2 })
    }
}

/// K=8 experts (compile-time constant).
pub const MOE_NUM_EXPERTS: usize = 8;

/// Expert bottleneck width (compile-time constant).
pub const MOE_EXPERT_BOTTLENECK: usize = 64;

/// Gate hidden width (compile-time constant).
pub const MOE_GATE_HIDDEN: usize = 64;
```

- [ ] **Step 2: Add module to lib.rs**

Add `pub mod moe;` near the other `pub mod` declarations.

- [ ] **Step 3: Add re-exports**

```rust
pub use moe::{MoeGate, MoeExpert, MOE_NUM_EXPERTS, MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN};
```

- [ ] **Step 4: Verify cargo check (xavier helper may not exist; if so, fall back to manual init)**

```bash
SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -10
```

If `GpuTensor::xavier` doesn't exist, replace with manual init pattern matching existing code in `crates/ml-dqn/src/dqn.rs::xavier_init_*` helpers.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-dqn/src/moe.rs crates/ml-dqn/src/lib.rs crates/ml-dqn/src/dqn.rs
git commit -m "feat(moe): add MoeGate and MoeExpert struct skeletons

Module crates/ml-dqn/src/moe.rs introduces:
- MoeGate: 42->64->8 softmax gating network, zero-init for uniform start.
- MoeExpert: 256->64->256 bottleneck, Xavier init per-expert seed.

DQN struct gains optional moe_gate / moe_experts fields, initialized to
None for now — wire-up in Phase 3.

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

### Task 1.6 — Extend params_buf layout in GpuDqnTrainer

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (params layout, Adam buffers, target sync)

- [ ] **Step 1: Locate the existing param tensor offset table**

```bash
grep -n "param_offsets\|PARAM_OFFSET\|fn build_param_layout" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | head -10
```

- [ ] **Step 2: Append MoE tensor offsets to the existing layout enum/table**

After the last existing param tensor entry, add 36 new entries:

```rust
// MoE gate weights (4 tensors, ~3.3k params total):
GateW1,        // [42, 64]
GateB1,        // [64]
GateW2,        // [64, 8]
GateB2,        // [8]
// MoE expert weights (8 experts × 4 tensors each = 32 tensors, ~264k params):
Expert0W1, Expert0B1, Expert0W2, Expert0B2,
Expert1W1, Expert1B1, Expert1W2, Expert1B2,
Expert2W1, Expert2B1, Expert2W2, Expert2B2,
Expert3W1, Expert3B1, Expert3W2, Expert3B2,
Expert4W1, Expert4B1, Expert4W2, Expert4B2,
Expert5W1, Expert5B1, Expert5W2, Expert5B2,
Expert6W1, Expert6B1, Expert6W2, Expert6B2,
Expert7W1, Expert7B1, Expert7W2, Expert7B2,
```

(Adapt to whatever the existing pattern is — Rust enum vs offset constants vs runtime table.)

- [ ] **Step 3: Update layout fingerprint hash to include new offsets**

Find the existing `LAYOUT_FINGERPRINT_SEED` / `layout_fingerprint_seed` function. Add the new shapes to the hash inputs.

- [ ] **Step 4: Update Adam state buffers (m_buf, v_buf) to extend by the same delta**

Locate `m_buf`, `v_buf` allocations. The total params count grows by ~268k. Adjust the alloc size accordingly (typically computed from total param count, so this may auto-extend if computed from the layout).

- [ ] **Step 5: Update target_params_buf to mirror**

- [ ] **Step 6: Update gradient scratch buffers**

- [ ] **Step 7: Verify cargo check**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10
```

- [ ] **Step 8: Verify Test 0.F still passes (architecture hash will fail loudly per design)**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib distributional_q_tests::test_0f -- --ignored --nocapture 2>&1 | tail -10
```

Expected: layout fingerprint mismatch error when loading the existing smoke checkpoint. This is **the no-fallback contract working**. Either:
- (a) Re-train the local smoke checkpoint per Phase 5 (preferred), or
- (b) `#[ignore]` the test temporarily — the Phase 5 L40S validation is the real gate. Add a comment marking the temporary ignore as Phase 5 dependency.

- [ ] **Step 9: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "feat(moe): extend params_buf layout with 36 new tensors

Gate (4 tensors, ~3.3k params) + 8 experts (32 tensors, ~264k params).
Layout fingerprint hash recomputes — old checkpoints fail to load with
fingerprint-mismatch error per the no-fallback contract in the spec.

Adam state, target_params, gradient scratch all extend in lockstep.

Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.1.

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

---

## Phase 2 — MoE kernels

Phase 2 implements the 4 new CUDA kernels with unit tests. Kernels are written and tested in isolation; not yet wired into training graph. **All test wrappers use mapped pinned memory exclusively** per `feedback_no_htod_htoh_only_mapped_pinned.md` — no `memcpy_stod`, no `memcpy_dtov`, no defensive Vec-to-Vec.

### Task 2.0 — Promote MappedF32Buffer / MappedI32Buffer to shared module

**Files:**
- Create: `crates/ml/src/cuda_pipeline/mapped_pinned.rs`
- Modify: `crates/ml/src/cuda_pipeline/mod.rs` (add `pub mod mapped_pinned;`)
- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (delete local definitions, import from shared)

The existing Test 0.F (commit `da632446c`) defines `MappedF32Buffer` and `MappedI32Buffer` locally in `distributional_q_tests.rs`. They use `cuMemHostAlloc(DEVICEMAP | PORTABLE)` + `cuMemHostGetDevicePointer_v2` — exactly the pattern this plan needs everywhere. Promote to a shared module so all kernel-test wrappers share one implementation; delete the local copies.

- [ ] **Step 1: Locate the existing buffer definitions**

```bash
grep -n "struct MappedF32Buffer\|struct MappedI32Buffer\|impl Mapped" crates/ml/src/trainers/dqn/distributional_q_tests.rs | head -10
```

- [ ] **Step 2: Create the shared module**

Create `crates/ml/src/cuda_pipeline/mapped_pinned.rs` by moving the entire `MappedF32Buffer` and `MappedI32Buffer` blocks (struct + `unsafe impl Send/Sync` + `impl new` + `impl read_all` + `impl Drop`) from `distributional_q_tests.rs` into the new file. Make the structs and methods `pub`.

Add a helper method per buffer type for "write from slice without memcpy":

```rust
impl MappedF32Buffer {
    /// Write CPU-side data via host_ptr — direct memory write, no memcpy.
    /// Caller must ensure `slice.len() <= self.len`.
    pub fn write_from_slice(&self, slice: &[f32]) {
        assert!(slice.len() <= self.len, "write_from_slice overflow");
        // Safety: host_ptr is valid for self.len f32s; writes are coherent
        // with subsequent kernel reads via dev_ptr (cuMemHostAlloc DEVICEMAP).
        unsafe {
            for (i, &v) in slice.iter().enumerate() {
                std::ptr::write_volatile(self.host_ptr.add(i), v);
            }
        }
    }
}
```

(Same for `MappedI32Buffer` with `i32`.)

- [ ] **Step 3: Add module to mod.rs**

```rust
pub mod mapped_pinned;
```

- [ ] **Step 4: Delete the local copies from distributional_q_tests.rs and import from shared module**

```rust
use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
```

- [ ] **Step 5: Verify Test 0.F still runs end-to-end with bit-identical output**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib distributional_q_tests::test_0f -- --ignored --nocapture 2>&1 | tail -10
```

Expected: σ_C51 = `[141.74062, 142.49536, 140.4468, 140.72096]`, argmax=2, P(active)=0.4330, panic with bias-hypothesis assertion. Bit-identical to commit `da632446c`'s output.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/mapped_pinned.rs \
        crates/ml/src/cuda_pipeline/mod.rs \
        crates/ml/src/trainers/dqn/distributional_q_tests.rs
git commit -m "refactor(cuda_pipeline): promote MappedF32/I32Buffer to shared module

Moves MappedF32Buffer + MappedI32Buffer from distributional_q_tests.rs
local definitions to crates/ml/src/cuda_pipeline/mapped_pinned.rs so all
kernel test wrappers (Phase 0.F, upcoming MoE tests) share one
implementation.

Adds write_from_slice helper for direct host_ptr write (no memcpy).

Test 0.F bit-identical: sigma_C51, argmax, Thompson counts unchanged.

Per feedback_no_htod_htoh_only_mapped_pinned.md: this is the only
allowed CPU<->GPU path. The shared module makes the mapped-pinned
pattern more discoverable and reduces duplication.

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

### Task 2.1 — moe_mixture_forward kernel

**Files:**
- Create: `crates/ml/src/cuda_pipeline/moe_kernels.cu`
- Create (or extend): `crates/ml/src/cuda_pipeline/gpu_moe_head.rs`
- Create: `crates/ml/tests/moe_kernels_test.rs`

- [ ] **Step 1: Write the failing unit test FIRST (TDD)**

Create `crates/ml/tests/moe_kernels_test.rs`:

```rust
//! Layer 1 unit tests for MoE kernels.
//! Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §8.

use ml::cuda_pipeline::gpu_moe_head::GpuMoeHead;
use std::sync::Arc;

#[test]
#[ignore]  // Requires CUDA; run with --ignored
fn moe_mixture_forward_matches_cpu_reference() {
    let context = ml_core::cuda_autograd::test_helpers::test_context();
    let stream = Arc::clone(context.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: upload, run kernel, download
    let mut 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]
        );
    }
}
```

- [ ] **Step 2: Run test to verify it fails (no kernel yet)**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_mixture_forward_matches_cpu_reference -- --ignored 2>&1 | tail -10
```

Expected: FAIL — `GpuMoeHead` doesn't exist yet.

- [ ] **Step 3: Create the kernel**

Create `crates/ml/src/cuda_pipeline/moe_kernels.cu`:

```cuda
/* MoE mixture forward / backward / load-balance / ISV-producer kernels.
 *
 * 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;
}
```

- [ ] **Step 4: Add kernel compilation to build.rs**

Locate `crates/ml/build.rs`, find the existing `.cu` file list (e.g., `aux_heads_kernel.cu`). Add `moe_kernels.cu` to the same list.

- [ ] **Step 5: Create the Rust wrapper**

Create `crates/ml/src/cuda_pipeline/gpu_moe_head.rs`:

```rust
//! Rust wrapper for MoE GPU kernels.
//! Loaded from cubin at runtime, like other cuda_pipeline heads.

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig};
use ml_core::MLError;
use std::sync::Arc;

static MOE_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));

pub struct GpuMoeHead {
    stream: Arc<CudaStream>,
    moe_mixture_forward: CudaFunction,
    // Other kernels added in subsequent tasks.
}

impl GpuMoeHead {
    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: stage CPU-side data into mapped pinned memory,
    /// run kernel, read GPU-written output via mapped pinned. **No HtoD,
    /// no HtoH** — `feedback_no_htod_htoh_only_mapped_pinned.md`.
    pub fn test_mixture_forward(
        &self,
        expert_outputs: &[f32],
        gate: &[f32],
        b: usize,
        k: usize,
        c: usize,
    ) -> Result<Vec<f32>, MLError> {
        // Allocate mapped pinned buffers (CPU-and-GPU coherent, no copy).
        // Pattern: cuMemHostAlloc(DEVICEMAP|PORTABLE) → cuMemHostGetDevicePointer_v2.
        // Reference: MappedF32Buffer in crates/ml/src/trainers/dqn/distributional_q_tests.rs.
        let expert_buf = unsafe { mapped_f32::MappedF32Buffer::new(k * b * c) }?;
        let gate_buf = unsafe { mapped_f32::MappedF32Buffer::new(b * k) }?;
        let h_s2_buf = unsafe { mapped_f32::MappedF32Buffer::new(b * c) }?;

        // Write CPU-side via host_ptr (no memcpy, just direct memory writes).
        expert_buf.write_from_slice(expert_outputs);
        gate_buf.write_from_slice(gate);

        let total = b * c;
        let block = 256_u32;
        let grid = ((total + block as usize - 1) / block as usize) as u32;

        unsafe {
            self.stream
                .launch_builder(&self.moe_mixture_forward)
                .arg(&expert_buf.dev_ptr)
                .arg(&gate_buf.dev_ptr)
                .arg(&h_s2_buf.dev_ptr)
                .arg(&(b as i32))
                .arg(&(k as i32))
                .arg(&(c as 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}")))?;
        }
        // Stream sync ensures __threadfence_system() side effects are visible
        // to the host_ptr read below.
        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())
    }
}

/// Promote the existing `MappedF32Buffer` from distributional_q_tests.rs
/// to a shared module so all kernel test wrappers share one implementation.
/// Lives at `crates/ml/src/cuda_pipeline/mapped_pinned.rs` and re-exports
/// `MappedF32Buffer`, `MappedI32Buffer`. The `distributional_q_tests`
/// module's local copies are deleted in favor of the shared module.
mod mapped_f32 {
    pub use crate::cuda_pipeline::mapped_pinned::MappedF32Buffer;
}
```

- [ ] **Step 6: Add module to mod.rs**

Edit `crates/ml/src/cuda_pipeline/mod.rs`, add `pub mod gpu_moe_head;`.

- [ ] **Step 7: Run test to verify it passes**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_mixture_forward_matches_cpu_reference -- --ignored --nocapture 2>&1 | tail -15
```
Expected: PASS.

- [ ] **Step 8: Commit**

```bash
git add crates/ml/src/cuda_pipeline/moe_kernels.cu \
        crates/ml/src/cuda_pipeline/gpu_moe_head.rs \
        crates/ml/src/cuda_pipeline/mod.rs \
        crates/ml/tests/moe_kernels_test.rs \
        crates/ml/build.rs
git commit -m "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.

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

### Task 2.2 — moe_mixture_backward kernel

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/moe_kernels.cu` (append kernel)
- Modify: `crates/ml/src/cuda_pipeline/gpu_moe_head.rs` (load + test wrapper)
- Modify: `crates/ml/tests/moe_kernels_test.rs` (finite-differences gradient test)

- [ ] **Step 1: Write failing finite-differences gradient test**

Append to `crates/ml/tests/moe_kernels_test.rs`:

```rust
#[test]
#[ignore]
fn moe_mixture_backward_finite_differences() {
    let context = ml_core::cuda_autograd::test_helpers::test_context();
    let stream = Arc::clone(context.stream());

    const B: usize = 2;
    const C: usize = 32;  // Smaller than 256 for FD test speed
    const K: usize = 4;   // Smaller than 8 for FD test speed
    const EPS: f32 = 1e-3;
    const TOL: f32 = 1e-3;

    // Setup: random expert_outputs and normalized gate
    let mut 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;
        }
    }

    // Random upstream gradient dh_s2 [B, C]
    let dh_s2: Vec<f32> = (0..B * C)
        .map(|i| ((i * 19 + 11) % 100) as f32 / 100.0 - 0.5)
        .collect();

    // Run analytic backward kernel
    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[k0, b0, c0] via finite differences over h_s2 = sum_k g_k * e_k
    // dh_s2 . d/d(e_k0[b0,c0]) (mixture) = dh_s2[b0, c0] * gate[b0, k0]
    // So de_k[k0, b0, c0] should equal that.
    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 via the same approach over gate.
    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
}
```

- [ ] **Step 2: Run test to verify it fails**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_mixture_backward_finite_differences -- --ignored 2>&1 | tail -10
```
Expected: FAIL — `test_mixture_backward` doesn't exist.

- [ ] **Step 3: Append the backward kernel to moe_kernels.cu**

```cuda
extern "C" __global__ void moe_mixture_backward(
    const float* __restrict__ dh_s2,           /* [B, C] */
    const float* __restrict__ gate,            /* [B, K] */
    const float* __restrict__ expert_outputs,  /* [K, B, C] */
    float*       __restrict__ de_k,            /* [K, B, C] */
    float*       __restrict__ dg,              /* [B, K] */
    int B,
    int K,
    int C
) {
    /* Two stages handled by separate launches in the wrapper:
     * Stage A: de_k[k, b, c] = g[b, k] * dh_s2[b, c]    (one thread per (k,b,c))
     * Stage B: dg[b, k] = sum_c (e_k[b, c] * dh_s2[b, c])   (one block per (b,k), reduce over 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;

    float g  = gate[b * K + k];
    de_k[tid] = g * dh_s2[b * C + c];
}

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
) {
    /* One block per (b, k) pair; each block reduces over c. */
    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];
    }
}
```

- [ ] **Step 4: Add kernel handles + test wrapper to gpu_moe_head.rs**

Extend `GpuMoeHead`:

```rust
moe_mixture_backward: CudaFunction,
moe_dgate_reduce: CudaFunction,
```

In `new()`:
```rust
let moe_mixture_backward = module.load_function("moe_mixture_backward")?;
let moe_dgate_reduce = module.load_function("moe_dgate_reduce")?;
```

Add `test_mixture_backward` method that runs both kernels and returns `(de_k, dg)` to host.

- [ ] **Step 5: Run test to verify it passes**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_mixture_backward_finite_differences -- --ignored --nocapture 2>&1 | tail -20
```
Expected: PASS.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/moe_kernels.cu \
        crates/ml/src/cuda_pipeline/gpu_moe_head.rs \
        crates/ml/tests/moe_kernels_test.rs
git commit -m "feat(moe): moe_mixture_backward kernels + finite-differences test

Two-stage backward: moe_mixture_backward computes de_k = g · dh_s2;
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.

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

### Task 2.3 — moe_load_balance_loss kernel

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/moe_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_moe_head.rs`
- Modify: `crates/ml/tests/moe_kernels_test.rs`

- [ ] **Step 1: Write the failing test**

Append:

```rust
#[test]
#[ignore]
fn moe_load_balance_loss_correctness() {
    let context = ml_core::cuda_autograd::test_helpers::test_context();
    let stream = Arc::clone(context.stream());

    const B: usize = 8;
    const K: usize = 8;
    const LAMBDA: f32 = 0.01;

    // Random gate, normalized rows
    let mut gate: Vec<f32> = (0..B * K)
        .map(|i| ((i * 23 + 5) % 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;
        }
    }

    // CPU reference: λ · K · sum_k (mean_b g[b, k])²
    let mut col_means = vec![0.0_f32; K];
    for b in 0..B {
        for k in 0..K {
            col_means[k] += gate[b * K + k] / B as f32;
        }
    }
    let expected_loss: f32 = LAMBDA * K as f32 * col_means.iter().map(|m| m * m).sum::<f32>();

    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
    let actual_loss = head.test_load_balance_loss(&gate, B, K, LAMBDA).unwrap();

    assert!(
        (actual_loss - expected_loss).abs() < 1e-5,
        "Loss mismatch: expected {:.6}, got {:.6}",
        expected_loss, actual_loss,
    );
}

#[test]
#[ignore]
fn moe_load_balance_loss_uniform_minimum() {
    // At uniform gate (every entry = 1/K), loss should be exactly λ.
    let context = ml_core::cuda_autograd::test_helpers::test_context();
    let stream = Arc::clone(context.stream());

    const B: usize = 8;
    const K: usize = 8;
    const LAMBDA: f32 = 0.01;

    let gate: Vec<f32> = vec![1.0 / K as f32; B * K];

    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
    let loss = head.test_load_balance_loss(&gate, B, K, LAMBDA).unwrap();

    assert!(
        (loss - LAMBDA).abs() < 1e-5,
        "Uniform-gate loss should equal λ ({}), got {}",
        LAMBDA, loss,
    );
}
```

- [ ] **Step 2: Run tests to confirm fail**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_load_balance_loss -- --ignored 2>&1 | tail -10
```

- [ ] **Step 3: Append kernel**

```cuda
extern "C" __global__ void moe_load_balance_loss(
    const float* __restrict__ gate,    /* [B, K] */
    float*       __restrict__ loss,    /* [1] (single scalar accumulator) */
    int B,
    int K,
    float lambda
) {
    /* Single-block kernel:
     *   col_mean[k] = (1/B) * sum_b gate[b, k]
     *   loss = lambda * K * sum_k (col_mean[k])^2
     *
     * Grid: (1, K, 1); threads: B per block (or 256 for stride-loop over B if B>256).
     */
    extern __shared__ float smem[];
    int k = blockIdx.x;
    if (k >= K) return;

    int tid = threadIdx.x;
    int block = blockDim.x;

    float local = 0.0f;
    for (int b = tid; b < B; b += block) {
        local += gate[b * K + k];
    }
    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) {
        float col_mean = smem[0] / (float)B;
        float contrib = lambda * (float)K * col_mean * col_mean;
        /* Atomic add into single-element loss; per `feedback_no_atomicadd.md`,
         * use 1-block-per-k design with K-element scratch then reduce host-
         * side OR keep loss as [K]-array and reduce in next kernel.
         * Simpler approach: separate output[K] then reduce in second small kernel. */
        loss[k] = contrib;
    }
}

extern "C" __global__ void moe_load_balance_reduce(
    const float* __restrict__ loss_per_k,   /* [K] */
    float*       __restrict__ loss_total,   /* [1] */
    int K
) {
    /* Single-thread reduce: K is small (=8). */
    if (threadIdx.x != 0 || blockIdx.x != 0) return;
    float acc = 0.0f;
    for (int k = 0; k < K; ++k) acc += loss_per_k[k];
    loss_total[0] = acc;
}
```

(The two-step design avoids `atomicAdd` per `feedback_no_atomicadd.md`.)

- [ ] **Step 4: Add wrapper + test method**

In `gpu_moe_head.rs`, add `moe_load_balance_loss`, `moe_load_balance_reduce` kernel handles + a `test_load_balance_loss(&gate, B, K, λ) -> f32` method that runs both kernels and downloads the scalar.

- [ ] **Step 5: Run tests**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_load_balance -- --ignored --nocapture 2>&1 | tail -15
```
Expected: 2 tests pass.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/moe_kernels.cu \
        crates/ml/src/cuda_pipeline/gpu_moe_head.rs \
        crates/ml/tests/moe_kernels_test.rs
git commit -m "feat(moe): moe_load_balance_loss kernel + correctness tests

Two-step design (per-k contribution then K-reduce) avoids atomicAdd per
feedback_no_atomicadd.md. Tests verify match against CPU reference and
that uniform gate gives exactly λ as the minimum.

Backward gradient (∂loss/∂g[b,k] = 2λK · mean_b g[b,k] / B) lands in
Phase 3 alongside the wire-up.

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

### Task 2.4 — moe_expert_util_ema_update ISV producer

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/moe_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_moe_head.rs`
- Modify: `crates/ml/tests/moe_kernels_test.rs`

- [ ] **Step 1: Write the failing test**

Append to `moe_kernels_test.rs`:

```rust
#[test]
#[ignore]
fn moe_expert_util_ema_update_writes_correct_values() {
    let context = ml_core::cuda_autograd::test_helpers::test_context();
    let stream = Arc::clone(context.stream());

    const B: usize = 16;
    const K: usize = 8;
    const ALPHA: f32 = 0.05;

    // Initial ISV: utilization slots at 1/8, entropy at ln(8)
    let mut isv = vec![0.0_f32; 127];
    let util_base = 118;
    for k in 0..K {
        isv[util_base + k] = 1.0 / K as f32;
    }
    isv[126] = (K as f32).ln();

    // Synthetic gate: not uniform — peaky on expert 3
    let mut gate = vec![0.0_f32; B * K];
    for b in 0..B {
        for k in 0..K {
            gate[b * K + k] = if k == 3 { 0.6 } else { 0.057143 };
        }
    }

    // CPU reference: new_util[k] = (1-α)*old + α*mean_b gate[b,k]
    let mut expected_util = vec![0.0_f32; K];
    for k in 0..K {
        let col_mean: f32 = (0..B).map(|b| gate[b * K + k]).sum::<f32>() / B as f32;
        expected_util[k] = (1.0 - ALPHA) * isv[util_base + k] + ALPHA * col_mean;
    }
    // Entropy of batch-averaged distribution
    let avg_dist: Vec<f32> = (0..K)
        .map(|k| (0..B).map(|b| gate[b * K + k]).sum::<f32>() / B as f32)
        .collect();
    let avg_entropy: f32 = avg_dist
        .iter()
        .map(|p| if *p > 1e-9 { -p * p.ln() } else { 0.0 })
        .sum();
    let expected_entropy = (1.0 - ALPHA) * isv[126] + ALPHA * avg_entropy;

    // Run kernel
    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
    let new_isv = head
        .test_expert_util_ema_update(&isv, &gate, B, K, ALPHA, util_base, 126)
        .unwrap();

    // Verify
    for k in 0..K {
        assert!(
            (new_isv[util_base + k] - expected_util[k]).abs() < 1e-5,
            "util[{}]: expected {:.6}, got {:.6}",
            k, expected_util[k], new_isv[util_base + k],
        );
    }
    assert!(
        (new_isv[126] - expected_entropy).abs() < 1e-5,
        "entropy: expected {:.6}, got {:.6}",
        expected_entropy, new_isv[126],
    );
}
```

- [ ] **Step 2: Run, expect fail**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_expert_util_ema -- --ignored 2>&1 | tail -10
```

- [ ] **Step 3: Append kernel**

```cuda
extern "C" __global__ void moe_expert_util_ema_update(
    const float* __restrict__ gate,            /* [B, K] */
    float*       __restrict__ isv,             /* [ISV_TOTAL_DIM] */
    int B,
    int K,
    int isv_util_base,
    int isv_entropy_index,
    float alpha
) {
    /* Single-block, single-thread cold-path-cadence kernel.
     * Same shape as h_s2_rms_ema_update / aux_heads_loss_ema_update.
     */
    if (blockIdx.x != 0) return;
    if (threadIdx.x != 0) return;

    const float one_minus_alpha = 1.0f - alpha;

    /* Per-expert utilization EMA. */
    for (int k = 0; k < K; ++k) {
        float sum = 0.0f;
        for (int b = 0; b < B; ++b) sum += gate[b * K + k];
        float col_mean = sum / (float)B;
        float prev = isv[isv_util_base + k];
        isv[isv_util_base + k] = one_minus_alpha * prev + alpha * col_mean;
    }

    /* Gate-entropy EMA. */
    float entropy = 0.0f;
    for (int k = 0; k < K; ++k) {
        float sum = 0.0f;
        for (int b = 0; b < B; ++b) sum += gate[b * K + k];
        float col_mean = sum / (float)B;
        if (col_mean > 1e-9f) {
            entropy += -col_mean * logf(col_mean);
        }
    }
    float prev_e = isv[isv_entropy_index];
    isv[isv_entropy_index] = one_minus_alpha * prev_e + alpha * entropy;
}
```

- [ ] **Step 4: Add wrapper + test method**

- [ ] **Step 5: Run test**

```bash
SQLX_OFFLINE=true cargo test -p ml --test moe_kernels_test moe_expert_util_ema -- --ignored --nocapture 2>&1 | tail -10
```
Expected: PASS.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/moe_kernels.cu \
        crates/ml/src/cuda_pipeline/gpu_moe_head.rs \
        crates/ml/tests/moe_kernels_test.rs
git commit -m "feat(moe): moe_expert_util_ema_update ISV producer kernel + test

Single-thread cold-path-cadence kernel writes 8 per-expert utilization
EMAs + 1 gate-entropy EMA into ISV slots [118..127), α=0.05. Same shape
as h_s2_rms_ema_update / aux_heads_loss_ema_update.

GPU-resident, CPU read-only per pearl_cold_path_no_exception_to_gpu_drives.md.

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

---

## Phase 3 — Wire MoE into the training graph

Phase 3 replaces the existing `h_s2` producer with the gated expert mixture. After Phase 3, training updates gate + experts; HEALTH_DIAG emits MoE utilization line.

### Task 3.1 — Add gate forward path to fused_training

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Locate where h_s2 is currently computed**

```bash
grep -n "h_s2\|h_s1.*forward\|trunk.*forward" crates/ml/src/trainers/dqn/fused_training.rs | head -10
```

- [ ] **Step 2: Insert gate forward call between trunk and existing h_s2 producer**

Add after the trunk forward and before whatever currently produces `h_s2`:

```rust
// Gate forward: state[B, 42] -> 64 -> 8 -> softmax -> g[B, 8]
// Stored in moe_gate_out_buf for backward pass.
let gate_logits_h1 = self.moe_gate_w1.matmul(&states_padded, &cublas, &stream)?;
let gate_logits_h1 = act_kernels.leaky_relu_fwd(&gate_logits_h1, 0.01, &stream)?;
let gate_logits_h2 = self.moe_gate_w2.matmul(&gate_logits_h1, &cublas, &stream)?;
let gate_softmax = act_kernels.softmax_fwd(&gate_logits_h2, &stream)?;
self.moe_gate_out_buf.copy_from(&gate_softmax, &stream)?;
```

(Adapt to actual existing helper signatures.)

- [ ] **Step 3: Add buffers to FusedTrainingCtx struct**

Add fields:
```rust
moe_gate_out_buf: CudaSlice<f32>,        // [B, 8]
moe_expert_outputs_buf: CudaSlice<f32>,  // [8, B, 256]
moe_h_s2_buf: CudaSlice<f32>,            // [B, 256] — replaces existing h_s2 producer
```

Allocate in constructor sized to `batch_size * K`, `batch_size * K * h_s2_dim`, `batch_size * h_s2_dim`.

- [ ] **Step 4: cargo check**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat(moe): wire gate forward into fused_training graph

Gate softmax stored in moe_gate_out_buf for downstream mixture +
backward + ISV producer. Expert outputs and mixed h_s2 buffers added.

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

### Task 3.2 — Add expert forward path

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: After gate forward, add 8 expert MLP forwards**

```rust
// 8 experts: each h_s1[B, 256] -> 64 -> h_s2[B, 256]
// Outputs concatenated into moe_expert_outputs_buf [8, B, 256].
for k in 0..8 {
    let expert_h1 = self.moe_experts[k].w1.matmul(&h_s1, &cublas, &stream)?;
    let expert_h1 = act_kernels.leaky_relu_fwd(&expert_h1, 0.01, &stream)?;
    let expert_out = self.moe_experts[k].w2.matmul(&expert_h1, &cublas, &stream)?;
    // Copy expert_out into the kth slab of moe_expert_outputs_buf
    self.moe_expert_outputs_buf
        .slice_mut(k * batch_size * h_s2_dim..(k + 1) * batch_size * h_s2_dim)
        .copy_from(&expert_out, &stream)?;
}
```

- [ ] **Step 2: cargo check**

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat(moe): wire 8 expert forward passes into fused_training graph

Each expert: h_s1[B,256] -> 64 -> h_s2[B,256]. Outputs stacked into
[8, B, 256] for the mixture kernel.

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

### Task 3.3 — Replace existing h_s2 producer with MoE mixture

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Locate where h_s2 currently feeds the branching heads**

```bash
grep -n "h_s2" crates/ml/src/trainers/dqn/fused_training.rs | head -20
```

- [ ] **Step 2: Launch moe_mixture_forward kernel after expert forward**

Replace the existing h_s2 production line with:

```rust
// Mixture: h_s2[b, c] = Σ_k g[b, k] · expert_outputs[k, b, c]
self.moe_head.launch_mixture_forward(
    &self.moe_expert_outputs_buf,
    &self.moe_gate_out_buf,
    &mut self.moe_h_s2_buf,
    batch_size,
    8,
    h_s2_dim,
)?;

// Subsequent code uses self.moe_h_s2_buf as h_s2 input to branching heads.
```

(Add `launch_mixture_forward` method to `GpuMoeHead`.)

- [ ] **Step 3: Update branching head input to use moe_h_s2_buf**

Find every reference to the previous `h_s2` source; replace with `moe_h_s2_buf`.

- [ ] **Step 4: cargo check**

- [ ] **Step 5: Smoke test (will fail loudly due to architecture-hash mismatch on existing checkpoint, but should at least compile)**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
```

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs \
        crates/ml/src/cuda_pipeline/gpu_moe_head.rs
git commit -m "feat(moe): wire mixture kernel; h_s2 is now MoE-mixed

Branching heads consume moe_h_s2_buf produced by the gated expert mixture.
Replaces the previous direct-from-trunk h_s2 producer.

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

### Task 3.4 — Wire load-balance aux loss

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (loss aggregation)

- [ ] **Step 1: Locate the loss aggregation section**

```bash
grep -n "total_loss\|loss_acc\|cql_alpha.*loss\|iqn_lambda.*loss" crates/ml/src/trainers/dqn/fused_training.rs | head -5
```

- [ ] **Step 2: Add load-balance loss kernel call to total loss**

```rust
// MoE load-balance aux loss: λ · K · Σ_k (mean_b g[b, k])²
let moe_lambda = hyperparams.moe_lambda.unwrap_or(0.01);
self.moe_head.launch_load_balance_loss(
    &self.moe_gate_out_buf,
    &mut self.moe_load_balance_loss_buf,
    batch_size,
    8,
    moe_lambda,
)?;

// loss accumulator already adds component buffers; add MoE term.
add_to_total_loss(&self.moe_load_balance_loss_buf)?;
```

- [ ] **Step 3: cargo check**

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat(moe): add load-balance aux loss to total training loss

λ · K · Σ_k (mean_b g[b, k])² added per training step. Reads
hyperparams.moe_lambda (default 0.01).

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

### Task 3.5 — Wire ISV producer + HEALTH_DIAG line

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`

- [ ] **Step 1: Locate the existing ISV-producer launch block**

```bash
grep -n "launch_aux_heads_loss_ema\|launch_h_s2_rms_ema\|launch_vsn_mask_ema" crates/ml/src/trainers/dqn/trainer/training_loop.rs | head -5
```

- [ ] **Step 2: Add MoE ISV producer launch alongside the existing producers**

```rust
// MoE per-expert utilization + gate-entropy ISV producer.
if let Some(ref fused) = self.fused_ctx {
    if let Err(e) = fused.trainer().launch_moe_expert_util_ema(ema_alpha) {
        tracing::warn!("MoE expert util ema launch failed: {e}");
    }
}
```

- [ ] **Step 3: Add HEALTH_DIAG aux_moe line in epoch loop**

After the existing aux line emit, add:

```rust
{
    use crate::cuda_pipeline::gpu_dqn_trainer::{
        MOE_EXPERT_UTIL_EMA_BASE, MOE_EXPERT_UTIL_EMA_COUNT, MOE_GATE_ENTROPY_EMA_INDEX,
    };
    let (utils, entropy) = if let Some(ref fused) = self.fused_ctx {
        let trainer = fused.trainer();
        let mut utils = [0.0_f32; 8];
        for k in 0..MOE_EXPERT_UTIL_EMA_COUNT {
            utils[k] = trainer.read_isv_signal_at(MOE_EXPERT_UTIL_EMA_BASE + k);
        }
        (utils, trainer.read_isv_signal_at(MOE_GATE_ENTROPY_EMA_INDEX))
    } else {
        ([0.0; 8], 0.0)
    };
    tracing::info!(
        "HEALTH_DIAG[{}]: aux_moe [util=[{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3}] ent={:.3}]",
        epoch,
        utils[0], utils[1], utils[2], utils[3], utils[4], utils[5], utils[6], utils[7],
        entropy,
    );
}
```

- [ ] **Step 4: cargo check + commit**

```bash
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
git add crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat(moe): HEALTH_DIAG aux_moe line + ISV producer launch

Per-step launch of moe_expert_util_ema_update kernel; per-epoch
HEALTH_DIAG line shows 8 utilization values + gate entropy. Operators
can watch live whether experts are differentiating or collapsing.

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

### Task 3.6 — Wire backward pass

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (backward graph extension)

- [ ] **Step 1: Locate the backward graph construction**

```bash
grep -n "submit_backward\|compute_grad_h_s2\|backward.*h_s2" crates/ml/src/trainers/dqn/fused_training.rs | head -10
```

- [ ] **Step 2: Insert MoE mixture backward + expert backward + gate backward**

After `dh_s2` is computed by branching heads' backward:

```rust
// MoE mixture backward: dh_s2 -> de_k + dg
self.moe_head.launch_mixture_backward(
    &dh_s2,
    &self.moe_gate_out_buf,
    &self.moe_expert_outputs_buf,
    &mut self.moe_de_k_buf,
    &mut self.moe_dg_buf,
    batch_size,
    8,
    h_s2_dim,
)?;

// Expert backward: each expert's params get gradient via cuBLAS SGEMM backward
for k in 0..8 {
    let de_k_slice = self.moe_de_k_buf.slice(
        k * batch_size * h_s2_dim..(k + 1) * batch_size * h_s2_dim,
    );
    backward_through_expert_k(&de_k_slice, &h_s1, &mut self.moe_experts[k], &cublas, &stream)?;
}

// Gate backward: dg + softmax backward + MLP backward
backward_through_gate(&self.moe_dg_buf, &states_padded, &mut self.moe_gate, &cublas, &stream)?;
```

(Implementations for `backward_through_expert_k` and `backward_through_gate` follow the existing cuBLAS-backward pattern in fused_training.)

- [ ] **Step 3: cargo check + commit**

```bash
SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat(moe): backward pass through gate + experts + mixture

Mixture kernel produces de_k and dg; each is propagated through
respective MLP backward via cuBLAS SGEMM. All 8 experts and gate get
gradient updates per step.

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

### Task 3.7 — Verify Phase 3 with smoke test

**Files:**
- Verify: full DQN training loop runs end-to-end

- [ ] **Step 1: Delete the old smoke checkpoint (architecture incompatible)**

```bash
rm -rf target/smoke_multi_fold_output
```

- [ ] **Step 2: Run smoke test fresh**

```bash
FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- multi_fold_convergence --ignored --nocapture 2>&1 | tee /tmp/moe_smoke.log | tail -50
```

Expected:
- Smoke completes (≤ 25 minutes).
- HEALTH_DIAG `aux_moe [util=[...]` lines present in log.
- Per-expert utilization values are NOT all 0.125 (gate has differentiated).
- Gate entropy < ln(8) (some peakiness emerged).
- Final val metrics not catastrophic (val_loss finite, finite Sharpe).

- [ ] **Step 3: Commit if all pass; if collapse criteria triggered, debug and re-commit**

```bash
git add . # any incidental fixes
git commit -m "feat(moe): Phase 3 smoke verification — MoE forward+backward end-to-end

Local smoke (5 ep × 3 folds) completes with healthy utilization spread
across experts (no collapse). HEALTH_DIAG aux_moe line emits per epoch.

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

---

## Phase 4 — Cleanup cascade: delete vestigial RegimeConditionalDQN

Phase 4 is the atomic deletion of the legacy infrastructure. Single commit per `feedback_no_partial_refactor.md`.

### Task 4.1 — Delete regime_conditional.rs and rewrite DQNAgentType

**Files:**
- Delete: `crates/ml-dqn/src/regime_conditional.rs`
- Modify: `crates/ml-dqn/src/lib.rs`
- Modify: `crates/ml/src/trainers/dqn/config.rs` (DQNAgentType rewrite)

- [ ] **Step 1: Delete the regime_conditional module**

```bash
git rm crates/ml-dqn/src/regime_conditional.rs
```

- [ ] **Step 2: Drop regime exports from lib.rs**

Remove `pub mod regime_conditional;` and `pub use regime_conditional::{...};` lines.

- [ ] **Step 3: Rewrite DQNAgentType to wrap DQN directly**

In `crates/ml/src/trainers/dqn/config.rs`:

```rust
/// DQN agent wrapper around `DQN`.
///
/// Provides the unified API surface that the training loop, fused CUDA trainer,
/// and checkpoint code call. All methods delegate to `self.agent`.
pub struct DQNAgentType {
    pub(crate) agent: ml_dqn::DQN,
}

impl DQNAgentType {
    pub fn new(agent: ml_dqn::DQN) -> Self {
        Self { agent }
    }
}

// Delegate all methods to self.agent directly. No more 3-head wrapper.
impl DQNAgentType {
    pub fn select_action(&mut self, state: &[f32]) -> Result<FactoredAction, MLError> {
        self.agent.select_action(state)
    }
    pub fn forward(&self, state: &GpuTensor) -> Result<GpuTensor, MLError> {
        self.agent.forward(state)
    }
    pub fn batch_q_values(&self, states: &GpuTensor) -> Result<GpuTensor, MLError> {
        self.agent.q_values_for_batch(states)
    }
    // ... and so on for every method that previously delegated to RegimeConditionalDQN.
    pub fn get_count_bonuses_branched(&self) -> ([f32; 4], [f32; 3], [f32; 3]) {
        self.agent.get_count_bonuses_branched()
    }
    pub fn reset_count_bonus(&mut self) {
        self.agent.reset_count_bonus();
    }
    pub fn set_noise_sigma_scale(&mut self, scale: f64) {
        self.agent.set_noise_sigma_scale(scale);
    }
    // ... etc.
}
```

- [ ] **Step 4: Drop 4 regime threshold fields from DQNConfig**

In `crates/ml-dqn/src/dqn.rs::DQNConfig`, delete:
```rust
pub regime_adx_idx: usize,
pub regime_cusum_idx: usize,
pub regime_adx_threshold: f32,
pub regime_cusum_threshold: f32,
```

And from all 4 default builders. And from any test that sets them.

- [ ] **Step 5: cargo check workspace**

```bash
SQLX_OFFLINE=true cargo check --workspace --tests 2>&1 | tail -20
```

Expect many errors about missing `RegimeType`, etc. Fix each one in turn:
- Tests referencing `RegimeType` enum: delete the test or rewrite as MoE-aware.
- Trainer call sites using `set_regime_metrics`: delete (no analog).
- Examples: drop regime config setters.

- [ ] **Step 6: Run smoke test**

```bash
FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- multi_fold_convergence --ignored --nocapture 2>&1 | tail -20
```

Expected: completes with same MoE health profile as Phase 3.

- [ ] **Step 7: Commit (atomic — single commit for ALL deletions)**

```bash
git add -A
git commit -m "$(cat <<'EOF'
refactor(dqn): delete vestigial RegimeConditionalDQN — MoE replaces it

Atomic cleanup per feedback_no_partial_refactor.md:

- Delete crates/ml-dqn/src/regime_conditional.rs (~700 LOC).
- Drop RegimeConditional* exports from lib.rs.
- Rewrite DQNAgentType to wrap DQN directly (no RegimeConditionalDQN
  field, no 3-head delegation gymnastics).
- Drop 4 regime-threshold fields from DQNConfig (regime_adx_idx,
  regime_cusum_idx, regime_adx_threshold, regime_cusum_threshold) +
  matching default-builder entries.
- Migrate all trainer/test/example call sites to direct-DQN methods.

The MoE redesign (Phases 1-3) replaces the regime-conditional purpose:
gate handles routing, experts specialize. Per pearl_learned_gate_subsumes_
handcoded.md, the threshold classifier was a strict information
bottleneck; the gate sees ADX (40) and CUSUM (41) as part of the full
state vector and learns its own decomposition.

Old per-regime checkpoint format (3 files: trending/ranging/volatile)
gone — single safetensors with shared trunk + 8 experts + gate +
branching heads.

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

---

## Phase 5 — Tests + L40S validation

### Task 5.1 — Layer 3 gate-differentiation validation test

**Files:**
- Create: `crates/ml/src/trainers/dqn/smoke_tests/moe_gate_differentiation.rs`

- [ ] **Step 1: Create the test file**

```rust
//! Layer 3 gate-differentiation validation per spec §8.
//!
//! On a held-out batch of states from the smoke val fold, verify:
//! - mean_b H[b] < 0.9·ln(8) (gate is peakier than uniform on average)
//! - Per-expert weight differs across ADX-low vs ADX-high state bins
//!   by > 0.05 for at least one expert (gate uses state info meaningfully)

use std::path::PathBuf;

#[test]
#[ignore]
fn moe_gate_differentiates_after_smoke() {
    // Load smoke checkpoint
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .ancestors()
        .find(|p| p.join("Cargo.toml").exists() && p.join("crates").exists())
        .unwrap()
        .to_path_buf();
    let checkpoint = workspace_root.join("target/smoke_multi_fold_output/dqn_fold0_best.safetensors");

    if !checkpoint.exists() {
        panic!(
            "Test requires smoke checkpoint at {}. Run smoke first: \
             FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test \
             -p ml --lib -- multi_fold_convergence --ignored",
            checkpoint.display()
        );
    }

    // Load checkpoint, run forward on val batch, extract gate outputs
    // ... (implementation specifics depend on whether MoE inference is exposed)

    // Assertion 1: mean entropy < 0.9·ln(8) ≈ 1.871
    // Assertion 2: gate weights differ across ADX bins by >0.05 for ≥1 expert
    // ... (remaining implementation per spec §8.3)
}
```

(Full test implementation depends on the MoE inference path landing — adapt to match the exposed API.)

- [ ] **Step 2: Run the test**

```bash
SQLX_OFFLINE=true cargo test -p ml --lib moe_gate_differentiates_after_smoke -- --ignored --nocapture 2>&1 | tail -15
```

Expected: PASS (gate differentiates) or FAIL with concrete numbers showing why.

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/trainers/dqn/smoke_tests/moe_gate_differentiation.rs
git commit -m "test(moe): Layer 3 gate-differentiation validation

Asserts on ~1000 val states: mean entropy < 0.9·ln(8) and per-expert
weight differs across ADX-low/high bins by >0.05 for at least one
expert. Failures are findings (concrete numbers reported per
feedback_kill_runs_on_anomaly_quickly.md).

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

### Task 5.2 — Layer 5 architecture-hash backward-incompat test

**Files:**
- Modify: `crates/ml-dqn/tests/architecture_hash_test.rs` (or create if absent)

- [ ] **Step 1: Create test asserting old checkpoint rejection**

```rust
#[test]
fn architecture_hash_old_checkpoint_rejected() {
    // Use an existing pre-MoE checkpoint OR construct a synthetic safetensors
    // with the old-format metadata (3-head regime, no MoE tensors).
    // Assert that DQN::load_from_safetensors returns Err with message
    // containing "layout fingerprint mismatch".

    // Implementation: build minimal safetensors with old metadata, load it,
    // assert .is_err() and error string contains the expected substring.
}
```

- [ ] **Step 2: Run + commit**

### Task 5.3 — L40S production validation

**Files:**
- N/A (deployment via Argo)

- [ ] **Step 1: Push branch with all MoE commits**

```bash
git push origin main
```

- [ ] **Step 2: Trigger L40S 30-epoch validation**

```bash
./scripts/argo-train.sh dqn --baseline --epochs 30 --folds 3 --gpu-pool ci-training-l40s --tag moe-validation
```

- [ ] **Step 3: Monitor HEALTH_DIAG live with kill criteria from spec §8.4**

```bash
kubectl logs -f -l workflows.argoproj.io/workflow=<workflow-name> | grep -E "HEALTH_DIAG|aux_moe|val_loss|Sharpe"
```

Watch for:
- `aux_moe [util=[g0,...,g7] ent=...]` line per epoch with finite values.
- Gate entropy stays > 0.1 in first 3 epochs (no snap-collapse).
- Per-expert utilization stays > 0.1% for all 8 (no dead experts after 5 epochs).
- Training Sharpe at epoch 5 within 10% of baseline.

If any kill criterion trips, kill the run immediately:
```bash
kubectl delete workflow <workflow-name>
```

- [ ] **Step 4: Document results**

If validation passes, append findings to `docs/dqn-wire-up-audit.md` with the `train-<id>` workflow name and final metrics. If validation fails with specific numbers, document the failure, root-cause, and follow-up commit.

---

## Self-Review

**Spec coverage check:**

- §1 Background → not directly a task; informs Task 0.1 + Phase 4 motivation. ✓
- §2 Decision → realized across all phases. ✓
- §3 Architecture → realized in Tasks 1.4–3.6. ✓
- §4.1 Gate subnetwork → Task 1.5 (skeleton), Task 3.1 (forward), Task 3.6 (backward). ✓
- §4.2 Expert MLPs → Task 1.5 (skeleton), Task 3.2, Task 3.6. ✓
- §4.3 moe_lambda hyperparameter → Task 1.1. ✓
- §4.4 Existing components retained → no task needed (unchanged). ✓
- §4.5 Components deleted → Task 4.1. ✓
- §4.6 ISV slot extension → Tasks 1.2, 1.3. ✓
- §5.1 Forward → Tasks 3.1–3.4. ✓
- §5.2 Backward → Task 3.6. ✓
- §5.3 Optimizer → automatic via params_buf extension (Task 1.6). ✓
- §5.4 Replay buffer → no change needed (already single-shared after Phase 0). ✓
- §5.5 Checkpoint → automatic via params_buf hash (Task 1.6); legacy 3-file format gone in Task 4.1. ✓
- §6 GPU integration → Tasks 1.6, 2.1–2.4, 3.1–3.6. ✓
- §7 Cleanup cascade → Task 4.1. ✓
- §8 Testing strategy → Tasks 2.1–2.4 (Layer 1), 3.7 (Layer 2), 5.1 (Layer 3), 5.2 (Layer 5), 5.3 (Layer 4). ✓
- §9 Success criteria → checked in Task 5.3. ✓
- §10 Out of scope → not implemented (correct). ✓

**Placeholder scan:** searched for "TBD", "TODO", "Similar to Task" — none present. Some tasks reference "(adapt to actual existing helper signatures)" which is intentional flexibility for places where I cannot read the existing code in this conversation.

**Type consistency:** `MoeGate`, `MoeExpert`, `MOE_NUM_EXPERTS=8`, `MOE_GATE_HIDDEN=64`, `MOE_EXPERT_BOTTLENECK=64`, `MOE_EXPERT_UTIL_EMA_BASE=118`, `MOE_GATE_ENTROPY_EMA_INDEX=126`, `ISV_TOTAL_DIM=127` — used consistently across all tasks.

---

## Execution Handoff

**Plan complete and saved to `docs/superpowers/plans/2026-04-27-moe-regime-redesign.md`. Two execution options:**

**1. Subagent-Driven (recommended)** — fresh subagent per task, two-stage review (spec compliance, then code quality) between tasks, fast iteration.

**2. Inline Execution** — execute tasks in this session using executing-plans, batch execution with checkpoints for review.

**Which approach?**
