From 1e70cd5e591ae4285c9b185aaf8e3f8906f9be6d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 8 May 2026 22:45:40 +0200 Subject: [PATCH] feat(sp17-3.1): A_var_ema per-branch producer + HEALTH_DIAG emit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 of SP17 dueling-Q identifiability — first of three diagnostic producers landed atomically with kernel + launcher + Rust wrapper + HEALTH_DIAG emit + GPU oracle test per `feedback_wire_everything_up`. Per branch d ∈ {dir, mag, ord, urg}: Var_d = (1/(B × n_d × NA)) Σ_{i, a, z} (A[i, a, z] − mean_a A[*, z])² Block tree-reduce (no atomicAdd, `feedback_no_atomicadd`); 4 blocks × 256 threads. Pearl-A first-observation bootstrap (sentinel 0.0 → REPLACE on first launch); steady-state α = WELFORD_ALPHA_MIN=0.4 per `pearl_wiener_alpha_floor_for_nonstationary` — the structural-control floor preserves catch-up bandwidth without storing 24 Welford accumulator slots for a cold-path-cadence diagnostic. Cold-path emit: single launch per HEALTH_DIAG cadence (epoch boundary) right after `v_a_means`. New line: HEALTH_DIAG[N]: dueling [a_var=(d=X m=Y o=Z u=W)] The line will be extended with V_share + advantage_clip_bound in Phase 3.2, then finalised in Phase 3.3. GPU oracle test on RTX 3050 Ti: synthetic A constructed so each branch d has a closed-form Var(A_centered); kernel readback matches expected value within ε=1e-4 (f32 rounding budget for ~8×n×51 accumulator length). Plan: docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md Phase 3.1. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 16 ++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 139 +++++++++++ .../cuda_pipeline/sp17_a_var_ema_kernel.cu | 217 ++++++++++++++++++ .../src/trainers/dqn/trainer/training_loop.rs | 34 +++ crates/ml/tests/sp17_dueling_oracle_tests.rs | 183 +++++++++++++++ docs/dqn-wire-up-audit.md | 60 +++++ 6 files changed, 649 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/sp17_a_var_ema_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 6cbe25bb1..a703d8fb1 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -925,6 +925,22 @@ fn main() { // the diagnostic baseline for the SP17 Phase 0 kill criterion (V // dominance pre-centering). "v_a_means_diag_kernel.cu", + // SP17 Phase 3.1 (2026-05-08): per-branch EMA of the over-actions + // variance of `A_centered`. 4 blocks × 256 threads (one block per + // branch — dir/mag/ord/urg). Reads the same BRANCH-MAJOR + // `on_b_logits_buf` that `compute_expected_q` consumes; computes + // per-atom mean of A across actions in the branch; sums squared + // deviations from the per-atom mean across (B × n_branch × NA); + // writes to ISV[A_VAR_EMA_DIR/MAG/ORD/URG_INDEX] with Pearl-A + // first-observation bootstrap (sentinel SENTINEL_A_VAR_EMA=0.0) + // + α=WELFORD_ALPHA_MIN=0.4 EMA blend per + // `pearl_wiener_alpha_floor_for_nonstationary`. Block tree-reduce + // (no atomicAdd) per `feedback_no_atomicadd.md`. Diagnostic + // observability — does NOT modify Q-head logits or any consumer + // path. Consumed by HEALTH_DIAG `dueling [a_var=...]` line via + // `read_a_var_ema_per_branch()` in gpu_dqn_trainer.rs. See + // docs/dqn-wire-up-audit.md § "SP17 Phase 3.1" for rationale. + "sp17_a_var_ema_kernel.cu", // SP14 Layer C Phase C.1 (2026-05-08): cubin entries for // alpha_grad_compute_kernel.cu, gradient_hack_detect_kernel.cu, // and sp14_scale_wire_col_kernel.cu deleted atomically with diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 44031d16c..d024197bf 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -674,6 +674,18 @@ pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] = pub(crate) static SP17_V_A_MEANS_DIAG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/v_a_means_diag_kernel.cubin")); +/// SP17 Phase 3.1 (2026-05-08): per-branch EMA of the over-actions +/// variance of `A_centered`. 4 blocks × 256 threads (one block per +/// branch). Cold-path (epoch boundary) producer for the HEALTH_DIAG +/// `dueling [a_var=...]` line; signals dueling-Q identifiability health +/// (Var → 0 = `all actions equivalent, V dominates` = pre-SP17 pathology +/// re-emerging). Block tree-reduce per `feedback_no_atomicadd`. Pearl-A +/// first-observation bootstrap on slots 474..478. Loaded from +/// `sp17_a_var_ema_kernel.cubin`. Consumed by +/// `read_a_var_ema_per_branch()`. +pub(crate) static SP17_A_VAR_EMA_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/sp17_a_var_ema_kernel.cubin")); + // SP14 Layer C Phase C.1 (2026-05-08): α-machinery cubin statics deleted // atomically with the alpha_grad_compute / gradient_hack_detect / // sp14_scale_wire_col kernel files. SP14_ALPHA_GRAD_CUBIN, @@ -6712,6 +6724,15 @@ pub struct GpuDqnTrainer { /// host reads via `read_all()` after a stream sync. Per /// `feedback_no_htod_htoh_only_mapped_pinned`. sp17_v_a_means_diag_buf: super::mapped_pinned::MappedF32Buffer, + // ── SP17 Phase 3.1 (2026-05-08): per-branch A_var EMA producer ──────── + /// 4-block × 256-thread kernel computing per-branch + /// `Var(A_centered)` over (B × n_branch × NA) and writing the EMA + /// into ISV[A_VAR_EMA_DIR/MAG/ORD/URG_INDEX] (slots 474..478). + /// Pearl-A first-observation bootstrap on sentinel 0.0; α=0.4 + /// post-bootstrap per `pearl_wiener_alpha_floor_for_nonstationary`. + /// Loaded from `sp17_a_var_ema_kernel.cubin`. Consumed by + /// `read_a_var_ema_per_branch()`. + sp17_a_var_ema_kernel: CudaFunction, // SP14 Layer C Phase C.1 (2026-05-08): α-machinery struct fields deleted // atomically with the kernel files (alpha_grad_compute_kernel, // gradient_hack_detect_kernel, sp14_scale_wire_col_kernel) per @@ -9009,6 +9030,112 @@ impl GpuDqnTrainer { Ok([host[0], host[1], host[2], host[3], host[4]]) } + /// SP17 Phase 3.1 (2026-05-08) — per-branch A_var EMA producer launch. + /// + /// Cold-path (epoch boundary) launcher for `sp17_a_var_ema_kernel`. + /// Reads the same online advantage buffer (`on_b_logits_buf`) that + /// `compute_expected_q` consumes; computes per-branch + /// `Var(A_centered)` over (B × n_branch × NA) where + /// `A_centered[i, a, z] = A[i, a, z] − (1/n_d) Σ_a' A[i, a', z]` + /// and EMA-blends into ISV[A_VAR_EMA_*_INDEX] with Pearl-A + /// first-observation bootstrap (sentinel + /// `SENTINEL_A_VAR_EMA=0.0`) + α=`WELFORD_ALPHA_MIN=0.4` per + /// `pearl_wiener_alpha_floor_for_nonstationary`. + /// + /// 4 blocks × 256 threads (one block per branch — dir/mag/ord/urg). + /// Block tree-reduce per `feedback_no_atomicadd`. Returns Ok(()) on + /// success; HEALTH_DIAG reads the ISV slots via + /// `read_isv_signal_at(A_VAR_EMA_*_INDEX)` after a stream sync. + /// + /// Diagnostic only — does NOT modify any consumer path. Per + /// `feedback_no_htod_htoh_only_mapped_pinned` the ISV bus is already + /// mapped-pinned; this kernel writes via `dev_ptr + + /// __threadfence_system()` and the host reads after sync. + pub fn launch_a_var_ema_update(&self) -> Result<(), MLError> { + use super::sp14_isv_slots::{ + A_VAR_EMA_DIR_INDEX, A_VAR_EMA_MAG_INDEX, + A_VAR_EMA_ORD_INDEX, A_VAR_EMA_URG_INDEX, + SENTINEL_A_VAR_EMA, WELFORD_ALPHA_MIN, + }; + let b = self.config.batch_size; + let na = self.config.num_atoms; + let b0 = self.config.branch_0_size; + let b1 = self.config.branch_1_size; + let b2 = self.config.branch_2_size; + let b3 = self.config.branch_3_size; + if b == 0 || na == 0 || (b0 + b1 + b2 + b3) == 0 { + return Ok(()); + } + let adv_ptr = self.on_b_logits_buf.raw_ptr(); + let isv_ptr = self.isv_signals_dev_ptr; + let b_i32 = b as i32; + let na_i32 = na as i32; + let b0_i32 = b0 as i32; + let b1_i32 = b1 as i32; + let b2_i32 = b2 as i32; + let b3_i32 = b3 as i32; + let var_idx_dir = A_VAR_EMA_DIR_INDEX as i32; + let var_idx_mag = A_VAR_EMA_MAG_INDEX as i32; + let var_idx_ord = A_VAR_EMA_ORD_INDEX as i32; + let var_idx_urg = A_VAR_EMA_URG_INDEX as i32; + let alpha: f32 = WELFORD_ALPHA_MIN; + let sentinel: f32 = SENTINEL_A_VAR_EMA; + // 4 blocks × 256 threads. Dynamic shmem = (BLOCK_SIZE + NA) × 4 + // bytes for the tree-reduce tile + per-atom mean cache. + let block_dim: u32 = 256; + let shmem_bytes: u32 = + (block_dim + na as u32) * std::mem::size_of::() as u32; + unsafe { + self.stream + .launch_builder(&self.sp17_a_var_ema_kernel) + .arg(&adv_ptr) + .arg(&isv_ptr) + .arg(&b_i32) + .arg(&na_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) + .arg(&var_idx_dir) + .arg(&var_idx_mag) + .arg(&var_idx_ord) + .arg(&var_idx_urg) + .arg(&alpha) + .arg(&sentinel) + .launch(LaunchConfig { + grid_dim: (4, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: shmem_bytes, + }) + .map_err(|e| MLError::ModelError(format!( + "sp17_a_var_ema_update launch: {e}" + )))?; + } + Ok(()) + } + + /// SP17 Phase 3.1 (2026-05-08) — per-branch A_var EMA readback. + /// + /// Convenience wrapper: launches the producer, syncs, returns + /// `[a_var_dir, a_var_mag, a_var_ord, a_var_urg]` from the ISV bus. + /// Cold-path only (HEALTH_DIAG cadence). + pub fn read_a_var_ema_per_branch(&self) -> Result<[f32; 4], MLError> { + use super::sp14_isv_slots::{ + A_VAR_EMA_DIR_INDEX, A_VAR_EMA_MAG_INDEX, + A_VAR_EMA_ORD_INDEX, A_VAR_EMA_URG_INDEX, + }; + self.launch_a_var_ema_update()?; + self.stream + .synchronize() + .map_err(|e| MLError::ModelError(format!("sp17_a_var_ema sync: {e}")))?; + Ok([ + self.read_isv_signal_at(A_VAR_EMA_DIR_INDEX), + self.read_isv_signal_at(A_VAR_EMA_MAG_INDEX), + self.read_isv_signal_at(A_VAR_EMA_ORD_INDEX), + self.read_isv_signal_at(A_VAR_EMA_URG_INDEX), + ]) + } + /// Update per-branch liquid tau modulation — GPU kernel (RK4/Euler adaptive ODE). /// /// Reads per_branch_q_gaps (pinned device-mapped) and qlstm_context from device. @@ -20117,6 +20244,17 @@ impl GpuDqnTrainer { } .map_err(|e| MLError::ModelError(format!("sp17_v_a_means_diag_buf alloc (5 f32): {e}")))?; + // SP17 Phase 3.1 (2026-05-08): per-branch A_var EMA producer. + // Cold-path (epoch boundary) — one launch per HEALTH_DIAG emit. + // 4 blocks × 256 threads; block tree-reduce per + // `feedback_no_atomicadd`. Pearl-A bootstrap on slots 474..478. + let sp17_a_var_ema_module = stream.context() + .load_cubin(SP17_A_VAR_EMA_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("sp17 a_var_ema cubin: {e}")))?; + let sp17_a_var_ema_kernel = sp17_a_var_ema_module + .load_function("sp17_a_var_ema_update") + .map_err(|e| MLError::ModelError(format!("sp17_a_var_ema_update: {e}")))?; + // SP14 Layer C Phase C.1 (2026-05-08): α-machinery cubin loads // (alpha_grad_compute, gradient_hack_detect, sp14_scale_wire_col) // deleted atomically with the kernel files. Coupling A @@ -24355,6 +24493,7 @@ impl GpuDqnTrainer { // producer launched at HEALTH_DIAG emit only. sp17_v_a_means_diag_kernel, sp17_v_a_means_diag_buf, + sp17_a_var_ema_kernel, sp14_dir_concat_qaux_kernel, sp14_dir_qaux_concat_scratch, // SP14 backward dX scratch — column-SH2 wire is dropped at diff --git a/crates/ml/src/cuda_pipeline/sp17_a_var_ema_kernel.cu b/crates/ml/src/cuda_pipeline/sp17_a_var_ema_kernel.cu new file mode 100644 index 000000000..159b541a1 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/sp17_a_var_ema_kernel.cu @@ -0,0 +1,217 @@ +/* ══════════════════════════════════════════════════════════════════════════ + * SP17 Phase 3.1 (2026-05-08) — A_var_ema producer kernel. + * + * Per-branch EMA of the over-actions variance of the centered advantage + * A_centered[i, a, z] = A[i, a, z] − (1/n_d) Σ_a' A[i, a', z] + * where the per-atom mean is taken over the b_d direction-of-branch + * actions (b0_size for dir, b1_size for mag, ...). + * + * Writes the per-branch variance EMA to four ISV slots: + * ISV[A_VAR_EMA_DIR_INDEX=474] (branch 0 — dir) + * ISV[A_VAR_EMA_MAG_INDEX=475] (branch 1 — mag) + * ISV[A_VAR_EMA_ORD_INDEX=476] (branch 2 — ord) + * ISV[A_VAR_EMA_URG_INDEX=477] (branch 3 — urg) + * + * Var_d = (1 / (B × n_d × NA)) Σ_{i, a, z} A_centered[i, a, z]^2 + * (since E[A_centered]=0 by construction over actions, the second + * moment IS the variance up to f32 rounding.) + * + * Producer cadence: cold-path (epoch boundary) — single launch per + * HEALTH_DIAG emit. Mirrors the v_a_means_diag_kernel.cu cadence and + * dispatch pattern. + * + * ── Pearls + invariants ──────────────────────────────────────────────── + * - `feedback_no_atomicadd.md` — block tree-reduce only; no atomicAdd. + * - `pearl_first_observation_bootstrap.md` — sentinel = SENTINEL_A_VAR_EMA + * (=0.0); first valid observation REPLACES directly. After bootstrap, + * blend with caller-supplied α (the host derives a Wiener-α with floor + * WELFORD_ALPHA_MIN_F=0.4 per the SP16 T3 fix-up; for cold-path + * cadence the host passes the steady-state floor directly because the + * variance signal is stable enough that a Welford accumulator state + * machine inside the kernel would add 6 ISV slots × 4 branches = 24 + * slots for a cold-path cadence diagnostic — not worth the complexity. + * Instead, the floor IS the α used after Pearl-A bootstrap.). + * - `pearl_wiener_alpha_floor_for_nonstationary` — α is the structural + * floor 0.4 (NOT zero); preserves catch-up bandwidth when the policy + * drifts. + * - `pearl_no_host_branches_in_captured_graph.md` — kernel runs + * cold-path; no host branches inside. + * - `feedback_isv_for_adaptive_bounds.md` — α is the structural + * bandwidth floor (NOT a tuned constant — it's the + * `pearl_wiener_alpha_floor_for_nonstationary` MIN value), and the + * sentinel/bootstrap is the cold-start contract. + * + * Algorithm (per launch, four branches in parallel via gridDim.x = 4): + * + * Per branch d (block d, threads 256): + * ── Pass 1: per-atom mean of A across n_d actions ── + * For each atom z in [0, NA): + * m_a[z] = (1 / (B × n_d)) Σ_{i, a} A[i, a, z] + * + * ── Pass 2: variance of centered A ── + * Var_d = (1 / (B × n_d × NA)) Σ_{i, a, z} (A[i, a, z] - m_a[z])^2 + * + * ── Pass 3: Pearl-A bootstrap + α blend (thread 0 only) ── + * cur = isv[var_idx_d]; + * if (fabsf(cur - SENTINEL_A_VAR_EMA) < EPS_F) { + * next = Var_d; // first observation REPLACES + * } else { + * next = (1 - alpha) * cur + alpha * Var_d; + * } + * isv[var_idx_d] = next; + * __threadfence_system(); + * + * Launch contract: + * grid_dim = (4, 1, 1) — one block per branch + * block_dim = (SP17_VA_BLOCK_SIZE, 1, 1) — 256 threads + * shared mem (dynamic): (BLOCK_SIZE + NA) × sizeof(float) + * = (256 + NA) × 4 bytes + * = 1024 bytes for the tree-reduce tile + * + NA × 4 bytes for the per-atom mean cache. + * + * Args: + * b_logits — Online branch logits, BRANCH-MAJOR layout + * [B×b0×NA | B×b1×NA | B×b2×NA | B×b3×NA]. Same buffer + * `compute_expected_q` reads from. f32 device ptr. + * isv — [ISV_TOTAL_DIM] f32 device ptr (mapped-pinned bus). + * B — batch size. + * NA — C51 atoms per action. + * b0_size, b1_size, b2_size, b3_size — per-branch action counts. + * var_idx_dir, var_idx_mag, var_idx_ord, var_idx_urg — + * ISV slot indices for each branch (474..478). + * alpha — EMA blend rate post-bootstrap. Caller passes + * WELFORD_ALPHA_MIN (=0.4) per + * `pearl_wiener_alpha_floor_for_nonstationary`. + * sentinel_a_var — SENTINEL_A_VAR_EMA = 0.0 (cold-start sentinel). + * ══════════════════════════════════════════════════════════════════════════ */ + +#include + +#define SP17_VA_BLOCK_SIZE 256 +#define SP17_EPS_F 1e-6f + +extern "C" __global__ void sp17_a_var_ema_update( + const float* __restrict__ b_logits, + float* __restrict__ isv, + int B, + int NA, + int b0_size, + int b1_size, + int b2_size, + int b3_size, + int var_idx_dir, + int var_idx_mag, + int var_idx_ord, + int var_idx_urg, + float alpha, + float sentinel_a_var) +{ + /* One block per branch (gridDim.x == 4). */ + const int branch = (int)blockIdx.x; + if (branch >= 4) return; + + /* Branch dispatch: select per-branch action count + ISV slot index + + * BRANCH-MAJOR offset (in float elements) into b_logits. */ + int n_branch; + int var_idx; + long long off_branch; + const long long Bll = (long long)B; + const long long NAll = (long long)NA; + const long long n0 = (long long)b0_size; + const long long n1 = (long long)b1_size; + const long long n2 = (long long)b2_size; + const long long off0 = 0; + const long long off1 = off0 + Bll * n0 * NAll; + const long long off2 = off1 + Bll * n1 * NAll; + const long long off3 = off2 + Bll * (long long)b2_size * NAll; + if (branch == 0) { n_branch = b0_size; var_idx = var_idx_dir; off_branch = off0; } + else if (branch == 1) { n_branch = b1_size; var_idx = var_idx_mag; off_branch = off1; } + else if (branch == 2) { n_branch = b2_size; var_idx = var_idx_ord; off_branch = off2; } + else { n_branch = b3_size; var_idx = var_idx_urg; off_branch = off3; } + + const int tid = (int)threadIdx.x; + const int bdim = (int)blockDim.x; + + /* Degenerate guard: zero-sized branch or empty buffers — skip the + * EMA update entirely. Leaves the ISV slot at its current value + * (sentinel on first launch, prior EMA on subsequent launches). */ + if (B <= 0 || NA <= 0 || n_branch <= 0) { + return; + } + + /* Shared memory layout (dynamic): + * shmem[0..bdim) — tree-reduce tile (per pass) + * shmem[bdim..bdim+NA) — per-atom mean a_mean[z] cached + * across pass 1 → pass 2. + */ + extern __shared__ float shmem[]; + float* s_reduce = shmem; + float* s_a_mean = shmem + bdim; + + const long long total = Bll * (long long)n_branch * NAll; + + /* ── Pass 1: per-atom mean over (B × n_branch) for each atom z ───── + * Loop over atoms; per-atom block-wide tree reduction. Cost ≤ 51 + * reductions per branch × 4 branches; cold-path cadence makes it + * irrelevant. + */ + for (int z = 0; z < NA; ++z) { + float local = 0.0f; + const long long count_ia = Bll * (long long)n_branch; + /* Linearise (i, a) → idx, read b_logits[off_branch + idx*NA + z]. */ + for (long long idx = (long long)tid; idx < count_ia; idx += (long long)bdim) { + const long long off = off_branch + idx * NAll + (long long)z; + local += b_logits[off]; + } + s_reduce[tid] = local; + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) s_reduce[tid] += s_reduce[tid + s]; + __syncthreads(); + } + if (tid == 0) { + const float total_count = (float)count_ia; + s_a_mean[z] = (total_count > 0.0f) ? (s_reduce[0] / total_count) : 0.0f; + } + __syncthreads(); + } + + /* ── Pass 2: variance over (B × n_branch × NA) of (A − a_mean[z]) ─ + * Single strided pass with per-thread sum-of-squared-deviations. */ + float local_sumsq = 0.0f; + for (long long idx = (long long)tid; idx < total; idx += (long long)bdim) { + const long long ia = idx / NAll; + const long long z = idx - ia * NAll; + const long long off = off_branch + ia * NAll + z; + const float a = b_logits[off]; + const float dev = a - s_a_mean[z]; + local_sumsq += dev * dev; + } + s_reduce[tid] = local_sumsq; + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) s_reduce[tid] += s_reduce[tid + s]; + __syncthreads(); + } + + /* ── Pass 3: Pearl-A bootstrap + α blend → ISV write ──────────── */ + if (tid == 0) { + const float total_f = (float)total; + const float var_d = (total_f > 0.0f) ? (s_reduce[0] / total_f) : 0.0f; + + const float current = isv[var_idx]; + float next; + if (fabsf(current - sentinel_a_var) < SP17_EPS_F) { + /* Pearl-A first-observation bootstrap: sentinel detected → + * REPLACE directly (no blend). */ + next = var_d; + } else { + /* Steady-state: blend with caller-supplied α (host passes + * WELFORD_ALPHA_MIN=0.4 per + * `pearl_wiener_alpha_floor_for_nonstationary`). */ + next = (1.0f - alpha) * current + alpha * var_d; + } + isv[var_idx] = next; + __threadfence_system(); /* PCIe-visible write for mapped-pinned ISV */ + } +} diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 345b0c122..b56edbb1d 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -5273,6 +5273,40 @@ impl DQNTrainer { ); } + // SP17 Phase 3.1 (2026-05-08) — per-branch A_var EMA diagnostic. + // Cold-path (epoch boundary) producer launch; emits the + // dueling-Q identifiability variance signal so the post-deploy + // reviewer can confirm `A_centered` is doing real work + // (Var > 0 = the policy discriminates between actions) vs. + // regression to the pre-SP17 V-dominated regime (Var → 0 = "all + // actions equivalent, V dominates"). Per `feedback_no_atomicadd` + // the kernel uses 4-block tree-reduce; per + // `pearl_first_observation_bootstrap` slots 474..478 use + // sentinel SENTINEL_A_VAR_EMA=0.0 with REPLACE on first + // observation. α post-bootstrap is WELFORD_ALPHA_MIN=0.4 per + // `pearl_wiener_alpha_floor_for_nonstationary`. Plan: + // docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md + // Phase 3.1. + { + let a_var = if let Some(ref fused) = self.fused_ctx { + match fused.trainer().read_a_var_ema_per_branch() { + Ok(v) => v, + Err(e) => { + tracing::warn!( + "HEALTH_DIAG read_a_var_ema_per_branch refresh failed: {e}" + ); + [0.0_f32; 4] + } + } + } else { + [0.0_f32; 4] + }; + tracing::info!( + "HEALTH_DIAG[{}]: dueling [a_var=(d={:.4} m={:.4} o={:.4} u={:.4})]", + epoch, a_var[0], a_var[1], a_var[2], a_var[3], + ); + } + // SP7 observability (grad_decomp_pinned): surface the exact contents of // grad_decomp_result_pinned[0..3, 12..15, 36..48] (iqn/cql_raw/c51 × // [mag, dir, trunk]) that the SP7 loss-balance controller read in the diff --git a/crates/ml/tests/sp17_dueling_oracle_tests.rs b/crates/ml/tests/sp17_dueling_oracle_tests.rs index 6448e5b1f..75320697a 100644 --- a/crates/ml/tests/sp17_dueling_oracle_tests.rs +++ b/crates/ml/tests/sp17_dueling_oracle_tests.rs @@ -2113,4 +2113,187 @@ mod gpu { counts[DIR_LONG], counts, ); } + + // ── F.1: A_var_ema_per_branch_tracks_centered_variance ───────────── + /// Phase 3.1 (2026-05-08) — synthetic A_centered with known per-branch + /// variances. Launch `sp17_a_var_ema_update` against a 4-branch + /// `b_logits` buffer constructed so each branch d has a distinct + /// closed-form Var(A_centered) over actions. + /// + /// Construction (per branch d, n_d actions, NA atoms): + /// - For each (i, a, z): A[i, a, z] = K_d × (a - (n_d - 1)/2), + /// i.e. linearly spaced about zero per atom z. + /// - The per-atom mean over actions is then 0 (A is mean-zero + /// already), so A_centered[i, a, z] = A[i, a, z]. + /// - Var_d = (1/(B × n_d × NA)) Σ K_d^2 × (a - (n_d-1)/2)^2 + /// = K_d^2 × (1/n_d) Σ_a (a - (n_d-1)/2)^2 + /// = K_d^2 × var_a(a) (constant across atom z) + /// - For n_d in {4, 3, 3, 3}, var_a(a) = {1.25, 0.667, 0.667, 0.667}. + /// - With K_d = {1.0, 2.0, 3.0, 0.5}, Var_d = + /// {1.25, 2.667, 6.0, 0.1667}. + /// + /// Pearl-A bootstrap: the ISV slot starts at sentinel 0.0 → first + /// observation REPLACES directly, so the readback equals Var_d + /// exactly (within f32 rounding) on the very first launch. + #[test] + #[ignore = "requires GPU"] + fn a_var_ema_per_branch_tracks_centered_variance() { + use ml::cuda_pipeline::sp14_isv_slots::{ + A_VAR_EMA_DIR_INDEX, A_VAR_EMA_MAG_INDEX, + A_VAR_EMA_ORD_INDEX, A_VAR_EMA_URG_INDEX, + SENTINEL_A_VAR_EMA, WELFORD_ALPHA_MIN, + }; + + const SP17_A_VAR_EMA_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/sp17_a_var_ema_kernel.cubin")); + + let stream = make_test_stream(); + let module = stream + .context() + .load_cubin(SP17_A_VAR_EMA_CUBIN.to_vec()) + .expect("load sp17_a_var_ema_kernel cubin"); + let kernel = module + .load_function("sp17_a_var_ema_update") + .expect("load sp17_a_var_ema_update function"); + + // Production-realistic shapes: B=8, NA=51, branches (4, 3, 3, 3). + const B: usize = 8; + const NA: usize = 51; + const N0: usize = 4; // dir + const N1: usize = 3; // mag + const N2: usize = 3; // ord + const N3: usize = 3; // urg + const K_DIR: f32 = 1.0; + const K_MAG: f32 = 2.0; + const K_ORD: f32 = 3.0; + const K_URG: f32 = 0.5; + + // Build BRANCH-MAJOR b_logits with the linear pattern. + let n_per_branch = [N0, N1, N2, N3]; + let k_per_branch = [K_DIR, K_MAG, K_ORD, K_URG]; + let dir_len = B * N0 * NA; + let mag_len = B * N1 * NA; + let ord_len = B * N2 * NA; + let urg_len = B * N3 * NA; + let mut b_logits = Vec::with_capacity(dir_len + mag_len + ord_len + urg_len); + for branch in 0..4 { + let n = n_per_branch[branch]; + let k = k_per_branch[branch]; + let center = (n as f32 - 1.0) / 2.0; + for _i in 0..B { + for a in 0..n { + for _z in 0..NA { + let val = k * (a as f32 - center); + b_logits.push(val); + } + } + } + } + assert_eq!(b_logits.len(), dir_len + mag_len + ord_len + urg_len); + + let b_buf = unsafe { MappedF32Buffer::new(b_logits.len()) } + .expect("alloc b_logits buf"); + b_buf.write_from_slice(&b_logits); + + // ISV bus: pre-populate slots 474..478 with sentinel 0.0 so + // Pearl-A bootstrap fires (first observation REPLACES directly). + // Allocate the full ISV span so device pointer arithmetic is safe. + const ISV_SPAN: usize = 483; // matches ISV_TOTAL_DIM + let mut isv_host = vec![0.0_f32; ISV_SPAN]; + isv_host[A_VAR_EMA_DIR_INDEX] = SENTINEL_A_VAR_EMA; + isv_host[A_VAR_EMA_MAG_INDEX] = SENTINEL_A_VAR_EMA; + isv_host[A_VAR_EMA_ORD_INDEX] = SENTINEL_A_VAR_EMA; + isv_host[A_VAR_EMA_URG_INDEX] = SENTINEL_A_VAR_EMA; + let isv_buf = unsafe { MappedF32Buffer::new(ISV_SPAN) } + .expect("alloc isv buf"); + isv_buf.write_from_slice(&isv_host); + + // Launch. 4 blocks × 256 threads. Dynamic shmem = (256 + NA) × 4 bytes. + let b_i32: i32 = B as i32; + let na_i32: i32 = NA as i32; + let n0_i32: i32 = N0 as i32; + let n1_i32: i32 = N1 as i32; + let n2_i32: i32 = N2 as i32; + let n3_i32: i32 = N3 as i32; + let var_idx_dir: i32 = A_VAR_EMA_DIR_INDEX as i32; + let var_idx_mag: i32 = A_VAR_EMA_MAG_INDEX as i32; + let var_idx_ord: i32 = A_VAR_EMA_ORD_INDEX as i32; + let var_idx_urg: i32 = A_VAR_EMA_URG_INDEX as i32; + let alpha: f32 = WELFORD_ALPHA_MIN; + let sentinel: f32 = SENTINEL_A_VAR_EMA; + let block_dim: u32 = 256; + let shmem_bytes: u32 = + (block_dim + NA as u32) * std::mem::size_of::() as u32; + + unsafe { + stream + .launch_builder(&kernel) + .arg(&b_buf.dev_ptr) + .arg(&isv_buf.dev_ptr) + .arg(&b_i32) + .arg(&na_i32) + .arg(&n0_i32) + .arg(&n1_i32) + .arg(&n2_i32) + .arg(&n3_i32) + .arg(&var_idx_dir) + .arg(&var_idx_mag) + .arg(&var_idx_ord) + .arg(&var_idx_urg) + .arg(&alpha) + .arg(&sentinel) + .launch(LaunchConfig { + grid_dim: (4, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: shmem_bytes, + }) + .expect("launch sp17_a_var_ema_update"); + } + stream.synchronize().expect("sync after sp17_a_var_ema_update"); + + let isv_after = isv_buf.read_all(); + + // Closed-form expected variances per branch. + // var_a(a − center) = (1/n) Σ_a (a − center)^2. + fn expected_var(n: usize, k: f32) -> f32 { + let center = (n as f32 - 1.0) / 2.0; + let mut s = 0.0_f32; + for a in 0..n { + let dev = a as f32 - center; + s += dev * dev; + } + (s / n as f32) * k * k + } + let exp_dir = expected_var(N0, K_DIR); + let exp_mag = expected_var(N1, K_MAG); + let exp_ord = expected_var(N2, K_ORD); + let exp_urg = expected_var(N3, K_URG); + + let got_dir = isv_after[A_VAR_EMA_DIR_INDEX]; + let got_mag = isv_after[A_VAR_EMA_MAG_INDEX]; + let got_ord = isv_after[A_VAR_EMA_ORD_INDEX]; + let got_urg = isv_after[A_VAR_EMA_URG_INDEX]; + + let eps = 1e-4_f32; // f32 rounding budget for ~B*n*NA accumulation + assert!( + (got_dir - exp_dir).abs() < eps, + "dir Var(A_centered) expected {exp_dir:.6}, got {got_dir:.6} \ + (diff {:.3e})", (got_dir - exp_dir).abs(), + ); + assert!( + (got_mag - exp_mag).abs() < eps, + "mag Var(A_centered) expected {exp_mag:.6}, got {got_mag:.6} \ + (diff {:.3e})", (got_mag - exp_mag).abs(), + ); + assert!( + (got_ord - exp_ord).abs() < eps, + "ord Var(A_centered) expected {exp_ord:.6}, got {got_ord:.6} \ + (diff {:.3e})", (got_ord - exp_ord).abs(), + ); + assert!( + (got_urg - exp_urg).abs() < eps, + "urg Var(A_centered) expected {exp_urg:.6}, got {got_urg:.6} \ + (diff {:.3e})", (got_urg - exp_urg).abs(), + ); + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 40cf23ed6..8cb54a268 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,66 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-08 — SP17 Phase 3.1: A_var_ema per-branch producer kernel (additive observability) + +Phase 3 of the SP17 dueling-Q identifiability plan introduces the diagnostic chain that lets the post-deploy reviewer confirm the mean-zero centering (Phase 1-2) is doing real work. This commit lands the first of three producers: per-branch EMA of the over-actions variance of `A_centered`. + +**New file:** `crates/ml/src/cuda_pipeline/sp17_a_var_ema_kernel.cu` + +**Kernel:** `sp17_a_var_ema_update` +- Grid: `(4, 1, 1)` — one block per branch (dir/mag/ord/urg). +- Block: `(256, 1, 1)`. +- Dynamic shmem: `(BLOCK_SIZE + NA) × sizeof(f32)` — tree-reduce tile + per-atom mean cache. +- Reads: `on_b_logits_buf` (BRANCH-MAJOR `[B×b0×NA | B×b1×NA | B×b2×NA | B×b3×NA]`). +- Writes: `ISV[A_VAR_EMA_DIR/MAG/ORD/URG_INDEX]` (slots 474..478). + +**Algorithm (per branch d):** +1. Pass 1 — per-atom mean over (B × n_branch) actions: + `m_a[z] = (1 / (B × n_branch)) Σ_{i, a} A[i, a, z]` +2. Pass 2 — variance over (B × n_branch × NA) of `(A − m_a[z])`: + `Var_d = (1 / (B × n_branch × NA)) Σ_{i, a, z} (A[i, a, z] − m_a[z])²` +3. Pass 3 — Pearl-A bootstrap + α blend → ISV write. + +**Pearls + invariants:** +- `feedback_no_atomicadd` — block tree-reduce only (no atomicAdd anywhere). +- `pearl_first_observation_bootstrap` — sentinel `SENTINEL_A_VAR_EMA=0.0`; first observation REPLACES directly. Cold-start guard at the first launch of every fold (FoldReset registry resets the slot to 0.0). +- `pearl_wiener_alpha_floor_for_nonstationary` — α post-bootstrap is the structural floor `WELFORD_ALPHA_MIN=0.4` (NOT a tuned value — it's the minimum bandwidth needed to track a co-adapting policy without lagging). Welford accumulator state would add 6 ISV slots × 4 branches = 24 slots for a cold-path-cadence diagnostic; instead the floor IS the steady-state α. The cold-start path uses Pearl-A REPLACE so first-observation responsiveness is preserved. +- `pearl_no_host_branches_in_captured_graph` — pure on-device control flow (branch dispatch via `blockIdx.x`). +- `feedback_isv_for_adaptive_bounds` — α is the structural bandwidth floor, sentinel/bootstrap is the cold-start contract; no hardcoded magic constants. +- `pearl_symmetric_clamp_audit` — N/A (variance is structurally non-negative; no clamp needed). + +**Rust integration (`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`):** +- `SP17_A_VAR_EMA_CUBIN` static (cubin handle). +- `sp17_a_var_ema_kernel: CudaFunction` field on `GpuDqnTrainer` (loaded once at construct). +- `launch_a_var_ema_update()` cold-path launcher (wrapped to skip empty-batch case). +- `read_a_var_ema_per_branch() -> [f32; 4]` convenience wrapper: launches → syncs → reads ISV slots. + +**HEALTH_DIAG emit (`crates/ml/src/trainers/dqn/trainer/training_loop.rs`):** + +``` +HEALTH_DIAG[N]: dueling [a_var=(d=X m=Y o=Z u=W)] +``` + +Emitted once per epoch, immediately after the existing `v_a_means` HEALTH_DIAG line. The line will be extended in Phase 3.2 (V_share + advantage_clip_bound producers) and finalised in Phase 3.3. + +**State reset:** ISV slots 474..478 use `FoldReset` category (already wired in PP.2 commit `a225926e5`). At fold boundary the slot is REPLACED with `SENTINEL_A_VAR_EMA=0.0` so Pearl-A bootstrap fires fresh on the new fold's first epoch. + +**GPU oracle test:** `crates/ml/tests/sp17_dueling_oracle_tests.rs::a_var_ema_per_branch_tracks_centered_variance` + +Synthetic A constructed so each branch d has a closed-form variance: +`A[i, a, z] = K_d × (a − (n_d-1)/2)` (linearly spaced, mean-zero per atom). +Expected `Var_d = K_d² × var_a(a)`. With `n_d ∈ {4, 3, 3, 3}` and `K_d ∈ {1.0, 2.0, 3.0, 0.5}`: +- dir = 1.0² × 1.25 = **1.25** +- mag = 2.0² × 0.667 = **2.667** +- ord = 3.0² × 0.667 = **6.0** +- urg = 0.5² × 0.667 = **0.1667** + +Tolerance ε = 1e-4 (f32 rounding for ≈ 8 × n × 51 accumulator length). Test runs in ~1.7s on RTX 3050 Ti — passes on first launch via Pearl-A bootstrap (no warmup needed). + +**Atomicity envelope:** Phase 3.1 is observability-only — no consumer path is modified. The Phase 1 mean-zero contract (commits `eabcf8d52` through `6f53d676f`) already lands the producer-side semantics; this commit observes them. V_share + advantage_clip_bound producers land in Phase 3.2; the final canonical HEALTH_DIAG line format is consolidated in Phase 3.3. + +**Plan:** `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Phase 3.1. + ## 2026-05-08 — SP17 Commit E: c51_loss_batched + c51_grad_kernel SP17-compliance annotations (audit-only, no code-path change) User design call DD9: audit `c51_loss_batched` + `c51_grad_kernel` for SP17-contract compliance and annotate the existing already-centered code with explicit SP17 markers. The plan's Task 1.4 incorrectly claimed these kernels needed migration — they have been mean-zero centered since commit `56373f094`. Verifying this audit closes the contract gap that the plan opened by undercounting consumers.