From 6135dea3113678dfc6fe1e1fec57bb5c3d1b1035 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 25 Apr 2026 15:41:09 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn-v2):=20Plan=204=20Task=202c.3c.5=20?= =?UTF-8?q?=E2=80=94=20append=20H=5FS2=5FRMS=5FEMA=20ISV=20slot=20+=20prod?= =?UTF-8?q?ucer=20kernel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ISV bus tail-append: - [96] H_S2_RMS_EMA_INDEX — per-batch RMS(save_h_s2) EMA (alpha=0.05) - [97] ISV_LAYOUT_FINGERPRINT_LO_INDEX (shifted from 94) - [98] ISV_LAYOUT_FINGERPRINT_HI_INDEX (shifted from 95) - ISV_TOTAL_DIM 96 -> 99 - New LAYOUT_FINGERPRINT_CURRENT = 0x3e21acecd922e540 (was 0xcf3a24b0a1f70057) Producer kernel `h_s2_rms_ema_kernel.cu`: - Single-block reduction (256 threads, shmem-tree, no atomicAdd) of the online trunk's `save_h_s2 [B, SH2]` post-GRN activation - RMS = sqrt(sum_sq / (B*SH2)); EMA into ISV[96] with alpha=0.05 - Launched once per training step alongside reward_component_ema and the other Plan 3/4 ISV producers in training_loop.rs Constructor cold-start writes ISV[96]=1.0 (neutral RMS). StateResetRegistry: H_S2_RMS_EMA registered as FoldReset -> 1.0. `launch_h_s2_rms_ema` and the kernel both treat the unconditional constructor allocation as an invariant — debug_assert! on the host side, no defensive (N>0) ternary in the kernel — same proper-resolution pattern as the 2c.3c.4 followup. Producer-only commit. 2c.3c.6 wires the consumer in mag_concat_qdir's adaptive-scale path so the magnitude-branch decoder sees a scale- invariant residual stack regardless of GRN drift across training. Smoke deferred — producer-only with zero consumers, runtime behaviour structurally unchanged from 2c.3c.4 (which validated the GRN backward chain end-to-end with all 3 fold checkpoints). cargo check + cargo build (new h_s2_rms_ema_kernel.cubin compiles clean under nvcc sm_80) is the standing validation; meaningful multi_fold_convergence re-run lands with 2c.3c.6 when the consumer wires up. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 6 + .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 119 ++++++++++++++++-- .../src/cuda_pipeline/h_s2_rms_ema_kernel.cu | 69 ++++++++++ .../src/trainers/dqn/state_reset_registry.rs | 6 + .../src/trainers/dqn/trainer/training_loop.rs | 25 ++++ docs/dqn-wire-up-audit.md | 2 + 6 files changed, 215 insertions(+), 12 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 30ffdc938..b4fcc06e5 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -131,6 +131,12 @@ fn main() { // Module is additive — no production callers in this commit; Task 2c.3+4 // wires it into the trunk encoder. "grn_kernel.cu", + // Plan 4 Task 2c.3c.5: per-batch RMS EMA of `save_h_s2` into ISV[96]. + // Single-block shmem-reduce kernel (256 threads, no atomicAdd) launched + // alongside `reward_component_ema` from `training_loop.rs`. Producer-only + // in this commit — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s + // adaptive-scale path. + "h_s2_rms_ema_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index c3660789e..fad550d4f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -126,6 +126,12 @@ pub(crate) static TARGET_DRIFT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_D /// production callers added in 2c.3+4. pub(crate) static GRN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grn_kernel.cubin")); +/// Plan 4 Task 2c.3c.5: per-batch RMS EMA of `save_h_s2` into ISV[96]. +/// Single-block (256 threads, shmem-reduce, no atomicAdd) kernel launched +/// alongside `reward_component_ema` from `training_loop.rs`. Producer-only — +/// 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path. +static H_S2_RMS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_rms_ema_kernel.cubin")); + /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension @@ -289,11 +295,14 @@ const ISV_NETWORK_DIM: usize = 23; /// replaces legacy CPU-DtoH `per_branch_target_drift()`. Written by /// `target_drift_ema_update` GPU kernel. Magnitude branch. /// [93] TARGET_DRIFT_DIR_EMA_INDEX — same kernel/contract; direction branch. -/// [94] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint (shifted from 90). -/// [95] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint (shifted from 91). +/// [96] H_S2_RMS_EMA_INDEX — EMA of RMS(save_h_s2) per batch. Producer-only +/// in this commit (no consumers). Wired in 2c.3c.6 to drive the +/// adaptive scale of `mag_concat_qdir`'s residual stack. +/// [97] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint (shifted from 94). +/// [98] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint (shifted from 95). /// Written by the constructor; checked at checkpoint load. Fail-fast only — no migration /// path exists. See spec §4.A.2 and `LAYOUT_FINGERPRINT_CURRENT` for structural-hash rationale. -const ISV_TOTAL_DIM: usize = 96; +const ISV_TOTAL_DIM: usize = 99; /// Legacy alias preserved for call sites that haven't been audited for the /// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight /// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer). @@ -547,7 +556,21 @@ pub const TARGET_DRIFT_MAG_EMA_INDEX: usize = 92; /// ISV slot [93] — same producer/contract as [92]; direction branch. pub const TARGET_DRIFT_DIR_EMA_INDEX: usize = 93; -/// ISV slot [94] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits). +/// ISV slot [96] — Plan 4 Task 2c.3c.5: EMA of RMS(save_h_s2) per batch. +/// Producer: `h_s2_rms_ema_update` GPU kernel (single-block, 256 threads, +/// shmem-reduction, no atomicAdd). Cold-path-cadence — launched once per +/// experience-collection epoch alongside the other Plan 3/Plan 4 ISV producers +/// in `training_loop.rs`. EMA α = 0.05 (≈ 13-batch half-life). Cold-start +/// 1.0 (neutral RMS so first kernel fire EMAs measured RMS toward 1.0 +/// rather than 0). FoldReset → 1.0. +/// +/// Producer-only in this commit. 2c.3c.6 wires the consumer in +/// `mag_concat_qdir`'s adaptive-scale path (residual stack normalised by +/// the trunk-output RMS so the magnitude-branch decoder sees a scale- +/// invariant input regardless of the GRN trunk's drift). +pub const H_S2_RMS_EMA_INDEX: usize = 96; + +/// ISV slot [97] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits). /// /// Design note: this is NOT a version number. There is no ordered version space. /// The fingerprint is a structural hash that automatically changes whenever any @@ -560,14 +583,15 @@ pub const TARGET_DRIFT_DIR_EMA_INDEX: usize = 93; /// actively written by `isv_signal_update` (Q-drift and gradient-norm EMA). /// Shifted 61→69 in Plan 3 Task 1, 69→73 in Plan 3 Task 3, 73→76 in /// Plan 3 Task 4 B.4, 76→80 in Plan 3 Task 7 C.3, 80→85 in Plan 3 Task 8 B.3, -/// 85→90 in Plan 4 Task 5 Mode A, then 90→94 in Plan 4 follow-up (target-drift -/// EMA replacing legacy CPU-DtoH path). -pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 94; -/// ISV slot [95] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits). +/// 85→90 in Plan 4 Task 5 Mode A, 90→94 in Plan 4 follow-up (target-drift +/// EMA replacing legacy CPU-DtoH path), then 94→97 in Plan 4 Task 2c.3c.5 +/// (H_S2_RMS_EMA producer slot). +pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 97; +/// ISV slot [98] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits). /// Shifted 62→70 in Plan 3 Task 1, 70→74 in Plan 3 Task 3, 74→77 in Plan 3 Task 4 B.4, /// 77→81 in Plan 3 Task 7 C.3, 81→86 in Plan 3 Task 8 B.3, 86→91 in Plan 4 Task 5, -/// then 91→95 in Plan 4 follow-up. -pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 95; +/// 91→95 in Plan 4 follow-up, then 95→98 in Plan 4 Task 2c.3c.5. +pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 98; /// Canonical alias for the fingerprint slot (the low half). pub const ISV_LAYOUT_FINGERPRINT_INDEX: usize = ISV_LAYOUT_FINGERPRINT_LO_INDEX; @@ -645,8 +669,9 @@ const fn layout_fingerprint_seed() -> &'static [u8] { SEED_STEPS_TARGET=82;SEED_STEPS_DONE=83;SEED_FRAC_EMA=84;\ VSN_MAG_EMA=87;VSN_DIR_EMA=88;MAMBA2_RETENTION_EMA=89;\ TARGET_DRIFT_MAG_EMA=92;TARGET_DRIFT_DIR_EMA=93;\ - ISV_LAYOUT_FINGERPRINT_LO=94;ISV_LAYOUT_FINGERPRINT_HI=95;\ - ISV_TOTAL_DIM=96;\ + H_S2_RMS_EMA=96;\ + ISV_LAYOUT_FINGERPRINT_LO=97;ISV_LAYOUT_FINGERPRINT_HI=98;\ + ISV_TOTAL_DIM=99;\ PARAM_W_A_H_S1=0;PARAM_B_A_H_S1=1;PARAM_W_B_H_S1=2;PARAM_B_B_H_S1=3;\ PARAM_W_RESIDUAL_H_S1=4;PARAM_GAMMA_H_S1=5;PARAM_BETA_H_S1=6;\ PARAM_W_A_H_S2=7;PARAM_B_A_H_S2=8;PARAM_W_B_H_S2=9;PARAM_B_B_H_S2=10;\ @@ -2234,6 +2259,15 @@ pub struct GpuDqnTrainer { /// Cold-path: one launch per training step. Loaded from reward_component_ema_kernel.cubin. reward_component_ema_kernel: CudaFunction, + // ── Plan 4 Task 2c.3c.5: H_S2 RMS EMA producer ──────────────────── + /// Single-block kernel (256 threads, shmem-reduction, no atomicAdd). Reduces + /// RMS = sqrt(sum_sq / (B*SH2)) over `save_h_s2 [B, SH2]` and EMA-updates + /// ISV[H_S2_RMS_EMA_INDEX=96] with α = 0.05. Cold-path-cadence (one launch + /// per training step, alongside `reward_component_ema`). Producer-only — + /// 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path. + /// Loaded from `h_s2_rms_ema_kernel.cubin`. + h_s2_rms_ema_kernel: CudaFunction, + // ── Q-mean drift correction (Component 8) ────────────────────────── /// Phase 1 kernel: computes global mean of Q-values into q_mean_scratch (pinned). q_mean_reduce_kernel: CudaFunction, @@ -6674,6 +6708,49 @@ impl GpuDqnTrainer { Ok(()) } + /// Plan 4 Task 2c.3c.5: launch `h_s2_rms_ema_update` (single-block, + /// 256-thread shmem-reduction kernel). Reduces RMS = sqrt(sum_sq / (B*SH2)) + /// over the trainer's `save_h_s2 [B, SH2]` buffer (the online trunk's + /// post-GRN activation, valid after the most recent forward) and EMAs the + /// result into ISV[H_S2_RMS_EMA_INDEX=96]. + /// + /// Producer-only — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s + /// adaptive-scale path. Cold-path-cadence: launched once per training step + /// alongside `launch_reward_component_ema`. No atomicAdd, no DtoH. + pub fn launch_h_s2_rms_ema(&self, ema_alpha: f32) -> Result<(), MLError> { + // `isv_signals_dev_ptr` is set by the constructor via + // `cuMemHostGetDevicePointer_v2` on a `cuMemAllocHost_v2`-backed pinned + // buffer (with its own `assert_eq!` on success) and is never reassigned, + // so non-zero is an invariant. Silent-skip on `== 0` would mask a + // misconfiguration; `debug_assert!` flags it in dev builds and release + // inherits the loud kernel-launch error rather than a no-op. + debug_assert!(self.isv_signals_dev_ptr != 0, + "launch_h_s2_rms_ema: isv_signals_dev_ptr must be allocated by constructor"); + let h_s2_ptr = self.save_h_s2.raw_ptr(); + let b_i = self.config.batch_size as i32; + let sh2_i = self.config.shared_h2 as i32; + let isv_idx = H_S2_RMS_EMA_INDEX as i32; + let isv_dev_ptr = self.isv_signals_dev_ptr; + const BLOCK_DIM: u32 = 256; + let smem_bytes = BLOCK_DIM * std::mem::size_of::() as u32; + unsafe { + self.stream.launch_builder(&self.h_s2_rms_ema_kernel) + .arg(&h_s2_ptr) + .arg(&b_i) + .arg(&sh2_i) + .arg(&isv_dev_ptr) + .arg(&isv_idx) + .arg(&ema_alpha) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update: {e}")))?; + } + Ok(()) + } + /// Apply spectral normalization to all 10 weight matrices (trunk + 8 heads). /// /// One step of power iteration per call (standard practice — single step @@ -7345,6 +7422,16 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("reward_component_ema load: {e}")))? }; + // Plan 4 Task 2c.3c.5: load h_s2_rms_ema kernel (cold-path, per-step). + // Single-block 256-thread shmem-reduction kernel; no atomicAdd. Producer + // for ISV[H_S2_RMS_EMA_INDEX=96]; consumer wired in 2c.3c.6. + let h_s2_rms_ema_kernel = { + let module = stream.context().load_cubin(H_S2_RMS_EMA_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("h_s2_rms_ema cubin load: {e}")))?; + module.load_function("h_s2_rms_ema_update") + .map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update load: {e}")))? + }; + // Plan 2 Task 1 C.1: allocate branch-action offset/size device buffers for q_quantile_reduce. // Branch layout in q_out_buf: [dir | mag | ord | urg] row-major per sample. // dir occupies actions [0..b0), mag [b0..b0+b1), ord [b0+b1..b0+b1+b2), @@ -8699,6 +8786,13 @@ impl GpuDqnTrainer { // Plan 4 follow-up: target-drift EMA cold-start. *sig_ptr.add(TARGET_DRIFT_MAG_EMA_INDEX) = 0.0_f32; *sig_ptr.add(TARGET_DRIFT_DIR_EMA_INDEX) = 0.0_f32; + /* Plan 4 Task 2c.3c.5: H_S2 RMS EMA cold-start. Set to 1.0 + * (neutral RMS) so the first kernel fire's α=0.05 EMA pulls + * the slot toward the measured RMS rather than collapsing + * toward 0. Half-life ≈ 13 batches (log(0.5)/log(0.95)). + * Producer-only in this commit; 2c.3c.6 wires the consumer + * in `mag_concat_qdir`'s adaptive-scale path. FoldReset → 1.0. */ + *sig_ptr.add(H_S2_RMS_EMA_INDEX) = 1.0_f32; // Plan 2 C.1 Q-quantile bootstrap (2026-04-24): // q_p05 = v_min, q_p95 = v_max — matches the cold-start atom range // so the first update_eval_v_range call reads meaningful bounds. @@ -9714,6 +9808,7 @@ impl GpuDqnTrainer { popart_var, popart_count, reward_component_ema_kernel, + h_s2_rms_ema_kernel, q_mean_reduce_kernel, q_mean_subtract_kernel, q_mean_scratch_pinned, diff --git a/crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu b/crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu new file mode 100644 index 000000000..febc0c666 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu @@ -0,0 +1,69 @@ +/* h_s2_rms_ema — GPU-driven RMS(save_h_s2) EMA into ISV[96]. + * + * Plan 4 Task 2c.3c.5. GPU-drives-CPU-reads (spec §4.C.6) and + * pearl_cold_path_no_exception_to_gpu_drives.md compliant. + * + * Reduces RMS = sqrt(sum_sq / N) over the flattened `save_h_s2 [B, SH2]` + * buffer (the online trunk's post-GRN activation) and EMA-updates + * ISV[H_S2_RMS_EMA_INDEX] with rate `ema_alpha`. + * + * Single-block kernel (256 threads, shmem-reduction). No atomicAdd + * (per `feedback_no_atomicadd.md`); the in-block tree reduction is + * standard `s >>= 1` with `__syncthreads()`. Only thread 0 writes the + * EMA result. + * + * Producer-only — no consumer reads ISV[96] in this commit. 2c.3c.6 + * wires the consumer in `mag_concat_qdir`'s adaptive-scale path so the + * residual stack is normalised by the trunk-output RMS regardless of + * GRN drift across training. + * + * Cost: one launch per training step (cold-path cadence — same site as + * `reward_component_ema`). 256 threads × ceil(B*SH2/256) strided loads + * + log2(256)=8 reduce steps + 1 global write. + */ + +extern "C" __global__ void h_s2_rms_ema_update( + const float* __restrict__ h_s2, /* [B, SH2] row-major */ + int B, + int SH2, + float* __restrict__ isv, /* ISV bus [ISV_TOTAL_DIM] */ + int isv_h_s2_rms_ema_index, /* H_S2_RMS_EMA_INDEX = 96 */ + float ema_alpha /* per-batch EMA rate (0,1) */ +) { + /* Single-block contract — match reward_component_ema_kernel's grid guard. */ + if (blockIdx.x != 0) return; + + extern __shared__ float smem[]; + const int tid = (int)threadIdx.x; + const int block = (int)blockDim.x; + const int N = B * SH2; + + /* Per-thread sum of squares over a strided slice. */ + float local_sumsq = 0.0f; + for (int i = tid; i < N; i += block) { + const float v = h_s2[i]; + local_sumsq += v * v; + } + smem[tid] = local_sumsq; + __syncthreads(); + + /* In-block tree reduction (no atomicAdd). */ + for (int s = block / 2; s > 0; s >>= 1) { + if (tid < s) smem[tid] += smem[tid + s]; + __syncthreads(); + } + + /* Thread 0 writes the EMA. The host caller's trainer config requires + * batch_size > 0 and shared_h2 > 0, so N = B*SH2 > 0 is invariant. + * No defensive `(N > 0) ? N : 1` ternary here — masking the divide + * would silently EMA toward 0 if the invariant ever broke; instead the + * NaN that division-by-zero produces propagates to ISV[96] and + * downstream consumers (mag_concat_qdir in 2c.3c.6) where it is + * loudly visible. */ + if (tid == 0) { + const float rms = sqrtf(smem[0] / (float)N); + const int slot = isv_h_s2_rms_ema_index; + const float prev = isv[slot]; + isv[slot] = (1.0f - ema_alpha) * prev + ema_alpha * rms; + } +} diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index f6d20a5e9..6a2a92a78 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -273,6 +273,12 @@ impl StateResetRegistry { category: ResetCategory::FoldReset, description: "ISV[MAMBA2_RETENTION_EMA_INDEX=89] — Mamba2 state-transition magnitude EMA, retention proxy; GPU attention_focus_ema_update kernel fills (Plan 4 E.5 Mode A)", }, + // ───── Plan 4 Task 2c.3c.5: trunk-RMS EMA producer ───────────── + RegistryEntry { + name: "isv_h_s2_rms_ema", + category: ResetCategory::FoldReset, + description: "ISV[H_S2_RMS_EMA_INDEX=96] — per-batch RMS(save_h_s2) EMA (α=0.05); GPU h_s2_rms_ema_update kernel fills; producer-only in 2c.3c.5, consumer wired in 2c.3c.6 (mag_concat_qdir adaptive-scale). Cold-start 1.0 (neutral RMS) reapplied at fold boundary", + }, ]; Self { entries } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 68f83714d..c98e11493 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2803,6 +2803,19 @@ impl DQNTrainer { tracing::warn!("C.5 cql_alpha_seed_update launch failed: {e}"); } } + + // Plan 4 Task 2c.3c.5: per-step H_S2 RMS EMA into ISV[96]. + // Single-block 256-thread shmem-reduction kernel reads the + // online trunk's `save_h_s2 [B, SH2]` (populated by the most + // recent forward) and EMAs sqrt(sum_sq / N) into ISV[96] at + // α=0.05. Producer-only — 2c.3c.6 wires the consumer in + // `mag_concat_qdir`'s adaptive-scale path. Same launch + // cadence as the Plan 3 producers above. + if let Some(ref fused) = self.fused_ctx { + if let Err(e) = fused.trainer().launch_h_s2_rms_ema(ema_alpha) { + tracing::warn!("Plan 4 2c.3c.5 h_s2_rms_ema launch failed: {e}"); + } + } } // B.2 Plan 3 Task 3: freeze TRADE_TARGET_RATE at the configured @@ -4470,6 +4483,18 @@ impl DQNTrainer { fused.trainer().write_isv_signal_at(idx, 0.0); } } + // Plan 4 Task 2c.3c.5: trunk-RMS EMA producer slot. Reset to + // 1.0 (cold-start neutral RMS) at fold boundary so the first + // h_s2_rms_ema_update fire on the new fold EMAs the measured + // RMS toward 1.0 rather than collapsing toward 0. + "isv_h_s2_rms_ema" => { + if let Some(ref fused) = self.fused_ctx { + fused.trainer().write_isv_signal_at( + crate::cuda_pipeline::gpu_dqn_trainer::H_S2_RMS_EMA_INDEX, + 1.0, + ); + } + } "isv_gamma_dir_eff" | "isv_gamma_mag_eff" | "isv_gamma_ord_eff" | "isv_gamma_urg_eff" => { // D.2 per-branch gamma slots. Reset to 0.0 at fold boundary; // per_branch_gamma_update GPU kernel re-populates on next epoch-boundary launch. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index b9f018a2d..b5d00799b 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -291,6 +291,8 @@ Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two help Plan 4 Task 2c.3c.3 (2026-04-24): GRN trunk backward scratch buffer allocation on `GpuDqnTrainer`. 8 new `CudaSlice` device buffers added (4 per GRN block × 2 blocks for h_s1 + h_s2): `bw_grn_h_s2_d_pre_ln` [B, SH2], `bw_grn_h_s2_d_linear_b_out` [B, 2*SH2], `bw_grn_h_s2_d_elu_out` [B, SH2], `bw_grn_h_s2_d_linear_a_out` [B, SH2], plus the four h_s1 mirrors at [B, SH1] (and 2*SH1 for the linear_b_out scratch — GLU's pre-activation has 2× hidden width because value-path + gate_pre are concatenated). The `d_pre_LN` buffer aliases both `d_glu_out` and `d_residual` (pure pass-through under `GrnBlock::backward_raw_phase1`'s pointer aliasing convention from 2c.3c.1); `d_linear_b_out` carries the phase1 → caller's Linear_b backward GEMM gradient; `d_elu_out` carries Linear_b backward's output into phase2; `d_linear_a_out` carries phase2's output to the caller's Linear_a backward GEMM. Allocated alongside the existing `alloc_backward_scratch` call in the trainer constructor (mirrors that helper's `stream.alloc_zeros::(size).map_err(...)?` pattern). Total footprint at SH1=SH2=256, B=512: 8 × 5 × 256 × 4B = 5.0 MB per fold (negligible against existing ~32 MB per-branch cuBLAS workspaces). Each field is `#[allow(dead_code)]` until 2c.3c.4's `apply_grn_trunk_backward_raw` helper consumes them — the comment `wired by Task 2c.3c.4` annotates each suppression. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4. Adam state for the new GRN tensors is auto-allocated since `m_buf`/`v_buf` are sized to `total_params + cutlass_tile_pad` which already includes the 13 GRN tensors after 2c.3a. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +57 LOC. No new module / kernel / ISV slot / Orphan row. +Plan 4 Task 2c.3c.5 (2026-04-25): producer-only ISV slot append for the trunk-RMS adaptive scale. New ISV slot `H_S2_RMS_EMA_INDEX=96` tail-appended between the prior tail-data block and the layout-fingerprint pair; `ISV_LAYOUT_FINGERPRINT_LO_INDEX` shifted 94→97, `ISV_LAYOUT_FINGERPRINT_HI_INDEX` shifted 95→98, `ISV_TOTAL_DIM` 96→99. New CUDA kernel `h_s2_rms_ema_kernel.cu` registered in `build.rs::kernels_with_common` (kernel count 80 → 81): single-block 256-thread shmem-tree reduction (no atomicAdd per `feedback_no_atomicadd.md`) computing `RMS = sqrt(sum_sq / (B*SH2))` over the trainer's `save_h_s2 [B, SH2]` buffer (online trunk's post-GRN activation), then EMA-updating ISV[96] with α=0.05 (≈13-batch half-life). Cubin loaded in `GpuDqnTrainer::new` and stored in a new `h_s2_rms_ema_kernel: CudaFunction` field; `pub fn launch_h_s2_rms_ema(&self, ema_alpha)` mirrors `launch_reward_component_ema`'s style and reads `save_h_s2.raw_ptr()` + `config.shared_h2` + `isv_signals_dev_ptr` from the trainer (no plumbing through the experience collector — the buffer lives on the trainer). Launched from `training_loop.rs` immediately after the existing per-step ISV producer block (`reward_component_ema` + `trade_attempt_rate_ema` + `plan_threshold_update` + `seed_step_counter` + `cql_alpha_seed_update`), at the same launch cadence — once per `collect_experiences_gpu` epoch. Constructor cold-start writes ISV[96]=1.0 (neutral RMS) so the first kernel fire EMAs measured RMS toward 1.0 rather than collapsing toward 0; `StateResetRegistry::isv_h_s2_rms_ema` registered as `FoldReset` with the same 1.0 reapplication at fold boundary (dispatch arm added to `training_loop.rs::reset_named_state`). `layout_fingerprint_seed()` extended with `H_S2_RMS_EMA=96;ISV_LAYOUT_FINGERPRINT_LO=97;ISV_LAYOUT_FINGERPRINT_HI=98;ISV_TOTAL_DIM=99;` — new `LAYOUT_FINGERPRINT_CURRENT = 0x3e21acecd922e540` (was `0xcf3a24b0a1f70057`). **Producer-only — ZERO consumers in this commit**. 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path so the magnitude-branch decoder's residual stack is normalised by the trunk-output RMS regardless of GRN drift across training. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 81 cubins (was 80). +109 / -10 LOC across `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (slot constants + docstring + CUBIN ref + field + load + launcher + cold-start + fingerprint seed entry), `crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu` (new file, 60 LOC), `crates/ml/build.rs` (1 entry), `crates/ml/src/trainers/dqn/state_reset_registry.rs` (1 entry), `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (per-step launch + FoldReset dispatch arm). Smoke deferred — 2c.3c.5 is producer-only with zero consumers, so runtime behaviour is structurally unchanged from 2c.3c.4 (which passed `multi_fold_convergence` in 599.73s with all 3 fold checkpoints written, train Sharpes 7.52 / 60.94 / 10.40 and val_Sharpes 0.51 / 0.64 / 0.82). The new producer's only runtime effects are (a) writing ISV[96] each step — read by nobody yet — and (b) bumping `LAYOUT_FINGERPRINT_CURRENT`, with no checkpoint-load path exercised by the smoke (it starts from scratch). cargo check + cargo build (with the new `h_s2_rms_ema_kernel.cubin` compiling clean under nvcc sm_80) is the standing validation; the meaningful `multi_fold_convergence` re-run lands with 2c.3c.6 when the consumer wires up in `mag_concat_qdir`'s adaptive-scale path. No new Orphan row — the kernel is wired to a real launch in this commit (cold-start + per-step EMA), classified Wired-Producer-Only until 2c.3c.6 promotes it to Wired-Producer+Consumer. 0 panic gates added/removed. + Plan 4 Task 2c.3c.4 followup (2026-04-25): stub-return defensive guards replaced with proper invariants + signal escalation in `gpu_dqn_trainer.rs`. Three classes of cleanup, prompted by the pre-commit stub-return heuristic flagging code paths exposed by the 2c.3c.4 commit's surrounding edits. (1) **Unreachable null/bounds guards → `debug_assert!` + deletion.** `read_isv_signal_at`, `read_atom_utilization`, `compute_q_spectral_gap` each guarded `is_null()` on `isv_signals_pinned` / `q_readback_pinned`, but both pointers are unconditionally allocated by the constructor (`cuMemAllocHost_v2` + `assert_eq!`) and never reassigned after construction outside Drop — the guards masked the contract instead of enforcing it. Same for `read_isv_signal_at`'s `index >= ISV_TOTAL_DIM` early-return (every caller passes a named ISV-slot constant) and `compute_q_spectral_gap`'s `n_cols == 0` early-return (`total_actions()` is the sum of branch sizes which config requires non-zero). All four converted to `debug_assert!(...)` so dev/test catches invariant violations loudly; release builds inherit the unsafe-deref crash on the broken precondition rather than a silent stub return. (2) **Silent training-instability mask → warn-log + collapse sentinel.** `compute_q_spectral_gap` previously returned 1.0 (= "balanced/healthy" per the downstream `NormalizedComponents::from_raw` calibration) on the first non-finite Q-value, masking NaN/Inf gradient corruption with the same value a healthy network produces. New behaviour: scan all samples, flag non-finite, then emit `tracing::warn!` and return 100.0 (the same collapse sentinel emitted when `lambda_2 < 1e-12` or `sigma2 < 1e-6`). Operators see the failure in logs alongside HEALTH_DIAG rather than chasing a phantom balanced spectral gap. (3) **Genuine domain encodings annotated** with `// ok: domain encoding` + explanatory comment: `lambda_2 < 1e-12 → 100.0` and `sigma2 < 1e-6 → 100.0` (rank-1 collapse — sigma_1/sigma_2 diverges, calibrated collapse sentinel), `power_iteration_largest`'s `norm < 1e-20 → 0.0` (M·v vanishes ⇒ eigenvalue is zero). `power_iteration_largest`'s `n == 0 → 0.0` was unreachable from the spectral_gap caller; replaced with `debug_assert!(n > 0)` at the helper entry. No behaviour change in production runs: the unreachable guards weren't firing and the new invariants match the constructor's contract; the non-finite Q-value path strictly improves observability. cargo check clean at 11 warnings (baseline preserved). +57 / -14 LOC, all in `gpu_dqn_trainer.rs`. No new module / kernel / ISV slot / Orphan row. Tiny followup tightens the remaining `.unwrap()` in `compute_q_spectral_gap` to `.expect("q_sample_history just received push_back; back() cannot be None")` documenting the invariant that the immediately-prior `push_back` makes the `back()` return Some-by-construction. Plan 4 Task 2c.3c.4 (2026-04-25): GRN backward chain wired through all 3 panic-gated trunk-backward sites — smoke validates. Three panic gates removed: `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. New helper `BatchedForward::encoder_backward_chain` orchestrates the full GRN trunk backward composition (h_s2 → h_s1) symmetrically with `encoder_forward_only`: `backward_raw_phase1` (LN_dx + LN_dgamma/dbeta no-atomicAdd reduction + GLU_bwd) → cuBLAS Linear_b dW + dX → `backward_raw_phase2` (ELU_bwd) → cuBLAS Linear_a dW + dX → for h_s1 only: Linear_residual dW (no_bias_lda) + dX accumulation when bottleneck active → for h_s2 only: `saxpy_inplace(alpha=1.0)` accumulates `d_pre_ln_h_s2` into `d_h_s1` (identity residual gradient). Used by 4 callers writing to 3 different gradient targets: main backward (writes to `grad_buf` for the 13 GRN trunk tensors at indices [0..13)), CQL backward (writes to `cql_grad_scratch` — same layout as grad_buf), and the two auxiliary paths (`apply_iqn_trunk_gradient` and `apply_ensemble_diversity_backward` write to `iqn_trunk_m` then SAXPY into `grad_buf` with their respective scale factors). `iqn_trunk_m` grown from legacy 4-tensor element count (`sh1*sd + sh1 + sh2*sh1 + sh2`) to `padded_byte_offset(¶m_sizes, 13) / sizeof(f32)` so the layout matches `grad_buf`'s padded byte offsets exactly for the trunk portion — element-wise SAXPY across that span lands every per-tensor gradient at the correct location without needing per-tensor offset computation. `apply_ensemble_diversity_backward` additionally fixes value-head weight indices: `w_v1` 4→13, `w_v2` 6→15 (legacy indices were stale post-2c.3a — caught by the panic gate before they could write garbage). `BatchedBackward::launch_dw_only_no_bias_lda` added because Linear_residual (h_s1, no bias) needs the padded states stride that plain `launch_dw_only_no_bias` doesn't accept; `launch_dw_only` would dereference NULL inside the bias-grad kernel. ReLU masks on `bw_d_h_s2` and `bw_d_h_s1` are gone — the GRN trunk ends in LayerNorm which has no truncation derivative; the value-head FC retains its ReLU mask (value head still uses ReLU for its hidden activation). `launch_cublas_backward_to` migrated from `&self` to `&mut self` because the GRN backward needs to scribble per-block partial-reduction scratch inside `grn_h_s{1,2}_online`; the borrow split lets the trainer hand `&mut BatchedForward` + `&BatchedBackward` to `encoder_backward_chain` simultaneously. Smoke (`cargo test … multi_fold_convergence --ignored`, 599.73s, 3 folds × 5 epochs): all 3 `dqn_fold{N}_best.safetensors` checkpoints written; per-fold best train Sharpe 7.52 / 60.94 / 10.40 at epochs 2 / 4 / 2; per-fold val_Sharpe peaks 0.51 / 0.64 / 0.82; fold 2 epoch 5 hit PF=5.29 with +988% return — gradients flow cleanly through the GRN composition (no NaN, no Inf, no policy collapse, no fold-boundary explosion). Audit row #11 (IQN target trunk orphan) was already retired by 2c.3b; this commit retires the 3 backward panic gates that 2c.3a introduced. cargo check clean at 11 warnings (baseline preserved). +582 / -320 LOC across `batched_backward.rs` (+154/-154 — panic-gate removal + `launch_dw_only_no_bias_lda`), `batched_forward.rs` (+284/-0 — `encoder_backward_chain`), and `gpu_dqn_trainer.rs` (+221/-243 — 3 trunk-backward call sites rewritten + `iqn_trunk_m` grown + 4 main/CQL/IQN/ensemble backward call sites updated). No new module / kernel / ISV slot. Closes 3 of the 6 panic gates introduced by 2c.3a (3 forward gates were retired by 2c.3b); panic-gate count: 6 → 0.