feat(ml-alpha): Phase 2A-C+ — device-aggregated gate diag
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
200
crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu
Normal file
200
crates/ml-alpha/cuda/multi_head_policy_aggregate_diag.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
#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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<CudaModule>,
|
||||
reduce_axis0_fn: CudaFunction,
|
||||
_aggregate_diag_module: Arc<CudaModule>,
|
||||
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:
|
||||
|
||||
@@ -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<f32> =
|
||||
(0..K_STRIDE).map(|k| isv[gpm_base + k]).collect();
|
||||
let gate_argmax_mass: Vec<f32> =
|
||||
(0..K_STRIDE).map(|k| isv[gam_base + k]).collect();
|
||||
let per_head_entropy_mean: Vec<f32> =
|
||||
(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,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user