diff --git a/crates/ml/build.rs b/crates/ml/build.rs index ffe264541..3a75d54ce 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -421,6 +421,15 @@ fn main() { // to producer_step_scratch_buf[103..111). // apply_pearls_ad_kernel then smooths via Pearls A+D into ISV[210..218). "pearl_3_sigma_kernel.cu", + // SP5 Task A3 (2026-05-01): Pearl 2 per-branch loss budget controller. + // Single-block 4-thread kernel reads Q_VAR_PER_BRANCH (Pearl 1's + // q_branch_stats, ISV[222..226)) + NOISY_SIGMA (Pearl 3 output, + // ISV[210..214)); computes per-branch flatness = clamp(var_q/(σ²+EPS_DIV), + // 0, 1) and derives (budget_c51, budget_iqn, budget_cql=0, budget_ens, + // flatness) per branch. Writes 20 floats to scratch[111..131). + // apply_pearls_ad_kernel then smooths via Pearls A+D into ISV[190..210). + // Must run AFTER pearl_1_atom (Q_VAR) AND pearl_3_sigma (NOISY_SIGMA). + "pearl_2_budget_kernel.cu", // SP4 GPU-only Pearls A+D applicator (2026-05-01). Single-thread // device-side loop applies Pearls A+D to N consecutive ISV slots // (each with its own Wiener-state triple) on the producer's diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1b10539dc..06f38818f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -304,6 +304,20 @@ static SP5_PEARL_1_ATOM_CUBIN: &[u8] = static SP5_PEARL_3_SIGMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/pearl_3_sigma_kernel.cubin")); +/// SP5 Task A3 (2026-05-01): Pearl 2 per-branch loss budget cubin. +/// Single-block 4-thread kernel. Reads Q_VAR_PER_BRANCH[b] (Pearl 1's +/// q_branch_stats output, ISV[222..226)) + NOISY_SIGMA[b] (Pearl 3 output, +/// ISV[210..214)). Computes per-branch flatness = clamp(var_q / (σ²+EPS_DIV), +/// 0, 1) and derives (budget_iqn, budget_c51, budget_cql=0, budget_ens, +/// flatness) per branch. Writes 20 floats to +/// producer_step_scratch_buf[111..131). Chained apply_pearls_ad_kernel +/// (20 single-slot launches) smooths via Pearls A+D into ISV[190..210). +/// Must run AFTER launch_sp5_pearl_1_atom (Q_VAR) AND +/// launch_sp5_pearl_3_sigma (NOISY_SIGMA). +/// Layer A only — loss-budget consumer migration deferred to Layer B. +static SP5_PEARL_2_BUDGET_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/pearl_2_budget_kernel.cubin")); + /// Plan 4 Task 3 (E.3): IQN multi-quantile diagnostic EMAs into ISV[99..103). /// 4-block (256 threads/block, shmem-reduce, no atomicAdd) kernel launched /// alongside `h_s2_rms_ema_update` from `training_loop.rs`. Reads the IQN @@ -887,15 +901,22 @@ pub const SP5_WIENER_TOTAL_FLOATS: usize = (SP4_PRODUCER_COUNT + crate::cuda_pipeline::sp5_isv_slots::SP5_PRODUCER_COUNT) * SP4_WIENER_FLOATS_PER_SLOT; -/// SP5 Task A1+A2: producer_step_scratch_buf total slot count after growth. +/// SP5 Task A1+A2+A3: producer_step_scratch_buf total slot count after growth. /// SP4 had 71 scratch slots [0..71); Task A1 adds: /// 16 floats for q_branch_stats output (SCRATCH_Q_STATS=71, 4 outputs × 4 branches) /// 16 floats for pearl_1_atom output (SCRATCH_ATOM_OUT=87, 4 outputs × 4 branches) /// Task A2 adds: /// 4 floats for pearl_3_sigma output (SCRATCH_PEARL_3_SIGMA=103, sigma × 4 branches) /// 4 floats for pearl_3_sf output (SCRATCH_PEARL_3_SF=107, sigma_fraction × 4 branches) -/// Combined: 71 + 16 + 16 + 4 + 4 = 111 scratch slots [0..111). -pub const SP5_SCRATCH_TOTAL: usize = 111; +/// Task A3 adds: +/// 20 floats for pearl_2_budget output (5 outputs × 4 branches at [111..131)) +/// budget_c51[4] → SCRATCH_PEARL_2_C51=111 +/// budget_iqn[4] → SCRATCH_PEARL_2_IQN=115 +/// budget_cql[4] → SCRATCH_PEARL_2_CQL=119 +/// budget_ens[4] → SCRATCH_PEARL_2_ENS=123 +/// flatness[4] → SCRATCH_PEARL_2_FLATNESS=127 +/// Combined: 71 + 16 + 16 + 4 + 4 + 20 = 131 scratch slots [0..131). +pub const SP5_SCRATCH_TOTAL: usize = 131; /// SP5 Task A2: scratch index base for pearl_3_sigma sigma[4] output block. /// Slots [103..107): per-branch NoisyNet sigma. @@ -905,6 +926,27 @@ pub const SCRATCH_PEARL_3_SIGMA: usize = 103; /// Slots [107..111): per-branch sigma fraction (entropy-deficit controller state). pub const SCRATCH_PEARL_3_SF: usize = 107; +/// SP5 Task A3: scratch index base for pearl_2_budget budget_c51[4] output block. +/// Slots [111..115): per-branch C51 loss budget share. +pub const SCRATCH_PEARL_2_C51: usize = 111; + +/// SP5 Task A3: scratch index base for pearl_2_budget budget_iqn[4] output block. +/// Slots [115..119): per-branch IQN loss budget share. +pub const SCRATCH_PEARL_2_IQN: usize = 115; + +/// SP5 Task A3: scratch index base for pearl_2_budget budget_cql[4] output block. +/// Slots [119..123): per-branch CQL loss budget share (zero in Pearl 2; gated by +/// existing regime_stability allocator). +pub const SCRATCH_PEARL_2_CQL: usize = 119; + +/// SP5 Task A3: scratch index base for pearl_2_budget budget_ens[4] output block. +/// Slots [123..127): per-branch ensemble loss budget share (1 - IQN - C51 - CQL, ≥0). +pub const SCRATCH_PEARL_2_ENS: usize = 123; + +/// SP5 Task A3: scratch index base for pearl_2_budget flatness[4] output block. +/// Slots [127..131): per-branch flatness signal (var_q / (σ² + EPS_DIV), clamped [0,1]). +pub const SCRATCH_PEARL_2_FLATNESS: usize = 127; + /// SP5 Task A1: scratch index base for q_branch_stats 16-float output block. /// Slots [71..87): mean, max_abs_dev, var, entropy × 4 branches. /// SP4 slot 70 = q_dir_grad_p99; SP5 starts at 71 (no collision). @@ -4039,6 +4081,20 @@ pub struct GpuDqnTrainer { /// No atomicAdd (feedback_no_atomicadd), no CPU compute (feedback_no_cpu_compute_strict). /// Loaded from `pearl_3_sigma_kernel.cubin`. pearl_3_sigma_kernel: CudaFunction, + /// SP5 Task A3 (2026-05-01): Pearl 2 per-branch loss budget controller kernel. + /// Single-block 4-thread kernel (one thread per branch). Reads Q_VAR_PER_BRANCH + /// (Pearl 1's q_branch_stats, ISV[222..226)) + NOISY_SIGMA (Pearl 3 output, + /// ISV[210..214)); computes per-branch flatness = clamp(var_q / (σ²+EPS_DIV), 0, 1) + /// and derives (budget_c51, budget_iqn, budget_cql=0, budget_ens, flatness) per + /// branch. Writes 20 floats to `producer_step_scratch_buf[111..131)`. + /// Chained `apply_pearls_ad_kernel` (20 single-slot launches) then smooths via + /// Pearls A+D into ISV[190..210). + /// Must run AFTER launch_sp5_pearl_1_atom (Q_VAR) AND + /// launch_sp5_pearl_3_sigma (NOISY_SIGMA). Layer A only — loss-budget consumer + /// migration deferred to Layer B. + /// No atomicAdd (feedback_no_atomicadd), no CPU compute (feedback_no_cpu_compute_strict). + /// Loaded from `pearl_2_budget_kernel.cubin`. + pearl_2_budget_kernel: CudaFunction, /// SP5 Task A1 (2026-05-01): per-branch clip-count buffer for Pearl 1 /// headroom controller. [4 i32], one element per branch. Zero-initialized /// at construction; populated by `atoms_update_kernel` in Layer B. @@ -10608,6 +10664,106 @@ impl GpuDqnTrainer { Ok(()) } + /// SP5 Task A3 (2026-05-01): Pearl 2 — per-branch loss budget from per-branch + /// flatness signal. + /// + /// Chained AFTER `launch_sp5_pearl_1_atom` (reads Q_VAR_PER_BRANCH written by + /// Pearl 1's q_branch_stats) AND AFTER `launch_sp5_pearl_3_sigma` (reads + /// NOISY_SIGMA written by Pearl 3). Two-step sequence: + /// 1. `pearl_2_budget_update` (1×4 kernel) → scratch[111..131). + /// 2. `apply_pearls_ad_kernel` × 20 (one per output slot): + /// budget_c51[4] → ISV[BUDGET_C51_BASE=190..194) + /// budget_iqn[4] → ISV[BUDGET_IQN_BASE=194..198) + /// budget_cql[4] → ISV[BUDGET_CQL_BASE=198..202) + /// budget_ens[4] → ISV[BUDGET_ENS_BASE=202..206) + /// flatness[4] → ISV[FLATNESS_BASE=206..210) + /// + /// Wiener offsets: (SP4_PRODUCER_COUNT=71 + (isv_slot − SP5_SLOT_BASE=174)) × 3. + /// Layer A only — loss-budget consumer migration deferred to Layer B. + pub(crate) fn launch_sp5_pearl_2_budget(&self) -> Result<(), MLError> { + use crate::cuda_pipeline::sp4_wiener_ema::launch_apply_pearls; + use crate::cuda_pipeline::sp5_isv_slots::{ + SP5_SLOT_BASE, + BUDGET_C51_BASE, BUDGET_IQN_BASE, BUDGET_CQL_BASE, BUDGET_ENS_BASE, + FLATNESS_BASE, NOISY_SIGMA_BASE, Q_VAR_PER_BRANCH_BASE, + }; + + debug_assert!(self.isv_signals_dev_ptr != 0, + "launch_sp5_pearl_2_budget: isv_signals_dev_ptr must be allocated by constructor"); + + let isv_dev = self.isv_signals_dev_ptr; + let scratch_dev = self.producer_step_scratch_buf.dev_ptr; + let wiener_dev = self.wiener_state_buf.dev_ptr; + + // Step 1: pearl_2_budget_update → scratch[111..131). + // Kernel signature: (isv_signals, q_var_isv_base, sigma_isv_base, + // scratch_out, c51_idx_base, iqn_idx_base, + // cql_idx_base, ens_idx_base, flatness_idx_base) + let q_var_base_i32 = Q_VAR_PER_BRANCH_BASE as i32; // 222 + let sigma_base_i32 = NOISY_SIGMA_BASE as i32; // 210 + let c51_base_i32 = SCRATCH_PEARL_2_C51 as i32; // 111 + let iqn_base_i32 = SCRATCH_PEARL_2_IQN as i32; // 115 + let cql_base_i32 = SCRATCH_PEARL_2_CQL as i32; // 119 + let ens_base_i32 = SCRATCH_PEARL_2_ENS as i32; // 123 + let flatness_base_i32 = SCRATCH_PEARL_2_FLATNESS as i32; // 127 + + unsafe { + self.stream + .launch_builder(&self.pearl_2_budget_kernel) + .arg(&isv_dev) + .arg(&q_var_base_i32) + .arg(&sigma_base_i32) + .arg(&scratch_dev) + .arg(&c51_base_i32) + .arg(&iqn_base_i32) + .arg(&cql_base_i32) + .arg(&ens_base_i32) + .arg(&flatness_base_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (4, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("pearl_2_budget_update: {e}")))?; + } + + // Step 2: apply_pearls_ad_kernel × 20 — one single-slot launch per ISV output. + // Wiener offset: (SP4_PRODUCER_COUNT + (isv_slot - SP5_SLOT_BASE)) * 3. + // + // budget_c51[4] → scratch [111..115), ISV [190..194) (BUDGET_C51_BASE) + // budget_iqn[4] → scratch [115..119), ISV [194..198) (BUDGET_IQN_BASE) + // budget_cql[4] → scratch [119..123), ISV [198..202) (BUDGET_CQL_BASE) + // budget_ens[4] → scratch [123..127), ISV [202..206) (BUDGET_ENS_BASE) + // flatness[4] → scratch [127..131), ISV [206..210) (FLATNESS_BASE) + let base_wiener_offset = SP4_PRODUCER_COUNT as i32 * 3; // 213 + + for (isv_base, scratch_base) in [ + (BUDGET_C51_BASE, SCRATCH_PEARL_2_C51), + (BUDGET_IQN_BASE, SCRATCH_PEARL_2_IQN), + (BUDGET_CQL_BASE, SCRATCH_PEARL_2_CQL), + (BUDGET_ENS_BASE, SCRATCH_PEARL_2_ENS), + (FLATNESS_BASE, SCRATCH_PEARL_2_FLATNESS), + ] { + for b in 0..4_usize { + let scratch_idx = (scratch_base + b) as i32; + let isv_idx = (isv_base + b) as i32; + let wiener_off = base_wiener_offset + (isv_idx - SP5_SLOT_BASE as i32) * 3; + unsafe { + launch_apply_pearls( + &self.stream, + &self.apply_pearls_ad_kernel, + scratch_dev, scratch_idx, + isv_dev, isv_idx, + wiener_dev, wiener_off, + 1, + )?; + } + } + } + + Ok(()) + } + /// Plan C Phase 2 follow-up A.2 (2026-04-29): launch `q_drift_rate_ema_update` /// — single-thread single-block ISV producer for ISV[Q_DRIFT_RATE_INDEX=129]. /// @@ -12617,6 +12773,17 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("pearl_3_sigma_update load: {e}")))? }; + // SP5 Task A3 (2026-05-01): load pearl_2_budget loss budget controller kernel. + // Single-block 4-thread kernel; reads Q_VAR_PER_BRANCH (ISV[222..226)) + + // NOISY_SIGMA (ISV[210..214)) from ISV; writes 20 floats to scratch[111..131). + // Must be invoked AFTER pearl_1_atom (Q_VAR) AND pearl_3_sigma (NOISY_SIGMA). + let pearl_2_budget_kernel = { + let module = stream.context().load_cubin(SP5_PEARL_2_BUDGET_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("sp5 pearl_2_budget cubin load: {e}")))?; + module.load_function("pearl_2_budget_update") + .map_err(|e| MLError::ModelError(format!("pearl_2_budget_update load: {e}")))? + }; + // SP5 Task A1 (2026-05-01): allocate atoms_clip_count_buf [4 i32]. // Zero-initialized; populated by atoms_update_kernel in Layer B. // Until Layer B wires, clip_rate=0 → headroom holds at 6.0 bootstrap. @@ -15738,6 +15905,7 @@ impl GpuDqnTrainer { q_branch_stats_kernel, pearl_1_atom_kernel, pearl_3_sigma_kernel, + pearl_2_budget_kernel, atoms_clip_count_buf, action_counts_buf, action_offsets_buf, diff --git a/crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu b/crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu new file mode 100644 index 000000000..8bd2aa894 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu @@ -0,0 +1,69 @@ +// crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu +// +// SP5 Task A3: Pearl 2 — per-branch C51/IQN/CQL/Ens loss budgets driven by +// per-branch flatness signal. +// +// Reads Q_VAR_PER_BRANCH[b] (ISV[222..226), Pearl 1's q_branch_stats output) and +// NOISY_SIGMA[b] (ISV[210..214), Pearl 3 output). Computes: +// +// flatness[b] = clamp(var_q[b] / (σ[b]² + EPS_DIV), 0, 1) +// +// When flatness ≈ 1 (Q values are flat relative to noise σ), IQN dominates; +// when flatness ≈ 0 (Q values are well-separated), C51 takes a proportional share. +// +// budget_iqn[b] = BASE_IQN + (1 - flatness[b]) × (1 - BASE_IQN) +// budget_c51[b] = flatness[b] × (1 - BASE_IQN) × BASE_C51_SHARE +// budget_cql[b] = 0 // CQL stays gated by existing regime_stability allocator +// budget_ens[b] = max(0, 1 - iqn - c51 - cql) +// +// Writes per-branch (budget_c51, budget_iqn, budget_cql, budget_ens, flatness) +// to scratch_out. apply_pearls_ad_kernel then smooths these via Pearls A+D. +// +// 4-thread single-block (one thread per branch). No atomicAdd. +// __threadfence_system() after writes. +// +// BASE_IQN=0.11, BASE_C51_SHARE=0.84/(1-BASE_IQN), and EPS_DIV=1e-8 are +// Invariant 1 structural anchors (preserve SP4-baseline budget under non-flat +// regime; small numerical-stability epsilon). + +#include + +extern "C" __global__ void pearl_2_budget_update( + const float* __restrict__ isv_signals, + int q_var_isv_base, // Q_VAR_PER_BRANCH_BASE = 222 + int sigma_isv_base, // NOISY_SIGMA_BASE = 210 + float* __restrict__ scratch_out, + int c51_idx_base, + int iqn_idx_base, + int cql_idx_base, + int ens_idx_base, + int flatness_idx_base +) { + int b = threadIdx.x; + if (blockIdx.x != 0 || b >= 4) return; + + // BASE_IQN = 0.11 is an Invariant 1 anchor (SP4 baseline budget share). + const float BASE_IQN = 0.11f; + // BASE_C51_SHARE = 0.84 / (1 - BASE_IQN) is an Invariant 1 anchor + // (proportional remainder preserving SP4 baseline C51 share under flat regime). + const float BASE_C51_SHARE = 0.84f / (1.0f - BASE_IQN); + // EPS_DIV = 1e-8 is an Invariant 1 anchor (numerical stability denominator floor). + const float EPS_DIV = 1e-8f; + + float var_q = isv_signals[q_var_isv_base + b]; + float sigma = isv_signals[sigma_isv_base + b]; + float sigma_sq = sigma * sigma + EPS_DIV; + float flatness = fminf(1.0f, fmaxf(0.0f, var_q / sigma_sq)); + + float budget_iqn = BASE_IQN + (1.0f - flatness) * (1.0f - BASE_IQN); + float budget_c51 = flatness * (1.0f - BASE_IQN) * BASE_C51_SHARE; + float budget_cql = 0.0f; + float budget_ens = fmaxf(0.0f, 1.0f - budget_iqn - budget_c51 - budget_cql); + + scratch_out[c51_idx_base + b] = budget_c51; + scratch_out[iqn_idx_base + b] = budget_iqn; + scratch_out[cql_idx_base + b] = budget_cql; + scratch_out[ens_idx_base + b] = budget_ens; + scratch_out[flatness_idx_base + b] = flatness; + __threadfence_system(); +} diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index ee6052044..ddb743384 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -572,6 +572,31 @@ impl StateResetRegistry { category: ResetCategory::FoldReset, description: "ISV[ATOM_CLIP_RATE_BASE=186..190) — per-branch atom clip rate (atoms_clip_count / (batch_size*13)). SP5 Pearl A sentinel 0 at fold boundary (Task A1).", }, + RegistryEntry { + name: "sp5_budget_c51", + category: ResetCategory::FoldReset, + description: "ISV[BUDGET_C51_BASE=190..194) — per-branch C51 loss budget share (driven by flatness × (1 - BASE_IQN) × BASE_C51_SHARE). SP5 Pearl A sentinel 0 at fold boundary so the new fold's first pearl_2_budget_update launch fires the first-observation replacement (Task A3).", + }, + RegistryEntry { + name: "sp5_budget_iqn", + category: ResetCategory::FoldReset, + description: "ISV[BUDGET_IQN_BASE=194..198) — per-branch IQN loss budget share (BASE_IQN + (1-flatness)×(1-BASE_IQN), dominates when Q is flat relative to noise σ). SP5 Pearl A sentinel 0 at fold boundary (Task A3).", + }, + RegistryEntry { + name: "sp5_budget_cql", + category: ResetCategory::FoldReset, + description: "ISV[BUDGET_CQL_BASE=198..202) — per-branch CQL loss budget share (zero in Pearl 2; CQL gating remains under the existing regime_stability allocator). SP5 Pearl A sentinel 0 at fold boundary (Task A3).", + }, + RegistryEntry { + name: "sp5_budget_ens", + category: ResetCategory::FoldReset, + description: "ISV[BUDGET_ENS_BASE=202..206) — per-branch ensemble loss budget share (1 - IQN - C51 - CQL, clamped to ≥0). SP5 Pearl A sentinel 0 at fold boundary (Task A3).", + }, + RegistryEntry { + name: "sp5_flatness", + category: ResetCategory::FoldReset, + description: "ISV[FLATNESS_BASE=206..210) — per-branch flatness signal (var(Q[b]) / (σ[b]² + EPS_DIV), clamped to [0, 1]). Drives Pearl 2's per-branch IQN/C51 split. SP5 Pearl A sentinel 0 at fold boundary (Task A3).", + }, RegistryEntry { name: "sp5_noisy_sigma", category: ResetCategory::FoldReset, diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 279f66b55..0a5084607 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -3596,6 +3596,13 @@ impl DQNTrainer { if let Err(e) = fused.trainer().launch_sp5_pearl_3_sigma() { tracing::warn!("SP5 A2 launch_sp5_pearl_3_sigma failed: {e}"); } + // SP5 Task A3: pearl_2_budget → apply_pearls_ad ×20 + // Must run AFTER launch_sp5_pearl_1_atom (Q_VAR) AND + // launch_sp5_pearl_3_sigma (NOISY_SIGMA). + // Producer-only (Layer A) — no consumer reads ISV[190..210) in this commit. + if let Err(e) = fused.trainer().launch_sp5_pearl_2_budget() { + tracing::warn!(error = %e, "SP5 Pearl 2 producer launch failed"); + } } } diff --git a/crates/ml/tests/sp5_producer_unit_tests.rs b/crates/ml/tests/sp5_producer_unit_tests.rs index c96eac5ec..27f78b056 100644 --- a/crates/ml/tests/sp5_producer_unit_tests.rs +++ b/crates/ml/tests/sp5_producer_unit_tests.rs @@ -30,6 +30,9 @@ const SP5_PEARL_1_ATOM_CUBIN: &[u8] = const SP5_PEARL_3_SIGMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/pearl_3_sigma_kernel.cubin")); +const SP5_PEARL_2_BUDGET_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/pearl_2_budget_kernel.cubin")); + // ── Helpers ─────────────────────────────────────────────────────────────────── fn load_pearl_3_sigma_kernel(stream: &Arc) -> CudaFunction { @@ -42,6 +45,16 @@ fn load_pearl_3_sigma_kernel(stream: &Arc) -> CudaFunction { .expect("load pearl_3_sigma_update function") } +fn load_pearl_2_budget_kernel(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP5_PEARL_2_BUDGET_CUBIN.to_vec()) + .expect("load pearl_2_budget_kernel cubin"); + module + .load_function("pearl_2_budget_update") + .expect("load pearl_2_budget_update function") +} + fn make_test_stream() -> Arc { let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?"); ctx.default_stream() @@ -563,3 +576,271 @@ fn pearl_3_sigma_fraction_increases_when_entropy_below_target() { assert!(sf > 0.1_f32, "branch={b}: SF should have increased from 0.1, got {sf}"); } } + +// ── Tests 5 + 6: pearl_2_budget_update ──────────────────────────────────────── + +/// SP5 Task A3 unit test: `pearl_2_budget_update` on flat-Q regime. +/// +/// Synthetic ISV: σ=1.0 (all branches), var_q=2.0 (all branches). +/// flatness = clamp(2.0 / (1.0² + 1e-8), 0, 1) = clamp(~2.0, 0, 1) = 1.0 +/// +/// Analytical expected (flatness = 1.0): +/// BASE_IQN = 0.11 +/// BASE_C51_SHARE = 0.84 / (1 - 0.11) = 0.84 / 0.89 ≈ 0.94382... +/// budget_iqn = 0.11 + (1 - 1.0) * 0.89 = 0.11 +/// budget_c51 = 1.0 * 0.89 * 0.94382 ≈ 0.84 +/// budget_cql = 0 +/// budget_ens = max(0, 1 - 0.11 - 0.84 - 0) = 0.05 +/// +/// All 4 branches receive identical inputs, so identical outputs. +/// No CPU reference oracle per `feedback_no_cpu_test_fallbacks.md`. +#[test] +#[ignore = "requires GPU"] +fn pearl_2_budget_flat_regime_iqn_floor_c51_takes_remainder() { + // Scratch layout mirrors SCRATCH_PEARL_2_* constants: + // [111..115) = budget_c51[4] + // [115..119) = budget_iqn[4] + // [119..123) = budget_cql[4] + // [123..127) = budget_ens[4] + // [127..131) = flatness[4] + const SCRATCH_C51: usize = 111; + const SCRATCH_IQN: usize = 115; + const SCRATCH_CQL: usize = 119; + const SCRATCH_ENS: usize = 123; + const SCRATCH_FLATNESS: usize = 127; + const SCRATCH_SIZE: usize = 131; + + // ISV: Q_VAR_PER_BRANCH_BASE=222, NOISY_SIGMA_BASE=210. + // Need ISV size >= 222 + 4 = 226. + const ISV_SIZE: usize = 230; + // Q_VAR_PER_BRANCH_BASE = 222, NOISY_SIGMA_BASE = 210. + const Q_VAR_BASE: usize = 222; + const SIGMA_BASE: usize = 210; + + let stream = make_test_stream(); + let kernel = load_pearl_2_budget_kernel(&stream); + + let isv_buf = unsafe { MappedF32Buffer::new(ISV_SIZE) } + .expect("alloc isv_buf"); + let mut isv_data = vec![0.0_f32; ISV_SIZE]; + for b in 0..4_usize { + isv_data[Q_VAR_BASE + b] = 2.0_f32; // var_q = 2.0 (above σ² → clamp to 1.0 flatness) + isv_data[SIGMA_BASE + b] = 1.0_f32; // σ = 1.0 + } + isv_buf.write_from_slice(&isv_data); + + let scratch_buf = unsafe { MappedF32Buffer::new(SCRATCH_SIZE) } + .expect("alloc scratch_buf"); + + let isv_dev = isv_buf.dev_ptr; + let q_var_base_i32 = Q_VAR_BASE as i32; + let sigma_base_i32 = SIGMA_BASE as i32; + let scratch_dev = scratch_buf.dev_ptr; + let c51_base_i32 = SCRATCH_C51 as i32; + let iqn_base_i32 = SCRATCH_IQN as i32; + let cql_base_i32 = SCRATCH_CQL as i32; + let ens_base_i32 = SCRATCH_ENS as i32; + let flatness_base_i32 = SCRATCH_FLATNESS as i32; + + unsafe { + stream + .launch_builder(&kernel) + .arg(&isv_dev) + .arg(&q_var_base_i32) + .arg(&sigma_base_i32) + .arg(&scratch_dev) + .arg(&c51_base_i32) + .arg(&iqn_base_i32) + .arg(&cql_base_i32) + .arg(&ens_base_i32) + .arg(&flatness_base_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (4, 1, 1), + shared_mem_bytes: 0, + }) + .expect("launch pearl_2_budget_update (flat regime)"); + } + stream.synchronize().expect("sync after pearl_2_budget_update (flat regime)"); + + let scratch = scratch_buf.read_all(); + + // Analytical values for flatness = 1.0 (saturated ceiling): + // BASE_IQN = 0.11, BASE_C51_SHARE = 0.84 / 0.89 + const BASE_IQN: f32 = 0.11_f32; + const BASE_C51_SHARE: f32 = 0.84_f32 / (1.0_f32 - BASE_IQN); + let expected_flatness = 1.0_f32; + let expected_iqn = BASE_IQN + (1.0_f32 - expected_flatness) * (1.0_f32 - BASE_IQN); + let expected_c51 = expected_flatness * (1.0_f32 - BASE_IQN) * BASE_C51_SHARE; + let expected_cql = 0.0_f32; + let expected_ens = (1.0_f32 - expected_iqn - expected_c51 - expected_cql).max(0.0_f32); + + for b in 0..4_usize { + let c51 = scratch[SCRATCH_C51 + b]; + let iqn = scratch[SCRATCH_IQN + b]; + let cql = scratch[SCRATCH_CQL + b]; + let ens = scratch[SCRATCH_ENS + b]; + let flatness = scratch[SCRATCH_FLATNESS + b]; + + println!( + "branch={b} flatness={flatness:.6} iqn={iqn:.6} c51={c51:.6} \ + cql={cql:.6} ens={ens:.6}" + ); + + let rel_flatness = (flatness - expected_flatness).abs(); + assert!( + rel_flatness < 1e-5, + "branch={b}: flatness={flatness:.6} expected={expected_flatness:.6}" + ); + let rel_iqn = (iqn - expected_iqn).abs() / expected_iqn.max(1e-9); + assert!( + rel_iqn < 0.01, + "branch={b}: iqn={iqn:.6} expected={expected_iqn:.6} rel={rel_iqn:.5}" + ); + let rel_c51 = (c51 - expected_c51).abs() / expected_c51.max(1e-9); + assert!( + rel_c51 < 0.01, + "branch={b}: c51={c51:.6} expected={expected_c51:.6} rel={rel_c51:.5}" + ); + assert!( + cql.abs() < 1e-7, + "branch={b}: cql={cql:.6} expected=0" + ); + let rel_ens = (ens - expected_ens).abs() / expected_ens.max(1e-9); + assert!( + rel_ens < 0.01, + "branch={b}: ens={ens:.6} expected={expected_ens:.6} rel={rel_ens:.5}" + ); + } +} + +/// SP5 Task A3 unit test: `pearl_2_budget_update` on sharp-Q regime. +/// +/// Synthetic ISV: σ=2.0 (all branches), var_q=0.04 (all branches). +/// flatness = clamp(0.04 / (4.0 + 1e-8), 0, 1) = 0.04/4.0 = 0.01 +/// +/// Analytical expected (flatness = 0.01): +/// BASE_IQN = 0.11 +/// BASE_C51_SHARE = 0.84 / 0.89 ≈ 0.94382 +/// budget_iqn = 0.11 + (1 - 0.01) * 0.89 = 0.11 + 0.8811 = 0.9911 +/// budget_c51 = 0.01 * 0.89 * 0.94382 ≈ 0.0084 +/// budget_cql = 0 +/// budget_ens = max(0, 1 - 0.9911 - 0.0084 - 0) ≈ 0.0005 (rounds to ~0) +/// +/// IQN dominates in the sharp-Q regime, C51 share is minimal. +/// No CPU reference oracle per `feedback_no_cpu_test_fallbacks.md`. +#[test] +#[ignore = "requires GPU"] +fn pearl_2_budget_sharp_regime_iqn_dominates() { + const SCRATCH_C51: usize = 111; + const SCRATCH_IQN: usize = 115; + const SCRATCH_CQL: usize = 119; + const SCRATCH_ENS: usize = 123; + const SCRATCH_FLATNESS: usize = 127; + const SCRATCH_SIZE: usize = 131; + const ISV_SIZE: usize = 230; + const Q_VAR_BASE: usize = 222; + const SIGMA_BASE: usize = 210; + + let stream = make_test_stream(); + let kernel = load_pearl_2_budget_kernel(&stream); + + let isv_buf = unsafe { MappedF32Buffer::new(ISV_SIZE) } + .expect("alloc isv_buf"); + let mut isv_data = vec![0.0_f32; ISV_SIZE]; + for b in 0..4_usize { + isv_data[Q_VAR_BASE + b] = 0.04_f32; // var_q small + isv_data[SIGMA_BASE + b] = 2.0_f32; // σ large → flatness small + } + isv_buf.write_from_slice(&isv_data); + + let scratch_buf = unsafe { MappedF32Buffer::new(SCRATCH_SIZE) } + .expect("alloc scratch_buf"); + + let isv_dev = isv_buf.dev_ptr; + let q_var_base_i32 = Q_VAR_BASE as i32; + let sigma_base_i32 = SIGMA_BASE as i32; + let scratch_dev = scratch_buf.dev_ptr; + let c51_base_i32 = SCRATCH_C51 as i32; + let iqn_base_i32 = SCRATCH_IQN as i32; + let cql_base_i32 = SCRATCH_CQL as i32; + let ens_base_i32 = SCRATCH_ENS as i32; + let flatness_base_i32 = SCRATCH_FLATNESS as i32; + + unsafe { + stream + .launch_builder(&kernel) + .arg(&isv_dev) + .arg(&q_var_base_i32) + .arg(&sigma_base_i32) + .arg(&scratch_dev) + .arg(&c51_base_i32) + .arg(&iqn_base_i32) + .arg(&cql_base_i32) + .arg(&ens_base_i32) + .arg(&flatness_base_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (4, 1, 1), + shared_mem_bytes: 0, + }) + .expect("launch pearl_2_budget_update (sharp regime)"); + } + stream.synchronize().expect("sync after pearl_2_budget_update (sharp regime)"); + + let scratch = scratch_buf.read_all(); + + // Analytical values for flatness = 0.04 / (4.0 + 1e-8) ≈ 0.01: + let expected_sigma_sq = 2.0_f32 * 2.0_f32 + 1e-8_f32; + let expected_flatness = (0.04_f32 / expected_sigma_sq).clamp(0.0, 1.0); + const BASE_IQN: f32 = 0.11_f32; + const BASE_C51_SHARE: f32 = 0.84_f32 / (1.0_f32 - BASE_IQN); + let expected_iqn = BASE_IQN + (1.0_f32 - expected_flatness) * (1.0_f32 - BASE_IQN); + let expected_c51 = expected_flatness * (1.0_f32 - BASE_IQN) * BASE_C51_SHARE; + let expected_cql = 0.0_f32; + let expected_ens = (1.0_f32 - expected_iqn - expected_c51 - expected_cql).max(0.0_f32); + + for b in 0..4_usize { + let c51 = scratch[SCRATCH_C51 + b]; + let iqn = scratch[SCRATCH_IQN + b]; + let cql = scratch[SCRATCH_CQL + b]; + let ens = scratch[SCRATCH_ENS + b]; + let flatness = scratch[SCRATCH_FLATNESS + b]; + + println!( + "branch={b} flatness={flatness:.6} iqn={iqn:.6} c51={c51:.6} \ + cql={cql:.6} ens={ens:.6}" + ); + + let rel_flatness = (flatness - expected_flatness).abs() / expected_flatness.max(1e-9); + assert!( + rel_flatness < 0.01, + "branch={b}: flatness={flatness:.6} expected={expected_flatness:.6} rel={rel_flatness:.5}" + ); + let rel_iqn = (iqn - expected_iqn).abs() / expected_iqn.max(1e-9); + assert!( + rel_iqn < 0.01, + "branch={b}: iqn={iqn:.6} expected={expected_iqn:.6} rel={rel_iqn:.5}" + ); + // IQN must dominate: iqn >> c51 + assert!( + iqn > 0.9_f32, + "branch={b}: iqn={iqn:.6} should exceed 0.9 in sharp regime" + ); + let rel_c51 = (c51 - expected_c51).abs() / expected_c51.max(1e-9); + assert!( + rel_c51 < 0.01, + "branch={b}: c51={c51:.6} expected={expected_c51:.6} rel={rel_c51:.5}" + ); + assert!( + cql.abs() < 1e-7, + "branch={b}: cql={cql:.6} expected=0" + ); + // ens is near-zero in sharp regime; just verify it is non-negative + assert!( + ens >= 0.0_f32, + "branch={b}: ens={ens:.6} must be non-negative" + ); + let _ = expected_ens; // expected is ~0.0005, just verify non-negative above + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 7f0e2d8f9..44f83ae0d 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,8 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +SP5 Task A3 — Pearl 2 per-branch loss budget GPU producer (Layer A, 2026-05-01): one new CUDA kernel (`pearl_2_budget_kernel.cu`) lands as a Layer A additive producer that feeds ISV slots [190..210) with per-branch loss budget and flatness signals. No loss-budget consumer migration in this commit — Layer B wires consumers. `pearl_2_budget_update`: single-block 4-thread kernel (one thread per branch), reads Q_VAR_PER_BRANCH[b] (Pearl 1's q_branch_stats output, ISV[222..226)) + NOISY_SIGMA[b] (Pearl 3 output, ISV[210..214)); computes `flatness[b] = clamp(var_q[b] / (σ[b]² + EPS_DIV), 0, 1)` where EPS_DIV=1e-8 is an Invariant 1 anchor (numerical stability); derives `budget_iqn = BASE_IQN + (1 - flatness) * (1 - BASE_IQN)` (IQN dominates when Q is flat relative to noise), `budget_c51 = flatness * (1 - BASE_IQN) * BASE_C51_SHARE` (C51 takes proportional remainder when Q is well-separated), `budget_cql = 0` (gated by existing regime_stability allocator), `budget_ens = max(0, 1 - iqn - c51 - cql)`; writes 20 floats to scratch[111..131) (SCRATCH_PEARL_2_C51=111, SCRATCH_PEARL_2_IQN=115, SCRATCH_PEARL_2_CQL=119, SCRATCH_PEARL_2_ENS=123, SCRATCH_PEARL_2_FLATNESS=127). BASE_IQN=0.11 and BASE_C51_SHARE=0.84/(1-BASE_IQN) are Invariant 1 anchors (preserve SP4-baseline budget under non-flat regime). `launch_sp5_pearl_2_budget` fires 20 `launch_apply_pearls` calls: 5 output groups × 4 branches each, mapping to ISV[BUDGET_C51_BASE=190..194), ISV[BUDGET_IQN_BASE=194..198), ISV[BUDGET_CQL_BASE=198..202), ISV[BUDGET_ENS_BASE=202..206), ISV[FLATNESS_BASE=206..210). Wiener offset formula: `base_wiener_offset + (isv_slot - SP5_SLOT_BASE) * 3` where `base_wiener_offset = SP4_PRODUCER_COUNT * 3 = 213`. Dependencies: Pearl 2 must chain AFTER Pearl 1 (reads Q_VAR_PER_BRANCH) AND AFTER Pearl 3 (reads NOISY_SIGMA). Buffer growth: `producer_step_scratch_buf` 111→131 slots (SP5_SCRATCH_TOTAL updated 111→131; 5 new module-level constants SCRATCH_PEARL_2_{C51,IQN,CQL,ENS,FLATNESS}). wiener_state_buf already at 543 (A1 sized for entire SP5 block — no further growth). StateResetRegistry gains 5 new FoldReset entries: sp5_budget_c51 (ISV[190..194)), sp5_budget_iqn (ISV[194..198)), sp5_budget_cql (ISV[198..202)), sp5_budget_ens (ISV[202..206)), sp5_flatness (ISV[206..210)) — Pearl A sentinel 0 at fold boundary triggers first-observation replacement on new fold's first launch. Wire-up in training_loop.rs: `launch_sp5_pearl_2_budget()` called immediately after `launch_sp5_pearl_3_sigma()` with `tracing::warn!(error = %e, ...)` on error. Two GPU-only `#[ignore = "requires GPU"]` unit tests: (5) flat-Q regime (σ=1.0, var_q=2.0 → flatness=1.0 → iqn=0.11, c51≈0.84); (6) sharp-Q regime (σ=2.0, var_q=0.04 → flatness≈0.01 → iqn>0.9, c51≈0.0084). No CPU compute, no HtoD/HtoH, no atomicAdd, no stubs, no consumer migration. Touched: `cuda_pipeline/pearl_2_budget_kernel.cu` (new), `build.rs` (+1 cubin entry), `cuda_pipeline/gpu_dqn_trainer.rs` (static cubin + SP5_SCRATCH_TOTAL 111→131 + 5 new scratch constants + docstring update + struct field + cubin loader + struct initializer + launch method), `trainers/dqn/state_reset_registry.rs` (+5 FoldReset entries), `trainers/dqn/trainer/training_loop.rs` (+5 LOC after Pearl 3 launch), `tests/sp5_producer_unit_tests.rs` (+2 GPU-only tests + cubin loader helper). Hard rules: feedback_no_cpu_compute_strict, feedback_no_atomicadd, feedback_no_stubs, feedback_no_partial_refactor (all consumers of existing budget allocator untouched). + SP5 Task A2 — Pearl 3 per-branch NoisyNet sigma GPU producer (Layer A, 2026-05-01): one new CUDA kernel (`pearl_3_sigma_kernel.cu`) lands as a Layer A additive producer that feeds ISV slots [210..218) with per-branch NoisyNet sigma and sigma-fraction controller state. No NoisyLinear consumer migration in this commit — Layer B wires the noisy_linear_kernel.cu consumer. `pearl_3_sigma_update`: single-block 4-thread kernel (one thread per branch), reads ATOM_V_HALF[b] (Pearl 1, ISV[178..182)) + BRANCH_ENTROPY[b] (q_branch_stats, ISV[218..222)) + current SIGMA_FRACTION[b] (ISV[214..218), Pearl A sentinel 0 → bootstrap to 0.1); applies entropy-deficit controller `new_sf = clamp(sf + SF_LR * (target_entropy - entropy), 0.001, 0.5)` where `target_entropy = log(n_actions[b]) * 0.7` (SF_LR=0.005 and target_entropy_factor=0.7 are Invariant 1 structural anchors); computes `new_sigma = new_sf * max(v_half, EPS_CLAMP_FLOOR=1.0)` (EPS_CLAMP_FLOOR is an Invariant 1 anchor); writes (sigma[4], SF[4]) to scratch[103..111) (SCRATCH_PEARL_3_SIGMA=103, SCRATCH_PEARL_3_SF=107). `launch_sp5_pearl_3_sigma` fires 8 `launch_apply_pearls` calls: sigma[4] → ISV[210..214) (NOISY_SIGMA_BASE), SF[4] → ISV[214..218) (SIGMA_FRACTION_BASE). Wiener offset formula: `base_wiener_offset + (isv_slot - SP5_SLOT_BASE) * 3` where `base_wiener_offset = SP4_PRODUCER_COUNT * 3 = 213`. Pearl 3 must chain AFTER Pearl 1 (reads ATOM_V_HALF which Pearl 1 populates via apply_pearls_ad). Buffer growth: `producer_step_scratch_buf` 103→111 (SP5_SCRATCH_TOTAL updated; new constants SCRATCH_PEARL_3_SIGMA=103, SCRATCH_PEARL_3_SF=107). wiener_state_buf already at 543 (A1 sized for entire SP5 block — no further growth). StateResetRegistry gains 2 new FoldReset entries: sp5_noisy_sigma (ISV[210..214)) and sp5_sigma_fraction (ISV[214..218)) — Pearl A sentinel 0 at fold boundary triggers bootstrap to 0.1 on new fold's first launch. Wire-up in training_loop.rs: `launch_sp5_pearl_3_sigma()` called immediately after `launch_sp5_pearl_1_atom()` with `warn!` on error. Two GPU-only `#[ignore = "requires GPU"]` unit tests: (3) bootstrap + target-entropy equality case (sentinel ISV SF=0→bootstrap 0.1, target==actual → no update, sigma=0.1*2.0=0.2); (4) entropy-below-target case (SF=0.1, entropy=0 → SF increases by SF_LR*target_entropy, sigma increases proportionally). No CPU compute, no HtoD/HtoH, no atomicAdd, no stubs. Touched: `cuda_pipeline/pearl_3_sigma_kernel.cu` (new), `build.rs` (+1 cubin entry), `cuda_pipeline/gpu_dqn_trainer.rs` (static cubin + SP5_SCRATCH_TOTAL 103→111 + 2 new constants + struct field + cubin loader + struct initializer + launch method), `trainers/dqn/state_reset_registry.rs` (+2 FoldReset entries), `trainers/dqn/trainer/training_loop.rs` (+4 LOC after Pearl 1 launch), `tests/sp5_producer_unit_tests.rs` (+2 GPU-only tests + cubin loader helper). cargo check -p ml --lib clean. No consumer migration in this commit. SP5 Task A1 fix-up — pearl_1 clip_rate denominator must be per-branch (2026-05-01): code-quality review caught the headroom controller's clip_rate computation using `batch_size × 13` (sum of per-branch action counts) instead of `batch_size × action_counts[b]` (per-branch). The TARGET_CLIP_RATE=0.01 anchor is calibrated against the per-branch event rate; using a denominator 3.25–4.3× too large made the controller systematically under-responsive — at actual 1% per-branch clipping it would read ~0.23–0.31% and contract headroom toward the 2.0 floor instead of holding the target. Currently masked because `atoms_clip_count` is zero in Layer A (Layer B's atoms_update_kernel migration populates it); without this fix Layer B would inherit a silently-wrong formula. Fix plumbs `action_counts[4]` ({4, 3, 3, 3}) through pearl_1_atom_kernel signature and uses `batch_size × action_counts[b]` as the per-branch denominator, matching q_branch_stats_kernel's existing parameter pattern. Test #2 launcher updated to allocate action_counts_buf and pass its dev_ptr through the new parameter slot. Touched: `cuda_pipeline/pearl_1_atom_kernel.cu` (+1 param, denominator math), `cuda_pipeline/gpu_dqn_trainer.rs` (+1 launch arg in `launch_sp5_pearl_1_atom`), `tests/sp5_producer_unit_tests.rs` (+5 LOC for action_counts_buf + extra arg). cargo check + cargo build --release + cargo test --no-run all clean.