From 5f10bcde3a1084e3a42495aeac7aaf6626816874 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 3 Jun 2026 12:15:39 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20Phase=202A-C+=20=E2=80=94=20d?= =?UTF-8?q?evice-aggregated=20gate=20diag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the specialization-signal diagnostics that Phase 2A-C deferred. 4 new stats are emitted under policy_diagnostic.multi_head_policy.* when FOXHUNT_USE_MULTI_HEAD_POLICY=1: gate_probs_mean[K], gate_argmax_mass[K], gate_entropy_mean, per_head_entropy_mean[K]. These are the load-bearing signals for Phase 2A-D's verdict — without them we'd see a pnl delta but couldn't distinguish "the mixture genuinely specializes by regime" from "the mixture accidentally acts as a single-head regularizer". ## ISV slots (25 new) * 765-772 RL_POLICY_GATE_PROBS_MEAN_BASE (8 slots, MAX_K_HEADS=8 stride) * 773-780 RL_POLICY_GATE_ARGMAX_MASS_BASE * 781 RL_POLICY_GATE_ENTROPY_MEAN * 782-789 RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE * RL_SLOTS_END = 790 ## Aggregator kernel (multi_head_policy_aggregate_diag.cu) * Grid = (MAX_K_HEADS + 1, 1, 1) = 9 blocks. Block = (128, 1, 1). * Per-head blocks (k_block 0..7): if k_block ≥ runtime K, thread 0 writes 0.0 to its two ISV destinations. Else parallel-sum over B in shared memory, tree-reduce by halving stride, thread 0 writes gate_probs_mean[k] + per_head_entropy_mean[k]. * Global block (k_block = MAX_K_HEADS): single block computes gate entropy + argmax-mass in one pass. Per-thread argmax counter array in registers scatters to s_am[MAX_K_HEADS][BLOCK_THREADS] shared mem; tree-reduce; thread 0 writes 9 ISV scalars. * No-atomicAdd: every ISV destination has a single writer thread. Tree-reduce via shared memory + __syncthreads(). Deterministic fixed-order pairwise sum. ## Trainer wiring * MultiHeadPolicy::emit_diag_stats(isv_dev_ptr) launches the aggregator on the struct stream. Called at all 3 train-path forward sites in integrated.rs immediately after mhp.forward() (gate_probs and pi_probs_k are forward outputs — must aggregate before next step's forward overwrites them). * build_diag_value emits the 4 fields inside the existing multi_head_policy.* block under if self.use_multi_head_policy. Flag-off schema unchanged (multi_head_policy key absent). ## Verification (all gates pass) * 13/13 invariants: 11 pre-existing + 2 new - aggregate_diag_gate_probs_mean_correct: uniform 1/K + asymmetric ramp case, max abs err < 1e-5 - aggregate_diag_entropy_pins_top_and_bottom: top (uniform gate → entropy = log(K) = 1.0986; uniform pi → per-head = log(11) = 2.398); bottom (one-hot gate → 0; one-hot pi → 0). Tie-broken- low argmax verified. * FOXHUNT_USE_MULTI_HEAD_POLICY=0 determinism-check.sh --quick: exit 0. Flag-off diag.jsonl preserved (multi_head_policy key ABSENT, EXPECTED_LEAVES=712 unchanged). * FOXHUNT_USE_MULTI_HEAD_POLICY=1 determinism-check.sh --quick: exit 0. Aggregator is deterministic. * 200-step flag-on smoke shows the diag working as designed: - gate_probs_mean ≈ [0.328, 0.339, 0.333, 0,0,0,0,0] sum ≈ 1.0 - gate_argmax_mass ≈ [0, 1.0, 0, ...] (Head 1 Long-bias dominant at init — matches init bias asymmetry) - gate_entropy_mean ≈ 1.0985 (at max log(3) ≈ 1.0986 — gate has NOT collapsed) - per_head_entropy_mean ≈ [2.36, 2.28, 2.29] vs max 2.398 — heads slightly less than uniform, expected at random init * Pre-commit: 0 atomicAdd, 0 raw memcpy_htod/dtoh, 0 TODO. ## Surprises handled None — all STOP-on-surprise predictions held: 1. Slot bootstrap zero-init contributes 0² to isv_state checksum, so unconditional bootstrap preserves flag-off bit-equality without needing a flag-gated block (unlike the 2A-C ISV surprise). 2. EXPECTED_LEAVES (712) preserved — flag-off schema unchanged. 3. K consistency (Rust MAX_K_HEADS=8 + CUDA #define) maintained. 4. emit_diag_stats placed after mhp.forward() and before any downstream consumer that would overwrite forward outputs. 5. CUDA graph capture: aggregator launch is deterministic (fixed grid/block, fixed reductions). Replay single-path. No conflict. ## Linked * Phase 2A-A: 0b3e40150 (MultiHeadPolicy foundation, inert) * Phase 2A-B: e22da61cf (backward + aux KL prior) * Phase 2A-C: 3e36f4a0e (trainer integration behind flag) * Plan: docs/superpowers/plans/2026-06-03-multi-head-policy- implementation.md * Spec ADDENDUM §R.5 (falsification gates for specialization) Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + .../cuda/multi_head_policy_aggregate_diag.cu | 200 +++++++++++++ crates/ml-alpha/src/rl/isv_slots.rs | 49 +++- crates/ml-alpha/src/rl/multi_head_policy.rs | 72 +++++ crates/ml-alpha/src/trainer/integrated.rs | 97 ++++-- .../tests/multi_head_policy_invariants.rs | 275 +++++++++++++++++- 6 files changed, 666 insertions(+), 28 deletions(-) create mode 100644 crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 99c443dbf..4beb147a8 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -159,6 +159,7 @@ const KERNELS: &[&str] = &[ "multi_head_policy_gate_forward",// Phase 2A-A (2026-06-03): gating head from raw regime features [B × 6] (parallel channel bypassing VSN/Mamba2). Spec ADDENDUM 2026-06-03 §R.3 Option B "multi_head_policy_backward", // Phase 2A-B (2026-06-03): backward through mixture + per-head softmax + gating softmax. Two kernels: `_backward_pi` and `_backward_gate`. Per-batch grad scratch; caller reduces via reduce_axis0. Spec ADDENDUM 2026-06-03 §R.6 "multi_head_policy_aux_prior", // Phase 2A-B (2026-06-03): per-head auxiliary KL prior gradient — additive into grad_pi_logits_k. β read from ISV slot 764. Spec ADDENDUM 2026-06-03 §R.4 + "multi_head_policy_aggregate_diag", // Phase 2A-D (2026-06-03): device-aggregated gate-and-head diagnostics. Reduces forward outputs (gate_probs, pi_probs_k) over B → 25 ISV slots (gate_probs_mean[K] / gate_argmax_mass[K] / gate_entropy_mean / per_head_entropy_mean[K]) for the multi-head specialization verdict signal. Tree-reduce, no atomicAdd. ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu b/crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu new file mode 100644 index 000000000..c79e03a52 --- /dev/null +++ b/crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu @@ -0,0 +1,200 @@ +// multi_head_policy_aggregate_diag.cu — Phase 2A-D (2026-06-03) +// +// Device-aggregated gate-and-head diagnostics for the MultiHeadPolicy +// path. Reads forward outputs (`gate_probs [B × K]`, `pi_probs_k +// [B × K × N_ACTIONS]`) and writes 25 scalars into the ISV bus: +// +// 765..773 gate_probs_mean[K] — mean over B of gate_probs[b,k] +// 773..781 gate_argmax_mass[K] — fraction of B where head k is argmax +// 781 gate_entropy_mean — mean over B of −Σ_k p·log p +// 782..790 per_head_entropy_mean[K] — mean over B of −Σ_a π_k·log π_k +// +// Tail entries past runtime K (in the 8-stride arrays) are written 0.0 +// so the diag JSON schema is constant for K ∈ [1, 8] (matches the +// MAX_K_HEADS=8 mirror in `rl/multi_head_policy.rs` and the kernel +// `#define MAX_K_HEADS 8` in `multi_head_policy_forward.cu` / +// `multi_head_policy_gate_forward.cu`). +// +// Constraints honoured: +// * `feedback_no_atomicadd` — no atomic ops. Each ISV destination cell +// has a single writer (thread 0 of the responsible block). +// * `feedback_cpu_is_read_only` — pure device kernel. +// * `feedback_nvidia_grade_perf_for_kernels` — no host branches in +// capture path; deterministic launch shape. +// * `pearl_fleet_fraction_not_aggregate` — these are explicit +// per-batch-fraction stats (mean over B / argmax-mass over B), not +// aggregate scalars hiding per-batch state. +// +// Determinism: every reduction uses a fixed-order sequential sum within +// a single thread. Inputs (`gate_probs`, `pi_probs_k`) are forward +// outputs already on the launching stream; aggregator runs on the same +// stream so the producer→consumer ordering is single-stream. +// +// Launch config: +// grid = (MAX_K_HEADS + 1, 1, 1) = 9 blocks +// block = (BLOCK_THREADS, 1, 1) = 128 threads +// smem = 0 (small static shared arrays declared inside) +// +// Block role assignment: +// blockIdx.x = 0..MAX_K_HEADS-1 → per-head stats for head k = blockIdx.x. +// Writes gate_probs_mean[k] + +// per_head_entropy_mean[k]. Tail +// heads (k >= runtime K) write 0.0. +// blockIdx.x = MAX_K_HEADS → global stats. Computes +// gate_entropy_mean + +// gate_argmax_mass[k] for all K +// heads. Tail entries write 0.0. + +#include + +#define MAX_K_HEADS 8 +#define N_ACTIONS 11 +#define BLOCK_THREADS 128 + +extern "C" __global__ void multi_head_policy_aggregate_diag( + const float* __restrict__ gate_probs, // [B × K] + const float* __restrict__ pi_probs_k, // [B × K × N_ACTIONS] + float* __restrict__ isv, // ISV bus + int b_size, + int k_runtime, // ∈ [1, MAX_K_HEADS] + int gate_probs_mean_base, // 765 + int gate_argmax_mass_base,// 773 + int gate_entropy_mean_slot,// 781 + int per_head_entropy_mean_base // 782 +) { + const int k_block = blockIdx.x; // ∈ [0, MAX_K_HEADS] + const int tid = threadIdx.x; // ∈ [0, BLOCK_THREADS) + const float inv_b = (b_size > 0) ? (1.0f / (float)b_size) : 0.0f; + + // ── Per-head blocks (k_block = 0 .. MAX_K_HEADS-1) ──────────────── + if (k_block < MAX_K_HEADS) { + // Tail heads past runtime K — sole writer (thread 0) clears its + // two ISV destinations to 0.0 and exits. + if (k_block >= k_runtime) { + if (tid == 0) { + isv[gate_probs_mean_base + k_block] = 0.0f; + isv[per_head_entropy_mean_base + k_block] = 0.0f; + } + return; + } + + // ── Stat 1: gate_probs_mean[k_block] ───────────────────────── + // Each thread sums a strided subset of batches; tree-reduce in + // shared memory. + __shared__ float s_gpm[BLOCK_THREADS]; + __shared__ float s_phe[BLOCK_THREADS]; + float my_gpm = 0.0f; + float my_phe = 0.0f; + + for (int b = tid; b < b_size; b += BLOCK_THREADS) { + // gate_probs is [B × K] row-major over (b, k). + const float p = gate_probs[b * k_runtime + k_block]; + my_gpm += p; + + // ── Stat 2: per-head entropy ───────────────────────────── + // pi_probs_k is [B × K × N_ACTIONS] row-major over + // (b, k, a). Per-head per-batch entropy: + // H_k(b) = -Σ_a π_k[b,a] · log(π_k[b,a]) + // Guard against log(0); the forward kernel writes a softmax + // output so π_k[b,a] > 0 (modulo fp32 underflow), but be + // defensive — entries < 1e-12 are clamped to a 0 contribution + // (matches the standard entropy convention 0·log(0) = 0). + const int row = (b * k_runtime + k_block) * N_ACTIONS; + float h = 0.0f; + for (int a = 0; a < N_ACTIONS; ++a) { + const float pa = pi_probs_k[row + a]; + if (pa > 1e-12f) { + h -= pa * logf(pa); + } + } + my_phe += h; + } + s_gpm[tid] = my_gpm; + s_phe[tid] = my_phe; + __syncthreads(); + + // Tree-reduce in shared memory. Fixed-order sequential pairwise + // sum — deterministic across launches (no warp shuffles, no + // atomic adds). + for (int stride = BLOCK_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + s_gpm[tid] += s_gpm[tid + stride]; + s_phe[tid] += s_phe[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + isv[gate_probs_mean_base + k_block] = s_gpm[0] * inv_b; + isv[per_head_entropy_mean_base + k_block] = s_phe[0] * inv_b; + } + return; + } + + // ── Global block (k_block == MAX_K_HEADS) ───────────────────────── + // Computes gate_entropy_mean (scalar) and gate_argmax_mass[K] (array). + // + // The argmax tally must avoid atomicAdd. Each thread maintains a + // private per-head counter in registers (small array indexed by + // 0..MAX_K_HEADS-1) while iterating its stride of batches. After + // the per-thread loop, tree-reduce each counter across threads via + // shared memory. + + __shared__ float s_ge[BLOCK_THREADS]; // gate entropy partial + __shared__ float s_am[MAX_K_HEADS][BLOCK_THREADS]; // argmax counts per (k, tid) + + int my_am[MAX_K_HEADS]; + float my_ge = 0.0f; + #pragma unroll + for (int k = 0; k < MAX_K_HEADS; ++k) my_am[k] = 0; + + for (int b = tid; b < b_size; b += BLOCK_THREADS) { + // Walk this batch row once: argmax + entropy in one pass. + int argmax_k = 0; + float max_p = gate_probs[b * k_runtime + 0]; + float h = 0.0f; + if (max_p > 1e-12f) h -= max_p * logf(max_p); + for (int k = 1; k < k_runtime; ++k) { + const float p = gate_probs[b * k_runtime + k]; + // Ties broken by lowest index — `>` not `>=`. Deterministic. + if (p > max_p) { + max_p = p; + argmax_k = k; + } + if (p > 1e-12f) h -= p * logf(p); + } + my_am[argmax_k] += 1; + my_ge += h; + } + + // Scatter per-thread counters into shared memory for tree reduction. + s_ge[tid] = my_ge; + #pragma unroll + for (int k = 0; k < MAX_K_HEADS; ++k) { + s_am[k][tid] = (float)my_am[k]; + } + __syncthreads(); + + for (int stride = BLOCK_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + s_ge[tid] += s_ge[tid + stride]; + #pragma unroll + for (int k = 0; k < MAX_K_HEADS; ++k) { + s_am[k][tid] += s_am[k][tid + stride]; + } + } + __syncthreads(); + } + + if (tid == 0) { + // gate_entropy_mean — scalar. + isv[gate_entropy_mean_slot] = s_ge[0] * inv_b; + + // gate_argmax_mass[k] — fraction over B. Tail entries (k >= K) + // are 0 by construction (my_am[k>=K] is never incremented). + #pragma unroll + for (int k = 0; k < MAX_K_HEADS; ++k) { + isv[gate_argmax_mass_base + k] = s_am[k][0] * inv_b; + } + } +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index 8a9aca3df..51c1b1a53 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -1840,6 +1840,52 @@ pub const RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX: usize = 763; /// (specialization driven only by gate + init bias). pub const RL_POLICY_AUX_PRIOR_BETA_INDEX: usize = 764; +// ───────────────────────────────────────────────────────────────────── +// Phase 2A-D (2026-06-03) — Device-aggregated gate-and-head diagnostics. +// Spec: dispatch 2026-06-03 §"add device-aggregated gate-and-head +// diagnostics to the MultiHeadPolicy path" +// Pearl: pearl_fleet_fraction_not_aggregate — these stats ARE the +// fleet-fraction pattern applied to per-batch gate routing +// (mean / argmax-mass / entropy across B). +// +// All 25 slots are produced by `multi_head_policy_aggregate_diag.cu` on +// every forward call (flag-on path only). With FOXHUNT_USE_MULTI_HEAD_POLICY +// off the slots stay sentinel-zero (MappedF32Buffer zero-init) — the +// `isv_state` checksum still includes them but adds 0² to the running +// sum-of-squares, so flag-off bit-equality vs HEAD 3e36f4a0e is preserved. +// +// MAX_K_HEADS=8 (mirror of `crate::rl::multi_head_policy::MAX_K_HEADS`) +// is the slot stride for the K-shaped arrays so the JSON schema stays +// constant when K varies in [1, 8]. Tail entries (k ≥ runtime K) are +// kernel-written 0.0. +// ───────────────────────────────────────────────────────────────────── + +/// Per-head mean gate probability: `mean_b gate_probs[b, k]` for k ∈ [0, 8). +/// Base of an 8-slot K-stride array (765..773). Tail entries past runtime K +/// are written 0.0 by the aggregator kernel. +/// Diag emit: `policy_diagnostic.multi_head_policy.gate_probs_mean[K]`. +/// Invariant (at runtime K): Σ_k gate_probs_mean[k] ≈ 1.0 (softmax-sum). +pub const RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX: usize = 765; + +/// Per-head argmax-mass: fraction of batches where head k is the gate +/// argmax. Base of an 8-slot K-stride array (773..781). Tail entries +/// past runtime K are written 0.0. Invariant (at runtime K): +/// Σ_k gate_argmax_mass[k] ≈ 1.0 (every batch contributes exactly one +/// argmax; ties broken by lowest index). +pub const RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX: usize = 773; + +/// Mean gate entropy over batches: `mean_b (-Σ_k p[b,k]·log(p[b,k]))`. +/// Scalar slot 781. Range [0, log(K)]; log(3) ≈ 1.0986 for default K=3. +/// Collapses toward 0 if the gate concentrates on one head. +pub const RL_POLICY_GATE_ENTROPY_MEAN_INDEX: usize = 781; + +/// Per-head mean action entropy: `mean_b (-Σ_a π_k[b,a]·log(π_k[b,a]))` +/// for head k. Base of an 8-slot K-stride array (782..790). Tail entries +/// past runtime K are written 0.0. Range [0, log(N_ACTIONS)]; +/// log(11) ≈ 2.398. Collapses toward 0 if a head determinizes onto a +/// single action. +pub const RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX: usize = 782; + /// Last RL-allocated slot index (exclusive). /// Pre-risk-stack: 662. Post-Fix-B: 685. Post-v9 (warmup boundary): 696. /// Post-regime-observer (F1.1): 716. Post-warmed-flag addendum: 717. @@ -1855,4 +1901,5 @@ pub const RL_POLICY_AUX_PRIOR_BETA_INDEX: usize = 764; /// Post-surfer-scaffold v5 adaptive (3 config slots): 757. /// Post-rollout-buffer+GAE Phase 1B-A (3 PPO control slots): 760. /// Post-multi-head-policy Phase 2A-A (4 policy-gating slots): 764. -pub const RL_SLOTS_END: usize = 765; +/// Post-multi-head-policy Phase 2A-D (25 aggregate diag slots): 789. +pub const RL_SLOTS_END: usize = 790; diff --git a/crates/ml-alpha/src/rl/multi_head_policy.rs b/crates/ml-alpha/src/rl/multi_head_policy.rs index e9838bce7..58bab04f6 100644 --- a/crates/ml-alpha/src/rl/multi_head_policy.rs +++ b/crates/ml-alpha/src/rl/multi_head_policy.rs @@ -78,6 +78,10 @@ const AUX_PRIOR_CUBIN: &[u8] = include_bytes!(concat!( )); const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); +const AGGREGATE_DIAG_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/multi_head_policy_aggregate_diag.cubin" +)); /// Compile-time upper bound on K — mirrors the `#define MAX_K_HEADS 8` in /// both kernel sources. The runtime K (from `RL_POLICY_NUM_HEADS_INDEX`) @@ -132,6 +136,8 @@ pub struct MultiHeadPolicy { aux_prior_fn: CudaFunction, _reduce_axis0_module: Arc, reduce_axis0_fn: CudaFunction, + _aggregate_diag_module: Arc, + aggregate_diag_fn: CudaFunction, // ── Policy-head weights ([K × N_ACTIONS × HIDDEN_DIM] row-major) ── /// `W_heads[k, a, j]` — row-major flat over `(k, a, j)`. @@ -268,6 +274,12 @@ impl MultiHeadPolicy { let reduce_axis0_fn = reduce_axis0_module .load_function("reduce_axis0") .context("load reduce_axis0 fn (multi_head_policy)")?; + let aggregate_diag_module = ctx + .load_cubin(AGGREGATE_DIAG_CUBIN.to_vec()) + .context("load multi_head_policy_aggregate_diag cubin")?; + let aggregate_diag_fn = aggregate_diag_module + .load_function("multi_head_policy_aggregate_diag") + .context("load multi_head_policy_aggregate_diag fn")?; // Install determinism guard before drawing any random samples // (pearl_scoped_init_seed_for_reproducibility). @@ -402,6 +414,8 @@ impl MultiHeadPolicy { aux_prior_fn, _reduce_axis0_module: reduce_axis0_module, reduce_axis0_fn, + _aggregate_diag_module: aggregate_diag_module, + aggregate_diag_fn, w_heads_d, b_heads_d, w_gate_d, @@ -533,6 +547,64 @@ impl MultiHeadPolicy { Ok(()) } + /// Launch the device-aggregated diagnostics kernel (Phase 2A-D). + /// + /// Reduces the forward outputs `gate_probs_d [B × K]` and + /// `pi_probs_k_d [B × K × N_ACTIONS]` over the batch axis and writes + /// 25 scalars into the ISV bus: + /// + /// * `gate_probs_mean[k]` — slots 765..773 (K-stride 8) + /// * `gate_argmax_mass[k]` — slots 773..781 (K-stride 8) + /// * `gate_entropy_mean` — slot 781 (scalar) + /// * `per_head_entropy_mean[k]` — slots 782..790 (K-stride 8) + /// + /// Tail entries past runtime K are written 0.0 so the diag JSON + /// schema is constant for K ∈ [1, 8] (`MAX_K_HEADS` mirror). + /// + /// Call **after** `forward()` and **before** the next step's forward + /// — the kernel reads forward outputs that are overwritten each step. + /// Launches on the struct's stream so the producer→consumer ordering + /// is single-stream (no extra sync needed). + /// + /// Per `feedback_no_atomicadd`: tree-reduce in shared memory; each + /// ISV destination has a single writer thread. + pub fn emit_diag_stats(&self, isv_dev_ptr: u64) -> Result<()> { + let b = self.cfg.b_size; + let k = self.cfg.k; + let b_i = b as i32; + let k_i = k as i32; + + // Launch shape: Grid = (MAX_K_HEADS + 1, 1, 1), Block = (128, 1, 1). + // Matches the role assignment in `multi_head_policy_aggregate_diag.cu`. + const BLOCK_THREADS: u32 = 128; + let grid_x: u32 = (MAX_K_HEADS as u32) + 1; + + let mut args = RawArgs::new(); + args.push_ptr(self.gate_probs_d.raw_ptr()); + args.push_ptr(self.pi_probs_k_d.raw_ptr()); + args.push_ptr(isv_dev_ptr); + args.push_i32(b_i); + args.push_i32(k_i); + args.push_i32(crate::rl::isv_slots::RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_POLICY_GATE_ENTROPY_MEAN_INDEX as i32); + args.push_i32(crate::rl::isv_slots::RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX as i32); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.aggregate_diag_fn.cu_function(), + (grid_x, 1, 1), + (BLOCK_THREADS, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_aggregate_diag: {e:?}"))?; + } + + Ok(()) + } + /// Backward pass through the mixture, per-head softmax, and gating softmax. /// /// Inputs: diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index a14a93d0a..901d5889d 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -5597,11 +5597,17 @@ impl IntegratedTrainer { // VSN). if self.use_multi_head_policy { let regime_h = self.perception.step_regime_d_view(); - self.multi_head_policy + let mhp = 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) + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)"); + mhp.forward(h_t_borrow, regime_h, &mut self.pi_logits_d) .context("multi_head_policy.forward [step_synthetic_body]")?; + // Phase 2A-D (2026-06-03) — device-aggregated gate/head diagnostics. + // Reads forward outputs (gate_probs_d, pi_probs_k_d) before any + // downstream consumer; writes 25 ISV slots (765..790) for the + // diag emit at step N+1. Single-stream — no extra sync needed. + mhp.emit_diag_stats(self.isv_dev_ptr) + .context("multi_head_policy.emit_diag_stats [step_synthetic_body]")?; } else { self.policy_head .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) @@ -7198,11 +7204,14 @@ impl IntegratedTrainer { // 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 + let mhp = 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) + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)"); + mhp.forward(h_t_borrow, regime_h, &mut self.pi_logits_d) .context("multi_head_policy.forward [step_with_lobsim]")?; + // Phase 2A-D — see comment in step_synthetic_body for rationale. + mhp.emit_diag_stats(self.isv_dev_ptr) + .context("multi_head_policy.emit_diag_stats [step_with_lobsim]")?; } else { self.policy_head .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) @@ -9362,11 +9371,16 @@ impl IntegratedTrainer { // 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 + let mhp = 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) + .expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)"); + mhp.forward(h_t_borrow, regime_h, &mut self.pi_logits_d) .context("multi_head_policy.forward [step_with_lobsim_gpu_body]")?; + // Phase 2A-D — see comment in step_synthetic_body. This forward + // path runs inside the captured CUDA graph; the aggregator launch + // is part of the graph and replays deterministically. + mhp.emit_diag_stats(self.isv_dev_ptr) + .context("multi_head_policy.emit_diag_stats [step_with_lobsim_gpu_body]")?; } else { self.policy_head .forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d) @@ -11784,30 +11798,63 @@ impl IntegratedTrainer { }, }); - // ── Phase 2A-C (2026-06-03) MultiHeadPolicy diag emit ──────── + // ── Phase 2A-C / 2A-D (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). + // flag-off diag.jsonl stays byte-equal to HEAD `3e36f4a0e`. // - // 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. + // Phase 2A-C: ISV-resident summary leaves (K, gating/head entropy + // floors, aux β). Phase 2A-D adds the device-aggregated gate- + // and-head specialization signals produced by the aggregator + // kernel `multi_head_policy_aggregate_diag`: + // + // * gate_probs_mean[K] — mean over B of gate_probs[b,k] + // * gate_argmax_mass[K] — fraction of B where head k is argmax + // * gate_entropy_mean — mean over B of −Σ_k p·log p + // * per_head_entropy_mean[K] — mean over B of −Σ_a π_k·log π_k + // + // The four array fields are emitted at MAX_K_HEADS=8 stride so + // the schema is constant for K ∈ [1, 8]; tail entries past + // runtime K are kernel-written 0.0. These four stats are the + // load-bearing specialization verdict signal — without them we + // can't tell whether the mixture is genuinely regime-conditional + // or accidentally acting as a single-head regularizer. if self.use_multi_head_policy { + use crate::rl::isv_slots::{ + RL_POLICY_AUX_PRIOR_BETA_INDEX, + RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX, + RL_POLICY_GATE_ENTROPY_MEAN_INDEX, + RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX, + RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, + RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, + RL_POLICY_NUM_HEADS_INDEX, + RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX, + }; + // MAX_K_HEADS stride — must match `multi_head_policy::MAX_K_HEADS` + // and the `#define MAX_K_HEADS 8` mirrors in the kernel sources. + const K_STRIDE: usize = crate::rl::multi_head_policy::MAX_K_HEADS; + let gpm_base = RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX; + let gam_base = RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX; + let phe_base = RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX; + let gate_probs_mean: Vec = + (0..K_STRIDE).map(|k| isv[gpm_base + k]).collect(); + let gate_argmax_mass: Vec = + (0..K_STRIDE).map(|k| isv[gam_base + k]).collect(); + let per_head_entropy_mean: Vec = + (0..K_STRIDE).map(|k| isv[phe_base + k]).collect(); 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], + "active": self.use_multi_head_policy, + "k": isv[RL_POLICY_NUM_HEADS_INDEX], + "gating_entropy_floor": isv[RL_POLICY_GATING_ENTROPY_FLOOR_INDEX], + "head_entropy_floor": isv[RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX], + "aux_prior_beta": isv[RL_POLICY_AUX_PRIOR_BETA_INDEX], + "gate_probs_mean": gate_probs_mean, + "gate_argmax_mass": gate_argmax_mass, + "gate_entropy_mean": isv[RL_POLICY_GATE_ENTROPY_MEAN_INDEX], + "per_head_entropy_mean": per_head_entropy_mean, }), ); } diff --git a/crates/ml-alpha/tests/multi_head_policy_invariants.rs b/crates/ml-alpha/tests/multi_head_policy_invariants.rs index 0b2236d73..27454c346 100644 --- a/crates/ml-alpha/tests/multi_head_policy_invariants.rs +++ b/crates/ml-alpha/tests/multi_head_policy_invariants.rs @@ -46,8 +46,15 @@ use ml_alpha::cfc::snap_features::REGIME_DIM; use ml_alpha::heads::HIDDEN_DIM; use ml_alpha::pinned_mem::MappedF32Buffer; use ml_alpha::rl::common::N_ACTIONS; -use ml_alpha::rl::isv_slots::{RL_POLICY_AUX_PRIOR_BETA_INDEX, RL_SLOTS_END}; -use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig}; +use ml_alpha::rl::isv_slots::{ + RL_POLICY_AUX_PRIOR_BETA_INDEX, + RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX, + RL_POLICY_GATE_ENTROPY_MEAN_INDEX, + RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX, + RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX, + RL_SLOTS_END, +}; +use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig, MAX_K_HEADS}; use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub}; use ml_core::device::MlDevice; use std::sync::Arc; @@ -1086,3 +1093,267 @@ fn backward_is_deterministic_across_contexts() -> Result<()> { ); Ok(()) } + +// ───────────────────────────────────────────────────────────────────── +// Phase 2A-D — device-aggregated gate-and-head diagnostics tests +// ───────────────────────────────────────────────────────────────────── +// +// These two tests exercise the `emit_diag_stats` kernel by synthesizing +// known `gate_probs` / `pi_probs_k` distributions directly into the +// MultiHeadPolicy's persistent forward-output buffers (bypassing the +// real forward kernel) and checking that the aggregator-written ISV +// slots match the analytical expectation. +// +// The synthetic write path is the minimal way to assert correctness of +// the reduction kernel itself, independent of the forward-pass softmax +// (already covered by `gate_probs_sum_to_one` / +// `pi_probs_mixture_sums_to_one`). + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn aggregate_diag_gate_probs_mean_correct() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 8usize; + let k = 3usize; + + let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) } + .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?; + for slot in 0..RL_SLOTS_END { + isv.write_record(slot, 0.0); + } + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0x2A2D_0001, + }, + )?; + + // ── Case 1: uniform gate (1/K per head) ────────────────────────── + // gate_probs_mean[k] = 1/K for every k ∈ [0, K), 0 for tail. + let uniform = 1.0_f32 / (k as f32); + let mut gate_probs_host = vec![uniform; b_size * k]; + // pi_probs_k: every (b, k) row a uniform 1/N_ACTIONS distribution + // (so per_head_entropy_mean[k] = log(N_ACTIONS) ≈ 2.398). + let uniform_a = 1.0_f32 / (N_ACTIONS as f32); + let pi_probs_k_host = vec![uniform_a; b_size * k * N_ACTIONS]; + write_slice_f32_d_pub(&stream, &gate_probs_host, &mut mhp.gate_probs_d)?; + write_slice_f32_d_pub(&stream, &pi_probs_k_host, &mut mhp.pi_probs_k_d)?; + + mhp.emit_diag_stats(isv_dev_ptr)?; + stream.synchronize()?; + + let gpm_base = RL_POLICY_GATE_PROBS_MEAN_BASE_INDEX; + for k_idx in 0..k { + let got = isv.read_record(gpm_base + k_idx); + assert!( + (got - uniform).abs() < 1e-5, + "uniform: gate_probs_mean[{k_idx}] = {got}, expected {uniform}" + ); + } + // Tail (k >= runtime K) must be zero. + for k_idx in k..MAX_K_HEADS { + let got = isv.read_record(gpm_base + k_idx); + assert_eq!( + got, 0.0, + "tail gate_probs_mean[{k_idx}] = {got}, expected exact 0.0" + ); + } + // ── Case 2: asymmetric per-batch (still deterministic mean) ───── + // + // gate_probs[b, k] = (b + k + 1) / Σ_k'(b + k' + 1) — row-normalized + // ramps. The analytical mean over B = mean_b [(b+k+1) / S_b] where + // S_b = Σ_k' (b + k' + 1) = K·(b+1) + K(K-1)/2. + gate_probs_host.fill(0.0); + let mut expected_gpm = vec![0.0_f64; k]; + for b in 0..b_size { + let s_b: f64 = (0..k).map(|k2| (b + k2 + 1) as f64).sum(); + for k2 in 0..k { + let p = ((b + k2 + 1) as f64) / s_b; + gate_probs_host[b * k + k2] = p as f32; + expected_gpm[k2] += p / (b_size as f64); + } + } + write_slice_f32_d_pub(&stream, &gate_probs_host, &mut mhp.gate_probs_d)?; + // Re-zero ISV slots so we read fresh writes. + for slot in gpm_base..gpm_base + MAX_K_HEADS { + isv.write_record(slot, 0.0); + } + + mhp.emit_diag_stats(isv_dev_ptr)?; + stream.synchronize()?; + + let mut max_err = 0.0_f64; + for k_idx in 0..k { + let got = isv.read_record(gpm_base + k_idx) as f64; + let err = (got - expected_gpm[k_idx]).abs(); + max_err = max_err.max(err); + assert!( + err < 1e-5, + "asymmetric: gate_probs_mean[{k_idx}] = {got}, expected {} (err {:.2e})", + expected_gpm[k_idx], + err + ); + } + eprintln!( + "PASS — gate_probs_mean correct under uniform + asymmetric inputs (max_abs_err = {max_err:.2e})" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn aggregate_diag_entropy_pins_top_and_bottom() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 16usize; + let k = 3usize; + + let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) } + .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?; + for slot in 0..RL_SLOTS_END { + isv.write_record(slot, 0.0); + } + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0x2A2D_0002, + }, + )?; + + // ── Top: uniform gate (max entropy = log(K)) ───────────────────── + let uniform = 1.0_f32 / (k as f32); + let gate_uniform = vec![uniform; b_size * k]; + // Uniform per-head action probs (max per_head_entropy = log(N_ACTIONS)). + let uniform_a = 1.0_f32 / (N_ACTIONS as f32); + let pi_uniform = vec![uniform_a; b_size * k * N_ACTIONS]; + write_slice_f32_d_pub(&stream, &gate_uniform, &mut mhp.gate_probs_d)?; + write_slice_f32_d_pub(&stream, &pi_uniform, &mut mhp.pi_probs_k_d)?; + + mhp.emit_diag_stats(isv_dev_ptr)?; + stream.synchronize()?; + + let log_k = (k as f32).ln(); + let log_a = (N_ACTIONS as f32).ln(); + let ge_top = isv.read_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX); + assert!( + (ge_top - log_k).abs() < 1e-5, + "top: gate_entropy_mean = {ge_top}, expected log(K) = {log_k}" + ); + let phe_base = RL_POLICY_PER_HEAD_ENTROPY_MEAN_BASE_INDEX; + for k_idx in 0..k { + let got = isv.read_record(phe_base + k_idx); + assert!( + (got - log_a).abs() < 1e-5, + "top: per_head_entropy_mean[{k_idx}] = {got}, expected log(N_ACTIONS) = {log_a}" + ); + } + // Tail must be zero. + for k_idx in k..MAX_K_HEADS { + let got = isv.read_record(phe_base + k_idx); + assert_eq!( + got, 0.0, + "top tail: per_head_entropy_mean[{k_idx}] = {got}, expected 0.0" + ); + } + // gate_argmax_mass with strict-tie-broken-low: all rows uniform, so + // argmax is index 0 for every batch → mass[0]=1, mass[k>=1]=0. + let gam_base = RL_POLICY_GATE_ARGMAX_MASS_BASE_INDEX; + let mass0 = isv.read_record(gam_base); + assert!( + (mass0 - 1.0).abs() < 1e-6, + "top: gate_argmax_mass[0] = {mass0}, expected 1.0 (ties broken by lowest index)" + ); + for k_idx in 1..MAX_K_HEADS { + let got = isv.read_record(gam_base + k_idx); + assert_eq!( + got, 0.0, + "top: gate_argmax_mass[{k_idx}] = {got}, expected 0.0" + ); + } + + // ── Bottom: one-hot gate (min entropy = 0) ──────────────────────── + // Spread argmax mass across the K heads so we exercise the per-head + // counter scatter. + let mut gate_onehot = vec![0.0_f32; b_size * k]; + for b in 0..b_size { + let chosen = b % k; + gate_onehot[b * k + chosen] = 1.0; + } + // One-hot per-head action probs (zero entropy per head). + let mut pi_onehot = vec![0.0_f32; b_size * k * N_ACTIONS]; + for b in 0..b_size { + for k_idx in 0..k { + pi_onehot[(b * k + k_idx) * N_ACTIONS + (k_idx % N_ACTIONS)] = 1.0; + } + } + write_slice_f32_d_pub(&stream, &gate_onehot, &mut mhp.gate_probs_d)?; + write_slice_f32_d_pub(&stream, &pi_onehot, &mut mhp.pi_probs_k_d)?; + + // Reset target slots so we read fresh writes. + for slot in gam_base..gam_base + MAX_K_HEADS { + isv.write_record(slot, 0.0); + } + for slot in phe_base..phe_base + MAX_K_HEADS { + isv.write_record(slot, 0.0); + } + isv.write_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX, 0.0); + + mhp.emit_diag_stats(isv_dev_ptr)?; + stream.synchronize()?; + + let ge_bot = isv.read_record(RL_POLICY_GATE_ENTROPY_MEAN_INDEX); + assert!( + ge_bot.abs() < 1e-6, + "bottom: gate_entropy_mean = {ge_bot}, expected ≈0" + ); + for k_idx in 0..k { + let got = isv.read_record(phe_base + k_idx); + assert!( + got.abs() < 1e-6, + "bottom: per_head_entropy_mean[{k_idx}] = {got}, expected ≈0" + ); + } + // argmax_mass should sum to 1 across the K heads (every batch gets + // one argmax). With chosen = b % k and b_size = 16, k = 3: + // head 0: b ∈ {0,3,6,9,12,15} → 6/16 + // head 1: b ∈ {1,4,7,10,13} → 5/16 + // head 2: b ∈ {2,5,8,11,14} → 5/16 + let mut expected_mass = vec![0.0_f32; k]; + for b in 0..b_size { + expected_mass[b % k] += 1.0 / (b_size as f32); + } + for k_idx in 0..k { + let got = isv.read_record(gam_base + k_idx); + assert!( + (got - expected_mass[k_idx]).abs() < 1e-6, + "bottom: gate_argmax_mass[{k_idx}] = {got}, expected {}", + expected_mass[k_idx] + ); + } + for k_idx in k..MAX_K_HEADS { + let got = isv.read_record(gam_base + k_idx); + assert_eq!( + got, 0.0, + "bottom tail: gate_argmax_mass[{k_idx}] = {got}, expected 0.0" + ); + } + eprintln!( + "PASS — entropy pins: top gate_entropy_mean = log(K) = {log_k:.4}, \ + bottom = 0; per_head_entropy top = log(N_ACTIONS) = {log_a:.4}, bottom = 0; \ + argmax_mass sums to 1.0" + ); + Ok(()) +}