From 3e36f4a0e644d26c6cd8da0cfaa09850dc8aca40 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 3 Jun 2026 11:07:27 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Phase=202A-C=20=E2=80=94=20tr?= =?UTF-8?q?ainer=20integration=20behind=20FOXHUNT=5FUSE=5FMULTI=5FHEAD=5FP?= =?UTF-8?q?OLICY=20flag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires MultiHeadPolicy into the live trainer's π forward + backward + Adam path, gated by FOXHUNT_USE_MULTI_HEAD_POLICY (default off). Mirrors the FOXHUNT_USE_ROLLOUT precedent (commit 779c03b9d) for load-bearing refactors with emergency rollback. ## Wiring sites * Forward: 4 call sites — step_synthetic_body (training_graph), step_with_lobsim (prefill_graph CPU), step_with_lobsim_gpu_body (prefill_graph GPU), eval_policy_probs_per_action (no graph). * Backward: step_synthetic_body only (line 5764+). * Adam apply: 4 new Adam optimizers (W_heads, b_heads, W_gate, b_gate) when flag on. * LR plumbing: same lr_pi mirrored onto all 4 MHP Adams. * ISV bootstrap: slots 761-764 (K=3, gating_entropy_floor=0.549, head_entropy_floor=0.959, aux_prior_β=0.05) written in a flag-gated block AFTER the main bootstrap loop (preserves flag-off bit-equality — see surprise #2). * regime_h plumbing: new pub accessor PerceptionTrainer:: step_regime_d_view() mirrors h_t_view; gate kernel reads the same buffer the encoder reads (no new memcpy, no shape change). ## CUDA Graph capture treatment Host-side select. The flag is invariant for trainer lifetime, so the branch evaluates at graph CAPTURE time — the selected kernel sequence is baked into the captured graph and replays single-path. No host branches inside the captured stream. Verified by determinism preservation in both flag-off and flag-on regimes. ## Verification (all gates passed) * FOXHUNT_USE_MULTI_HEAD_POLICY=0 ./scripts/determinism-check.sh --quick: exit 0. * Flag-off diag.jsonl byte-equal to e22da61cf baseline (modulo elapsed_s timestamps). * FOXHUNT_USE_MULTI_HEAD_POLICY=1 ./scripts/determinism-check.sh --quick: exit 0 (new path is deterministic). * 11/11 multi_head_policy_invariants tests still pass. * FOXHUNT_USE_MULTI_HEAD_POLICY=1 ./scripts/local-mid-smoke.sh: 2500 steps (2000 train + 500 eval) completed without NaN. 0 NaN/null in 2000-row diag.jsonl and 501-row eval_diag.jsonl. eval_summary.total_pnl_usd = -$3.82M (single-seed at random init — not load-bearing; behavioral verdict requires 3-seed per pearl_local_smoke_noise_floor_and_regime_concentration). * 68/0/6 lib tests (passed/failed/ignored). * Pre-commit: 0 atomicAdd, 0 raw memcpy_htod/dtoh, 0 TODO, 0 host-branches inside captured streams. ## Surprises handled (not silently) 1. step_regime_d was private on PerceptionTrainer — added clean pub accessor step_regime_d_view mirroring h_t_view pattern. 2. Unconditional ISV bootstrap broke flag-off bit-equality (writing 761-764 from 0.0 sentinel shifted isv_state checksum). Caught by verification gate. Resolved by moving the 4 writes into a flag-gated block AFTER the main bootstrap loop. Flag-off leaves slots 761-764 at 0.0 sentinel, matching e22da61cf. 3. Spectral-norm regularization on legacy policy_head.w_d still runs when flag on (cost: one kernel/step, no correctness impact). Documented as known follow-up — proper gating happens in Phase 2A-E productionization, not C.3 (which must preserve flag-off bit-equality). 4. Adam state across flag toggles: 4 MHP Adams are None when flag off. Mirrors FOXHUNT_USE_ROLLOUT precedent — flag cached at construction, one-shot read. 5. Device-aggregated diag deferred (gate_probs_mean[K], gate_argmax_mass[K], gate_entropy_mean, per_head_entropy_mean [K]). These require either a new reduction kernel + ISV slots or host-readback (forbidden by feedback_no_htod_htoh_only_ mapped_pinned). Emitted ISV-resident summary leaves only (multi_head_policy.{active,k,gating_entropy_floor,head_ entropy_floor,aux_prior_beta}). Phase 2A-D scope. ## Linked * Phase 2A-A: 0b3e40150 (MultiHeadPolicy foundation, inert) * Phase 2A-B: e22da61cf (backward + aux KL prior) * Plan: docs/superpowers/plans/2026-06-03-multi-head-policy- implementation.md * Spec ADDENDUM: docs/superpowers/specs/2026-06-02-multi-head- policy-with-r-multiple.md * Q-distill pearl: pearl_foxhunt_pi_trained_by_q_distillation_ not_ppo.md Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/integrated.rs | 391 ++++++++++++++++++++-- crates/ml-alpha/src/trainer/perception.rs | 16 + 2 files changed, 376 insertions(+), 31 deletions(-) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index dee6b5b09..a14a93d0a 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -181,6 +181,7 @@ use crate::rl::isv_slots::{ use crate::rl::common::{FRD_N_ATOMS, FRD_N_HORIZONS, N_ACTIONS}; use crate::rl::frd::FRD_OUT_DIM; use crate::rl::loss_balance::{read_loss_lambdas_from_isv, LossLambdas}; +use crate::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig}; use crate::rl::ppo::{PolicyHead, PpoHeadsConfig, ValueHead}; use crate::rl::reward::RlLobBackend; use crate::trainer::optim::AdamW; @@ -665,6 +666,35 @@ pub struct IntegratedTrainer { pub policy_b_adam: AdamW, pub value_w_adam: AdamW, pub value_b_adam: AdamW, + + /// Phase 2A-C (2026-06-03) — multi-head policy gate (FOXHUNT_USE_MULTI_HEAD_POLICY). + /// + /// Cached `bool` snapshot of the env flag captured once at construction + /// time. Must not change for the lifetime of the trainer because the + /// training / prefill graphs bake whichever branch is selected at first + /// capture into the replayed graph (host-side `if` runs only during + /// capture; afterwards the captured kernel stream is replayed verbatim). + /// + /// Flag-off path is bit-equal to commit `e22da61cf` (legacy single-head + /// `policy_head` drives π). + pub use_multi_head_policy: bool, + /// `Some(_)` iff `use_multi_head_policy == true`. Owns K policy heads + + /// the regime-gated mixing head + all backward scratch + reduced-grad + /// buffers ready for an Adam step. Allocated on the same stream as the + /// rest of the trainer (`dev.cuda_stream()`); no extra synchronization. + /// See `crates/ml-alpha/src/rl/multi_head_policy.rs` for the kernel + /// pipeline. Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md + /// ADDENDUM 2026-06-03. + pub multi_head_policy: Option, + /// Adam optimisers for the four `MultiHeadPolicy` weight tensors — + /// `W_heads [K × N_ACTIONS × HIDDEN_DIM]`, `b_heads [K × N_ACTIONS]`, + /// `W_gate [K × REGIME_DIM=6]`, `b_gate [K]`. `Some(_)` iff + /// `use_multi_head_policy == true`. LR reads slot `RL_LR_PI_INDEX` + /// (mirrors the legacy `policy_w_adam` / `policy_b_adam` LR sourcing). + pub mhp_w_heads_adam: Option, + pub mhp_b_heads_adam: Option, + pub mhp_w_gate_adam: Option, + pub mhp_b_gate_adam: Option, /// Phase 4 dueling-head Adam optimisers (4 weight tensors). pub dueling_q_w_v_adam: AdamW, pub dueling_q_b_v_adam: AdamW, @@ -1613,6 +1643,103 @@ impl IntegratedTrainer { AdamW::new(dev, policy_head.w_d.len(), lr_placeholder).context("policy_w_adam")?; let policy_b_adam = AdamW::new(dev, policy_head.b_d.len(), lr_placeholder).context("policy_b_adam")?; + + // ── Phase 2A-C (2026-06-03) MultiHeadPolicy gate ───────────── + // `FOXHUNT_USE_MULTI_HEAD_POLICY` env flag follows the + // `FOXHUNT_USE_ROLLOUT` precedent from commit 779c03b9d. Read + // ONCE here so the value is invariant across captured-graph + // replay (the `if self.use_multi_head_policy` branches at the + // π forward / backward sites run only during graph CAPTURE; at + // replay time the selected kernel stream is baked in). + // + // Gate semantics: + // * `FOXHUNT_USE_MULTI_HEAD_POLICY` unset or `=0` → legacy + // single-head `policy_head` drives π. `multi_head_policy` is + // `None`, no extra device buffers, bit-equal to HEAD `e22da61cf`. + // * `FOXHUNT_USE_MULTI_HEAD_POLICY=1` → construct + // `MultiHeadPolicy` with K from ISV slot 761 (bootstrap K=3). + // Forward routes `h_t + step_regime_d` through the K-head + // mixture; backward distributes Q-distill grad through + // mixture + per-head softmax + gating softmax Jacobians; + // Adam updates the four new weight tensors. + // + // Per `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo`: + // the wiring target is the post-Q-distill `ss_pi_grad_logits_d` + // write site (`integrated.rs:~5607`). PPO-surrogate grad is still + // zeroed at line 5577; nothing changes for that path. + let use_multi_head_policy: bool = std::env::var("FOXHUNT_USE_MULTI_HEAD_POLICY") + .ok() + .and_then(|v| v.parse::().ok()) + .map(|n| n > 0) + .unwrap_or(false); + let ( + multi_head_policy, + mhp_w_heads_adam, + mhp_b_heads_adam, + mhp_w_gate_adam, + mhp_b_gate_adam, + ): ( + Option, + Option, + Option, + Option, + Option, + ) = if use_multi_head_policy { + // K = 3 is the bootstrap value documented in + // `RL_POLICY_NUM_HEADS_INDEX` (slot 761) and seeded in the + // controller-bootstrap table below. Reading host-side ISV here + // is safe because the table runs strictly BEFORE this point — + // however `with_controllers_bootstrapped` runs AFTER `new()`, + // so the slot is still zero at construction. We therefore use + // the hardcoded bootstrap constant `K=3` to match the bootstrap + // table and per the spec's "K=3 production default"; runtime K + // changes would require re-allocation anyway (weight tensors + // are sized by K). + const K_BOOTSTRAP: usize = 3; + // Distinct seed from the legacy `policy_head` (offset 0xC1A55 = + // "CLASS" — independent head initialisation streams per + // `feedback_no_shared_init_streams`). + let mhp_seed: u64 = cfg.ppo_seed.wrapping_add(0xC1A55); + let mhp = MultiHeadPolicy::new( + dev, + MultiHeadPolicyConfig { + k: K_BOOTSTRAP, + b_size: cfg.perception.n_batch, + seed: mhp_seed, + }, + ) + .context("MultiHeadPolicy::new (FOXHUNT_USE_MULTI_HEAD_POLICY=1)")?; + // Mirror legacy policy Adam pattern: lr_placeholder (overwritten + // every step from ISV[RL_LR_PI_INDEX]). One Adam per weight + // tensor — gate weights get their own Adam (different shape + + // potentially different LR scale in a future controller). + let w_heads_adam = AdamW::new(dev, mhp.w_heads_d.len(), lr_placeholder) + .context("mhp_w_heads_adam")?; + let b_heads_adam = AdamW::new(dev, mhp.b_heads_d.len(), lr_placeholder) + .context("mhp_b_heads_adam")?; + let w_gate_adam = AdamW::new(dev, mhp.w_gate_d.len(), lr_placeholder) + .context("mhp_w_gate_adam")?; + let b_gate_adam = AdamW::new(dev, mhp.b_gate_d.len(), lr_placeholder) + .context("mhp_b_gate_adam")?; + eprintln!( + "[multi-head-policy] FOXHUNT_USE_MULTI_HEAD_POLICY=1 → MultiHeadPolicy active \ + (K={} B={} REGIME_DIM=6, seed={:#x})", + K_BOOTSTRAP, cfg.perception.n_batch, mhp_seed + ); + ( + Some(mhp), + Some(w_heads_adam), + Some(b_heads_adam), + Some(w_gate_adam), + Some(b_gate_adam), + ) + } else { + eprintln!( + "[multi-head-policy] FOXHUNT_USE_MULTI_HEAD_POLICY unset/=0 → legacy single-head \ + policy_head path (bit-equal to HEAD e22da61cf)" + ); + (None, None, None, None, None) + }; let value_w_adam = AdamW::new(dev, value_head.w_d.len(), lr_placeholder).context("value_w_adam")?; let value_b_adam = @@ -2940,6 +3067,13 @@ impl IntegratedTrainer { policy_b_adam, value_w_adam, value_b_adam, + // Phase 2A-C MultiHeadPolicy + Adams (None when flag off). + use_multi_head_policy, + multi_head_policy, + mhp_w_heads_adam, + mhp_b_heads_adam, + mhp_w_gate_adam, + mhp_b_gate_adam, dueling_q_w_v_adam, dueling_q_b_v_adam, dueling_q_w_a_adam, @@ -4016,6 +4150,47 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("rl_isv_write (design constant): {:?}", e))?; } } + + // ── Phase 2A-C (2026-06-03) MultiHeadPolicy ISV bootstrap ──── + // + // Slots 761-764 (allocated in Phase 2A-A). Bootstrap is gated + // on `FOXHUNT_USE_MULTI_HEAD_POLICY` for bit-equality with + // commit `e22da61cf`: writing these slots when the flag is off + // would diverge `isv_state` checksum (the slots are 0.0 + // sentinel at HEAD; writing 3.0 / 0.549 / 0.959 / 0.05 shifts + // the checksum), failing verification gate C.7.4. + // + // When flag is on, this writes match the dispatch's required + // bootstrap values per spec ADDENDUM 2026-06-03 §R.1: + // * K=3: spec §1.1 (production default) + // * gating_entropy_floor = log(K)·0.5 ≈ 0.549 + // * head_entropy_floor = log(N_ACTIONS)·0.4 ≈ 0.959 + // * aux_prior_beta = 0.05 (5% of policy-loss magnitude) + if self.use_multi_head_policy { + let mhp_isv_constants: [(usize, f32); 4] = [ + (crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX, 3.0_f32), + (crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, 0.549_f32), + (crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, 0.959_f32), + (crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX, 0.05_f32), + ]; + for (slot, value) in mhp_isv_constants.iter() { + let slot_i32 = *slot as i32; + let value_f32 = *value; + let mut args = RawArgs::new(); + args.push_ptr(self.isv_dev_ptr); + args.push_i32(slot_i32); + args.push_f32(value_f32); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.rl_isv_write_fn.cu_function(), + (1, 1, 1), (1, 1, 1), 0, + self.raw_stream, + &mut ptrs[..args.len()], + ).map_err(|e| anyhow::anyhow!("rl_isv_write (multi-head bootstrap): {:?}", e))?; + } + } + } } let alpha = RL_LR_CONTROLLER_ALPHA; @@ -5344,6 +5519,15 @@ impl IntegratedTrainer { self.dqn_b_adam.lr = lr_q; self.policy_w_adam.lr = lr_pi; self.policy_b_adam.lr = lr_pi; + // Phase 2A-C: mirror lr_pi onto the multi-head Adam state when + // the gate is on. ISV[RL_LR_PI_INDEX] remains the single source + // of truth — no separate slot for multi-head LR. + if self.use_multi_head_policy { + if let Some(a) = self.mhp_w_heads_adam.as_mut() { a.lr = lr_pi; } + if let Some(a) = self.mhp_b_heads_adam.as_mut() { a.lr = lr_pi; } + if let Some(a) = self.mhp_w_gate_adam.as_mut() { a.lr = lr_pi; } + if let Some(a) = self.mhp_b_gate_adam.as_mut() { a.lr = lr_pi; } + } self.value_w_adam.lr = lr_v; self.value_b_adam.lr = lr_v; // SP20 P3 FRD head — LR from dedicated ISV slot (498-503 block). @@ -5405,9 +5589,24 @@ impl IntegratedTrainer { self.dqn_head .forward(&self.sampled_h_t_d, b_size, &mut self.q_logits_d) .context("dqn_head.forward(sampled_h_t) [R7d off-policy]")?; - self.policy_head - .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) - .context("policy_head.forward_logits")?; + // Phase 2A-C: gate π forward between legacy single-head and + // MultiHeadPolicy K-head mixture. Host branch runs during graph + // CAPTURE only; the captured kernel sequence bakes the chosen + // path so replay is single-path. `regime_h` is read DIRECTLY + // from perception's `step_regime_d` (parallel channel — bypasses + // VSN). + if self.use_multi_head_policy { + let regime_h = self.perception.step_regime_d_view(); + self.multi_head_policy + .as_ref() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)") + .forward(h_t_borrow, regime_h, &mut self.pi_logits_d) + .context("multi_head_policy.forward [step_synthetic_body]")?; + } else { + self.policy_head + .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) + .context("policy_head.forward_logits")?; + } self.value_head .forward(h_t_borrow, &self.isv_dev_ptr, b_size, &mut self.v_pred_d) .context("value_head.forward")?; @@ -5604,19 +5803,52 @@ impl IntegratedTrainer { // π gradient is fully determined by Q→π distillation + SAC entropy // (rl_q_pi_distill_grad above). No PPO surrogate or KL penalty. - self.policy_head - .backward_to_w_b_h( - h_t_borrow, - &self.ss_pi_grad_logits_d, - b_size, - &mut self.ss_pi_grad_w_per_batch_d, - &mut self.ss_pi_grad_b_per_batch_d, - &mut self.ss_pi_grad_h_t_d, - ) - .context("policy_head.backward_to_w_b_h")?; + // Phase 2A-C: gate π backward between legacy and MultiHeadPolicy. + // Both paths consume `ss_pi_grad_logits_d` (Q-distill output) + // and produce reduced grads ready for Adam below. The + // MultiHeadPolicy path internally: + // * distributes grad through mixture + per-head softmax, + // * adds β-weighted aux KL prior gradient, + // * distributes through gating softmax, + // * reduces per-batch grads to [K × N_ACTIONS × HIDDEN_DIM], + // [K × N_ACTIONS], [K × REGIME_DIM], [K] (ready for Adam). + // Legacy path uses the existing single-head GEMM-style backward. + if self.use_multi_head_policy { + // Split disjoint borrows: `perception`, `ss_pi_grad_logits_d`, + // and `multi_head_policy` are separate fields, so the borrow + // checker can prove non-aliasing. `isv_dev_ptr` is `u64` (Copy). + let isv_dev_ptr = self.isv_dev_ptr; + let regime_h: &CudaSlice = self.perception.step_regime_d_view(); + let grad_pi: &CudaSlice = &self.ss_pi_grad_logits_d; + let mhp = self + .multi_head_policy + .as_mut() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)"); + mhp.backward(grad_pi, h_t_borrow, regime_h, isv_dev_ptr) + .context("multi_head_policy.backward [step_synthetic_body]")?; + // NOTE: per Phase 2A-B comment in `MultiHeadPolicy::backward`, + // `grad_h_t_scratch_d` is NOT folded into the encoder grad + // here — Phase 2A-C wires the encoder path conservatively + // (encoder still receives V-head + auxiliaries; π gradient + // into encoder via multi-head route would change the + // encoder learning signal materially and is deferred to a + // future phase per the dispatch's "no incidental changes" + // constraint). + } else { + self.policy_head + .backward_to_w_b_h( + h_t_borrow, + &self.ss_pi_grad_logits_d, + b_size, + &mut self.ss_pi_grad_w_per_batch_d, + &mut self.ss_pi_grad_b_per_batch_d, + &mut self.ss_pi_grad_h_t_d, + ) + .context("policy_head.backward_to_w_b_h")?; - reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_w_per_batch_d, b_size, N_ACTIONS * HIDDEN_DIM, &mut self.ss_pi_grad_w_d)?; - reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut self.ss_pi_grad_b_d)?; + reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_w_per_batch_d, b_size, N_ACTIONS * HIDDEN_DIM, &mut self.ss_pi_grad_w_d)?; + reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut self.ss_pi_grad_b_d)?; + } // ── Step 8: V backward ─────────────────────────────────────── self.value_head @@ -5644,12 +5876,45 @@ impl IntegratedTrainer { self.dqn_b_adam .step(&mut self.dqn_head.b_d, &self.ss_q_grad_b_d) .context("dqn_b_adam.step")?; - self.policy_w_adam - .step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d) - .context("policy_w_adam.step")?; - self.policy_b_adam - .step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d) - .context("policy_b_adam.step")?; + // Phase 2A-C: gate the policy Adam step between legacy single-head + // and MultiHeadPolicy four-tensor update. Both branches preserve + // the existing LR sourcing (slot 412 `RL_LR_PI_INDEX`) — see the + // `lr_pi` assignment near line 5345. + if self.use_multi_head_policy { + let mhp = self + .multi_head_policy + .as_mut() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)"); + // Same LR for all four tensors (mirrors legacy w/b sharing + // a single `lr_pi`). Future controller could split. + self.mhp_w_heads_adam + .as_mut() + .expect("mhp_w_heads_adam None with flag on") + .step(&mut mhp.w_heads_d, &mhp.grad_w_heads_d) + .context("mhp_w_heads_adam.step")?; + self.mhp_b_heads_adam + .as_mut() + .expect("mhp_b_heads_adam None with flag on") + .step(&mut mhp.b_heads_d, &mhp.grad_b_heads_d) + .context("mhp_b_heads_adam.step")?; + self.mhp_w_gate_adam + .as_mut() + .expect("mhp_w_gate_adam None with flag on") + .step(&mut mhp.w_gate_d, &mhp.grad_w_gate_d) + .context("mhp_w_gate_adam.step")?; + self.mhp_b_gate_adam + .as_mut() + .expect("mhp_b_gate_adam None with flag on") + .step(&mut mhp.b_gate_d, &mhp.grad_b_gate_d) + .context("mhp_b_gate_adam.step")?; + } else { + self.policy_w_adam + .step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d) + .context("policy_w_adam.step")?; + self.policy_b_adam + .step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d) + .context("policy_b_adam.step")?; + } self.value_w_adam .step(&mut self.value_head.w_d, &self.ss_v_grad_w_d) .context("value_w_adam.step")?; @@ -6928,9 +7193,21 @@ impl IntegratedTrainer { } // ── Step 2b: Forward π logits for log_pi_old. ───────────────── - self.policy_head - .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) - .context("step_with_lobsim: policy_head.forward_logits")?; + // Phase 2A-C: same gate as step_synthetic_body. Inside prefill + // graph capture; host branch runs once and bakes the chosen + // forward kernel sequence into the replayed graph. + if self.use_multi_head_policy { + let regime_h = self.perception.step_regime_d_view(); + self.multi_head_policy + .as_ref() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)") + .forward(h_t_borrow, regime_h, &mut self.pi_logits_d) + .context("multi_head_policy.forward [step_with_lobsim]")?; + } else { + self.policy_head + .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) + .context("step_with_lobsim: policy_head.forward_logits")?; + } // R9 audit — Q-vs-π agreement diag fires into ISV[407] EMA. // Needs both q_logits_d (from step 2) and pi_logits_d (just @@ -9082,9 +9359,19 @@ impl IntegratedTrainer { } // pi logits - self.policy_head - .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) - .context("step_with_lobsim_gpu: policy_head.forward_logits")?; + // Phase 2A-C: gate the GPU-path pi forward as well. + if self.use_multi_head_policy { + let regime_h = self.perception.step_regime_d_view(); + self.multi_head_policy + .as_ref() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)") + .forward(h_t_borrow, regime_h, &mut self.pi_logits_d) + .context("multi_head_policy.forward [step_with_lobsim_gpu_body]")?; + } else { + self.policy_head + .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) + .context("step_with_lobsim_gpu: policy_head.forward_logits")?; + } // Q-vs-pi agree { @@ -9652,9 +9939,21 @@ impl IntegratedTrainer { .context("eval_policy_probs: forward_encoder")?; let h_t = self.perception.h_t_view(); let mut logits_d = self.stream.alloc_zeros::(b_size * N_ACTIONS)?; - self.policy_head - .forward_logits(h_t, b_size, &mut logits_d) - .context("eval_policy_probs: policy_head.forward_logits")?; + // Phase 2A-C: gate the eval-time pi forward. NOT inside a graph + // capture (eval is single-shot), so the branch is run host-side + // every call — cost is one `bool` read. + if self.use_multi_head_policy { + let regime_h = self.perception.step_regime_d_view(); + self.multi_head_policy + .as_ref() + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)") + .forward(h_t, regime_h, &mut logits_d) + .context("multi_head_policy.forward [eval_policy_probs_per_action]")?; + } else { + self.policy_head + .forward_logits(h_t, b_size, &mut logits_d) + .context("eval_policy_probs: policy_head.forward_logits")?; + } let logits_host = read_slice_d(&self.stream, &logits_d, b_size * N_ACTIONS)?; let mut probs = vec![0.0_f32; N_ACTIONS]; @@ -11004,7 +11303,7 @@ impl IntegratedTrainer { let replay_len = self.gpu_replay.capacity.min((step as usize) + 1); - let record = serde_json::json!({ + let mut record = serde_json::json!({ "step": step, "elapsed_s": elapsed_s, "loss": { @@ -11484,6 +11783,36 @@ impl IntegratedTrainer { "mamba2_grad_w_c_l1": checksums[31], }, }); + + // ── Phase 2A-C (2026-06-03) MultiHeadPolicy diag emit ──────── + // + // Only emitted when `FOXHUNT_USE_MULTI_HEAD_POLICY=1` so the + // flag-off diag.jsonl stays byte-equal to HEAD `e22da61cf` + // (verification gate C.7.4). + // + // Phase 2A-C ships ISV-resident summary leaves only — K, + // gating/head entropy floors, aux β. The richer per-step + // device-aggregated stats (gate_probs_mean[K], gate_argmax_mass[K], + // gate_entropy_mean, per_head_entropy_mean[K]) require a new + // reduction kernel + new ISV slots and are deferred per the + // dispatch's STOP-on-surprise rule. The four-leaf summary here + // is the minimum cluster-verifiable signal that the gate path + // is active and bootstrapped correctly. + if self.use_multi_head_policy { + if let Some(obj) = record.as_object_mut() { + obj.insert( + "multi_head_policy".to_string(), + serde_json::json!({ + "active": self.use_multi_head_policy, + "k": isv[crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX], + "gating_entropy_floor": isv[crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX], + "head_entropy_floor": isv[crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX], + "aux_prior_beta": isv[crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX], + }), + ); + } + } + Ok(record) } } diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index a4f0716b3..03c4748e1 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -4169,6 +4169,22 @@ impl PerceptionTrainer { &self.h_t_d } + /// Phase 2A-C (2026-06-03) — read-only accessor for the per-batch raw + /// regime input `[n_batch × REGIME_DIM=6]` that the snap-feature + /// pipeline populates per step. Consumed by `MultiHeadPolicy::forward` + /// as a PARALLEL channel — bypassing VSN — to gate K policy heads on + /// vol/popart regime directly (`pearl_local_smoke_noise_floor_and_regime_concentration`). + /// + /// The slice's lifetime tracks `&self`; the buffer is overwritten on + /// every `snap_feature_assemble` so callers must finish their kernel + /// dispatch before the next encoder step. In the captured training + /// graph, multi-head policy forward runs in the same stream sequence + /// immediately after the encoder, so this is single-stream-safe by + /// construction (same property `h_t_view` relies on). + pub fn step_regime_d_view(&self) -> &CudaSlice { + &self.step_regime_d + } + /// Determinism Phase 2.2 sub-investigation: read-only accessors for /// encoder intermediate activations. Used exclusively by /// `IntegratedTrainer::compute_diag_checksums` to checksum each layer's