From 6c3a91c8781525bb4ba4f8e311abd75a5df6685d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 3 May 2026 14:55:40 +0200 Subject: [PATCH] fix(sp7): wire raw CQL norm kernel reading cql_grad_scratch directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SP7 controller's CQL reference signal was reading cql_sx (post-SAXPY, budget-scaled) which created a self-perpetuating deadlock at small budget values: cql_sx_norm = budget × raw_grad → small budget → small cql_sx → controller can't update → budget stays small. The earlier offset 6 → 3 attempt failed because grad_decomp_launch_cql was never effectively populating slot 3 — the snapshot pattern measures ‖grad_buf − snapshot‖, but apply_cql_gradient writes to cql_grad_scratch (separate buffer), not grad_buf, so the snapshot delta is always 0. This commit adds a real producer kernel cql_raw_norm_compute that reads cql_grad_scratch directly and computes ‖raw_cql‖ over mag/dir/trunk slices. Wired to fire AFTER apply_cql_gradient and BEFORE apply_cql_saxpy, populating grad_decomp_result_pinned[3..6] with the raw norm independent of cql_budget. SP7 launcher updated to read from offset 3 (now: real raw CQL norm, not the never-populated cql snapshot delta). HEALTH_DIAG label renamed cql_sx → cql_raw to reflect the new contract; component index in the cached grad_component_norms_* arrays switched 2 → 1. The historical grad_decomp_launch_cql() call is removed — keeping it would overwrite the slot with 0 after cql_raw_norm_compute fires. The paired grad_decomp_snapshot_cql snapshot is left in place to scope the diff to Path A; buffer cleanup (grad_snapshot_cql allocation + grad_decomp_launch_cql definition) belongs in a follow-up commit per feedback_no_partial_refactor. Files: cql_raw_norm_kernel.cu (new, 97 LOC), build.rs, gpu_dqn_trainer.rs (struct field + cubin static + load + launcher + SP7 read offset 6→3), loss_balance_controller_kernel.cu (docstring + arg comment), fused_training.rs (4 launch_cql_raw_norm call sites, 1 dead grad_decomp_launch_cql call removed), training_loop.rs (HEALTH_DIAG label + index), audit doc Fix 31 SP7 Path A entry. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 6 ++ .../src/cuda_pipeline/cql_raw_norm_kernel.cu | 97 +++++++++++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 90 ++++++++++++++++- .../loss_balance_controller_kernel.cu | 18 +++- crates/ml/src/trainers/dqn/fused_training.rs | 60 ++++++++++-- .../src/trainers/dqn/trainer/training_loop.rs | 22 +++-- docs/dqn-wire-up-audit.md | 54 +++++++++++ 7 files changed, 322 insertions(+), 25 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 10bd72314..1ed09c9b1 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -101,6 +101,12 @@ fn main() { "mamba2_temporal_kernel.cu", "graph_utility_kernels.cu", "grad_decomp_kernel.cu", + // SP7 Path A (2026-05-03): raw CQL norm reading `cql_grad_scratch` + // directly so the SP7 controller has a budget-independent CQL + // reference signal. Replaces the never-populated `cql` slot at + // element offset 3 in `grad_decomp_result_pinned`. See + // `cql_raw_norm_kernel.cu` and audit doc Fix 31 SP7 Path A. + "cql_raw_norm_kernel.cu", "branch_grad_balance_kernel.cu", "backtest_plan_kernel.cu", "tau_update_kernel.cu", diff --git a/crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu b/crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu new file mode 100644 index 000000000..a7e82c545 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/cql_raw_norm_kernel.cu @@ -0,0 +1,97 @@ +/** + * cql_raw_norm_kernel — SP7 Path A (2026-05-03) raw CQL gradient norm. + * + * Computes L2 norms of `cql_grad_scratch` over the same trunk/dir/mag + * slices `grad_decomp_kernel` uses, and writes three f32 results to a + * 3-float view inside the shared `grad_decomp_result_pinned` buffer: + * result_out[0] = ‖cql_scratch‖ on the magnitude branch slice + * result_out[1] = ‖cql_scratch‖ on the direction branch slice + * result_out[2] = ‖cql_scratch‖ on the trunk slice + * + * Why this exists. The SP7 loss-balance controller needs a CQL reference + * signal that scales with the *raw* CQL gradient magnitude — independent + * of the per-branch `cql_budget` the consumer applies via + * `apply_cql_saxpy`. Reading the post-SAXPY delta (`cql_sx` slot) + * creates a self-perpetuating deadlock: cql_sx_norm = budget × raw_grad, + * so when budget bootstraps small (≈0.02) the controller sees ≈0, + * cold-starts every step, and never updates the budget. + * + * The earlier offset-3 (`cql`) attempt failed because + * `grad_decomp_launch_cql` measures `‖grad_buf − snapshot‖`, but + * `apply_cql_gradient` writes to `cql_grad_scratch` (a separate buffer); + * the snapshot delta on `grad_buf` is therefore always 0. + * + * This kernel reads `cql_grad_scratch` *directly* (no snapshot pair), + * so the norm reflects the raw gradient magnitude before budget scaling + * — exactly what the SP7 controller needs as a reference signal. + * + * Layout: + * cql_grad_scratch — f32 device pointer, length total_params (matches + * `grad_buf`'s element layout). Same trunk_start/ + * dir_start/mag_start indices used by the + * grad_decomp pipeline apply unmodified. + * trunk_start/len, dir_start/len, mag_start/len + * — element offsets + lengths into `cql_grad_scratch` + * for each slice (identical to grad_decomp_kernel + * contract). + * result_out — [3] output: (mag_norm, dir_norm, trunk_norm). The + * caller computes the destination pointer as + * `grad_decomp_result_dev_ptr + 3 * sizeof(f32)`, + * replacing the never-populated `cql` slot at + * element offset 3 (was always zero — see audit + * doc Fix 31 SP7 Path A). + * + * All reductions are block-internal via shared-memory tree reduction — + * NO atomicAdd (`feedback_no_atomicadd`). + * + * Launch config: one block, 256 threads. + */ +#include + +extern "C" __global__ void cql_raw_norm_compute( + const float* __restrict__ cql_grad_scratch, + int trunk_start, int trunk_len, + int dir_start, int dir_len, + int mag_start, int mag_len, + float* __restrict__ result_out +) { + __shared__ float sum_mag[256]; + __shared__ float sum_dir[256]; + __shared__ float sum_trunk[256]; + int tid = threadIdx.x; + sum_mag[tid] = 0.0f; + sum_dir[tid] = 0.0f; + sum_trunk[tid] = 0.0f; + + // Trunk slice — read cql_grad_scratch[trunk_start..+trunk_len) directly. + for (int i = tid; i < trunk_len; i += blockDim.x) { + float v = cql_grad_scratch[trunk_start + i]; + sum_trunk[tid] += v * v; + } + // Direction slice — cql_grad_scratch[dir_start..+dir_len). + for (int i = tid; i < dir_len; i += blockDim.x) { + float v = cql_grad_scratch[dir_start + i]; + sum_dir[tid] += v * v; + } + // Magnitude slice — cql_grad_scratch[mag_start..+mag_len). + for (int i = tid; i < mag_len; i += blockDim.x) { + float v = cql_grad_scratch[mag_start + i]; + sum_mag[tid] += v * v; + } + __syncthreads(); + + // Block-level tree reduction (256 → 1). + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sum_mag[tid] += sum_mag[tid + s]; + sum_dir[tid] += sum_dir[tid + s]; + sum_trunk[tid] += sum_trunk[tid + s]; + } + __syncthreads(); + } + if (tid == 0) { + result_out[0] = sqrtf(sum_mag[0]); + result_out[1] = sqrtf(sum_dir[0]); + result_out[2] = sqrtf(sum_trunk[0]); + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index b61844e31..2b4a7ffa3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -91,6 +91,15 @@ static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_gra static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin")); pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin")); static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_decomp_kernel.cubin")); +/// SP7 Path A (2026-05-03): raw CQL gradient norm reading `cql_grad_scratch` +/// directly. Populates `grad_decomp_result_pinned[3..6]` (the previously +/// never-populated `cql` slot at element offset 3) with `‖raw_cql‖` over the +/// same trunk/dir/mag slices `grad_decomp_kernel` uses. The SP7 +/// loss-balance controller reads this slot as a budget-independent CQL +/// reference; the prior `cql_sx` reference at offset 6 was budget-scaled +/// (cql_sx_norm = budget × raw_grad), creating a self-perpetuating +/// deadlock at small bootstrap budgets. +static CQL_RAW_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_raw_norm_kernel.cubin")); static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/branch_grad_balance_kernel.cubin")); /// Plan 1 Task 13: GPU-driven Polyak-EMA tau coefficient. static TAU_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tau_update_kernel.cubin")); @@ -3220,6 +3229,13 @@ pub struct GpuDqnTrainer { grad_decomp_snapshot_len: usize, /// Reduction kernel (one-block 256-thread launch, no atomics). grad_decomp_kernel: CudaFunction, + /// SP7 Path A (2026-05-03) raw CQL norm kernel. Reads + /// `cql_grad_scratch` directly (no snapshot pair) and writes 3 floats + /// to `grad_decomp_result_pinned[3..6]` — the `cql` slot at element + /// offset 3 that was never populated by the (defined-but-unused) + /// `grad_decomp_launch_cql` snapshot pattern. Single block, 256 + /// threads, shared-memory tree reduction (no atomicAdd). + cql_raw_norm_kernel: CudaFunction, /// Adaptive per-branch gradient-norm balancer — caps any branch whose /// L2 norm exceeds `num_branches × median_branch_norm` (see /// `branch_grad_balance_kernel.cu`). `num_branches = 4` is the @@ -5867,6 +5883,52 @@ impl GpuDqnTrainer { self.launch_grad_decomp(&self.grad_snapshot_ens, 24, "ens") } + /// SP7 Path A (2026-05-03) — launch `cql_raw_norm_compute` reading + /// `cql_grad_scratch` directly. Writes `‖raw_cql‖` over (mag, dir, + /// trunk) slices to `grad_decomp_result_pinned[3..6]` — the `cql` + /// slot at element offset 3. + /// + /// Why a dedicated launcher rather than reusing `launch_grad_decomp` + /// with the `grad_decomp_launch_cql` snapshot: that pattern produces + /// `‖grad_buf − snapshot_cql‖`, which is structurally always 0 + /// because `apply_cql_gradient` writes to `cql_grad_scratch` (a + /// separate buffer), never to `grad_buf`. The previously-defined + /// `grad_decomp_launch_cql` is unwired (no call site); this launcher + /// supplants it with a real producer that the SP7 controller can + /// read at element offset 3 as a budget-independent reference signal. + /// + /// Must run AFTER `apply_cql_gradient` populates `cql_grad_scratch` + /// and BEFORE `apply_cql_saxpy` consumes it (the saxpy doesn't + /// modify `cql_grad_scratch`, but ordering keeps the contract + /// "raw norm reflects what gets SAXPYed this step" obvious). + pub(crate) fn launch_cql_raw_norm(&self) -> Result<(), MLError> { + let scratch_ptr = self.cql_grad_scratch.raw_ptr(); + let f32_size = std::mem::size_of::() as u64; + // Offset 3 (cql slot) — was never populated by the unwired + // `grad_decomp_launch_cql` snapshot pattern. Now produces the + // raw-CQL-norm SP7 needs. + let result_ptr = self.grad_decomp_result_dev_ptr + 3 * f32_size; + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + self.stream.launch_builder(&self.cql_raw_norm_kernel) + .arg(&scratch_ptr) + .arg(&self.grad_decomp_trunk_start) + .arg(&self.grad_decomp_trunk_len) + .arg(&self.grad_decomp_dir_start) + .arg(&self.grad_decomp_dir_len) + .arg(&self.grad_decomp_mag_start) + .arg(&self.grad_decomp_mag_len) + .arg(&result_ptr) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("cql_raw_norm_compute: {e}")))?; + } + Ok(()) + } + /// Task 2.0 — populate the host-side `grad_component_norms_mag/_dir/ /// _trunk` caches from the pinned result slot. Called at epoch /// boundary, before HEALTH_DIAG emission. @@ -11348,11 +11410,20 @@ impl GpuDqnTrainer { // grad_decomp pinned layout: 27 floats total, 9 components × 3 floats // each ([mag, dir, trunk]). Component element offsets per // launch_grad_decomp documentation: iqn=0, cql=3, cql_sx=6, c51=9. - // We use cql_sx (post-budget delta) for ratio parity with what landed - // in grad_buf. + // + // SP7 Path A (2026-05-03): the CQL reference now reads slot 3 (`cql`) + // populated by `launch_cql_raw_norm` — a real producer that reads + // `cql_grad_scratch` directly. Previously SP7 read slot 6 (`cql_sx`, + // post-budget delta = budget × raw_grad), which created a self- + // perpetuating deadlock at small bootstrap budgets: the controller + // saw ≈0 norm → cold-started every step → never updated the budget. + // Slot 3 is budget-independent — exactly the reference SP7 needs. + // The historical "cql at offset 3" was zero-valued (`grad_decomp_launch_cql` + // measured a delta on `grad_buf`, but `apply_cql_gradient` writes + // to `cql_grad_scratch`); the unwired path is now fully replaced. let f32_size = std::mem::size_of::() as u64; let iqn_dev = self.grad_decomp_result_dev_ptr + 0 * f32_size; - let cql_sx_dev = self.grad_decomp_result_dev_ptr + 6 * f32_size; + let cql_dev = self.grad_decomp_result_dev_ptr + 3 * f32_size; let c51_dev = self.grad_decomp_result_dev_ptr + 9 * f32_size; // Step 1: producer kernel. @@ -11379,7 +11450,7 @@ impl GpuDqnTrainer { self.stream .launch_builder(&self.loss_balance_controller_kernel) .arg(&iqn_dev) - .arg(&cql_sx_dev) + .arg(&cql_dev) .arg(&c51_dev) .arg(&isv_dev) .arg(&flatness_isv_base_i32) @@ -14067,6 +14138,16 @@ impl GpuDqnTrainer { module.load_function("grad_component_delta_norm") .map_err(|e| MLError::ModelError(format!("grad_component_delta_norm load: {e}")))? }; + // SP7 Path A (2026-05-03) — raw CQL norm kernel. Same shape as + // grad_decomp_kernel but reads `cql_grad_scratch` directly (no + // snapshot pair). Loaded into its own CUmodule for the same + // CUfunction-isolation reason. + let cql_raw_norm_kernel = { + let module = stream.context().load_cubin(CQL_RAW_NORM_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("cql_raw_norm cubin load: {e}")))?; + module.load_function("cql_raw_norm_compute") + .map_err(|e| MLError::ModelError(format!("cql_raw_norm_compute load: {e}")))? + }; info!( "GpuDqnTrainer: grad_decomp kernel loaded — snapshot_len={} trunk=[{}..+{}] dir=[{}..+{}] mag=[{}..+{}]", grad_decomp_snapshot_len, @@ -17391,6 +17472,7 @@ impl GpuDqnTrainer { grad_decomp_mag_len, grad_decomp_snapshot_len, grad_decomp_kernel, + cql_raw_norm_kernel, branch_grad_balance_reduce, branch_grad_balance_isv_update, branch_grad_balance_rescale, diff --git a/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu b/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu index c49ac7a15..fda09bc1c 100644 --- a/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu +++ b/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu @@ -3,7 +3,16 @@ // SP7 (2026-05-03): per-branch loss-balance controller for CQL and C51 budgets. // // Reads per-loss decomp pinned slots [mag_norm, dir_norm, trunk_norm] for -// IQN (reference), CQL_SX (managed, post-budget), C51 (managed, pre-budget). +// IQN (reference), CQL (managed, RAW pre-budget — Path A 2026-05-03), +// C51 (managed, pre-budget). +// +// Path A note: prior to 2026-05-03 the kernel read CQL_SX (post-SAXPY, +// budget-scaled = budget × raw_grad) at element offset 6, which created +// a self-perpetuating deadlock at small bootstrap budgets — the controller +// saw ≈0 norm → cold-started → never updated the budget. The new +// `cql_raw_norm_compute` kernel populates element offset 3 (the `cql` slot) +// directly from `cql_grad_scratch`, giving SP7 a budget-independent CQL +// reference. Launcher updated atomically to read offset 3. // Reads per-branch flatness from ISV[FLATNESS_BASE..+4]. Reads prior // budgets from ISV[BUDGET_{CQL,C51}_BASE..+4] and prior Wiener state from // ISV[LB_*_VAR_{CQL,C51}_BASE..+4]. Writes new budgets + new Wiener @@ -62,9 +71,10 @@ extern "C" __global__ void loss_balance_controller_update( // Pointers into the shared `grad_decomp_result_pinned` buffer at the // 3-float component offsets (the launcher computes these by adding // (offset_elems * sizeof(float)) to grad_decomp_result_dev_ptr): - // iqn_decomp = base + 0 (slot iqn, layout [mag, dir, trunk]) - // cql_decomp = base + 24 (slot cql_sx, layout [mag, dir, trunk]) - // c51_decomp = base + 36 (slot c51, layout [mag, dir, trunk]) + // iqn_decomp = base + 0 (slot iqn, layout [mag, dir, trunk]) + // cql_decomp = base + 12 (slot cql, Path A 2026-05-03 — raw norm, + // layout [mag, dir, trunk]) + // c51_decomp = base + 36 (slot c51, layout [mag, dir, trunk]) const float* __restrict__ iqn_decomp, const float* __restrict__ cql_decomp, const float* __restrict__ c51_decomp, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2aef0d4a9..c230b2122 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2381,16 +2381,40 @@ impl FusedTrainingCtx { let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0; let f5_barrier_weight = 0.20_f32 * temporal_amp; - // Task 2.0 — snapshot BEFORE the CQL gradient path. The snapshot runs - // whether or not `has_cql()` fires so the post-CQL reduction below - // can run unconditionally (keeps the pinned result slot's CQL field - // populated every step — 0.0 on steps where CQL was skipped). + // SP7 Path A (2026-05-03): the historical `grad_decomp_launch_cql` + // (paired with this snapshot) measured `‖grad_buf − snapshot_cql‖` + // over the CQL gradient path — structurally always 0 because + // `apply_cql_gradient` writes into `cql_grad_scratch` (a separate + // buffer), never touching `grad_buf` until `apply_cql_saxpy` fires. + // The launch is replaced below by `launch_cql_raw_norm`, which reads + // `cql_grad_scratch` directly and writes a real raw-CQL norm to + // `grad_decomp_result_pinned[3..6]`. The SP7 loss-balance controller + // reads slot 3 as a budget-independent CQL reference signal (was + // deadlocking on the budget-scaled `cql_sx` slot at offset 6). + // + // The snapshot itself is now functionally dead (nothing reads its + // output) but is left in place to preserve the `grad_snapshot_cql` + // buffer + `grad_decomp_snapshot_cql()` call site semantics — + // removing the buffer touches the constructor, struct definition, + // and 9 sibling `grad_snapshot_*` references; that cleanup belongs + // in a follow-up commit per `feedback_no_partial_refactor` (this + // commit is focused on Path A wire-up). self.trainer.grad_decomp_snapshot_cql() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql: {e}"))?; if self.trainer.has_cql() { match self.trainer.apply_cql_gradient(f5_barrier_weight) { Ok(true) => { + // SP7 Path A (2026-05-03) — populate `cql` slot at offset 3 + // with the raw CQL gradient norm read directly from + // `cql_grad_scratch`. Must run AFTER `apply_cql_gradient` + // populated the scratch and BEFORE `apply_cql_saxpy` + // consumes it. Replaces the unwired + // `grad_decomp_launch_cql` snapshot path (which always + // reported 0 because the snapshot was on `grad_buf` and + // CQL writes to `cql_grad_scratch`). + self.trainer.launch_cql_raw_norm() + .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm: {e}"))?; // Task 2.0 extension — isolate `apply_cql_saxpy`. The // preceding `apply_cql_gradient` only writes into // `cql_grad_scratch`, not `grad_buf`, so the cql_sx @@ -2413,6 +2437,13 @@ impl FusedTrainingCtx { } Ok(false) => { // Keep cql_sx slot populated on skip path — delta = 0. + // SP7 Path A: also keep `cql` slot at offset 3 populated. + // `cql_grad_scratch` was zeroed at the top of + // `apply_cql_gradient` (`cuMemsetD32Async`) before the + // early-return, so the raw-norm kernel reads all-zeros + // and writes 0.0 to the slot — honest semantic. + self.trainer.launch_cql_raw_norm() + .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (skip): {e}"))?; self.trainer.grad_decomp_snapshot_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (skip): {e}"))?; self.trainer.grad_decomp_launch_cql_sx() @@ -2420,7 +2451,12 @@ impl FusedTrainingCtx { } Err(e) => { tracing::warn!("CQL gradient failed (non-fatal): {e}"); - // Keep cql_sx slot populated on error path — delta = 0. + // Keep both slots populated on error path. The scratch + // may be in an indeterminate state after the partial + // backward; the raw-norm kernel still produces a + // numerically-finite L2 norm, which is what SP7 needs. + self.trainer.launch_cql_raw_norm() + .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (err): {e}"))?; self.trainer.grad_decomp_snapshot_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (err): {e}"))?; self.trainer.grad_decomp_launch_cql_sx() @@ -2429,16 +2465,22 @@ impl FusedTrainingCtx { } } else { // Keep cql_sx slot populated on CQL-disabled path — delta = 0. + // SP7 Path A: scratch is left at its prior contents (no + // `apply_cql_gradient` call to zero it), so the raw norm reads + // whatever survived from the last step where CQL fired. In the + // CQL-permanently-disabled regime the buffer is still zero from + // construction, so the kernel writes 0.0. In a transient-disable + // regime the slot stays at the most-recent active value, which + // is the appropriate "carry forward" semantic for a controller + // reading a gradient-magnitude reference. + self.trainer.launch_cql_raw_norm() + .map_err(|e| anyhow::anyhow!("SP7 Path A cql_raw_norm (no-cql): {e}"))?; self.trainer.grad_decomp_snapshot_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx (no-cql): {e}"))?; self.trainer.grad_decomp_launch_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx (no-cql): {e}"))?; } - // Task 2.0 — CQL reduction. Delta = grad_buf − snapshot over branch 0+1. - self.trainer.grad_decomp_launch_cql() - .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp CQL: {e}"))?; - // D1/N1: Temporal self-distillation — per-step kernel pull toward the // best historical snapshot. Fully GPU-native: alpha is computed // in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped), diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index d222b720b..7858bb35c 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -4101,7 +4101,7 @@ impl DQNTrainer { } // SP7 observability (grad_decomp_pinned): surface the exact contents of - // grad_decomp_result_pinned[0..12, 24..36, 36..48] (iqn/cql_sx/c51 × + // 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 // last step of this epoch. The values are already cached in the // DQNTrainer host arrays by `refresh_grad_component_norms` (called ~300 @@ -4109,23 +4109,29 @@ impl DQNTrainer { // zero extra synchronisation. // // Component index mapping (matches `grad_decomp_result_pinned` layout): - // IQN = component 0 → pinned offsets [0, 1, 2] - // CQL_SX = component 2 → pinned offsets [6, 7, 8] - // C51 = component 3 → pinned offsets [9, 10, 11] + // IQN = component 0 → pinned offsets [0, 1, 2] + // CQL_RAW = component 1 → pinned offsets [3, 4, 5] (Path A 2026-05-03) + // C51 = component 3 → pinned offsets [9, 10, 11] + // + // SP7 Path A (2026-05-03): label changed `cql_sx` → `cql_raw` to + // reflect the new producer. `cql_sx` (component 2, post-SAXPY + // budget-scaled delta on `grad_buf`) is no longer the SP7 reference; + // `cql_raw` (component 1, raw norm of `cql_grad_scratch` written by + // `cql_raw_norm_compute`) is — see audit doc Fix 31 SP7 Path A. // // This disambiguates "grad_decomp not populated at SP7 read time" // (all zeros here) from "values non-zero but cold-start branch fires // for another reason" (non-zeros here, lb_active stays 0.0). { let (iqn_m, iqn_d, iqn_t, - cql_sx_m, cql_sx_d, cql_sx_t, + cql_raw_m, cql_raw_d, cql_raw_t, c51_m, c51_d, c51_t) = if let Some(ref fused) = self.fused_ctx { let trainer = fused.trainer(); let mag = trainer.grad_component_norms_mag(); let dir = trainer.grad_component_norms_dir(); let trunk = trainer.grad_component_norms_trunk(); (mag[0], dir[0], trunk[0], // IQN - mag[2], dir[2], trunk[2], // CQL_SX + mag[1], dir[1], trunk[1], // CQL_RAW (Path A — was [2] CQL_SX) mag[3], dir[3], trunk[3]) // C51 } else { (0.0_f32, 0.0_f32, 0.0_f32, @@ -4133,10 +4139,10 @@ impl DQNTrainer { 0.0_f32, 0.0_f32, 0.0_f32) }; tracing::info!( - "HEALTH_DIAG[{}]: grad_decomp_pinned [iqn:[m={:.4} d={:.4} t={:.4}] cql_sx:[m={:.4} d={:.4} t={:.4}] c51:[m={:.4} d={:.4} t={:.4}]]", + "HEALTH_DIAG[{}]: grad_decomp_pinned [iqn:[m={:.4} d={:.4} t={:.4}] cql_raw:[m={:.4} d={:.4} t={:.4}] c51:[m={:.4} d={:.4} t={:.4}]]", epoch, iqn_m, iqn_d, iqn_t, - cql_sx_m, cql_sx_d, cql_sx_t, + cql_raw_m, cql_raw_d, cql_raw_t, c51_m, c51_d, c51_t, ); } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 0bccaf898..3d64bbb45 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -4066,5 +4066,59 @@ extended. No producer kernel yet — that arrives in the next commit. No new ISV slots. No kernel change. No StateResetRegistry entries. Purely additive. Both lines placed adjacent to `q_var_per_branch` in the per-epoch HEALTH_DIAG block (`training_loop.rs`). +- SP7 Path A: raw CQL norm kernel (2026-05-03, wt/sp7-cql-raw-kernel): + fixes a self-perpetuating deadlock in the SP7 controller's CQL reference + signal. Pre-fix: SP7 read `cql_decomp` from `grad_decomp_result_pinned` + at element offset 6 (the `cql_sx` slot, populated by + `grad_decomp_launch_cql_sx`). `cql_sx_norm` measures + `‖grad_buf_after_apply_cql_saxpy − grad_buf_before‖` = + `‖cql_budget × raw_cql_grad‖` = `cql_budget × ‖raw_cql_grad‖`. When + bootstrap `cql_budget ≈ 0.02` (`COLD_START_FLOOR_CQL` = + `CQL_BOOTSTRAP_BUDGET`), `cql_sx_norm ≈ 0`, the controller's + `h_n < EPS_DIV` cold-start guard fires every step, and the budget never + updates from its bootstrap value. Empirically observed on + `smoke-test-kl4lw` (HEAD `237b3dbfb`): all 4 CQL active flags stuck at + bootstrap across 5 epochs of fold 0 despite non-zero + `grad_split_bwd cql=6.48`. + + An earlier offset-3 attempt (read the `cql` slot instead of `cql_sx`) + was abandoned because: (a) `grad_decomp_launch_cql` measures + `‖grad_buf − snapshot_cql‖`, but `apply_cql_gradient` writes only into + `cql_grad_scratch` (a separate buffer) — the delta on `grad_buf` is + always 0; (b) the existing snapshot pattern can't be repointed at + `cql_grad_scratch` because the kernel takes a single `current` pointer + + `snapshot` pointer, designed around the `grad_buf` layout. + + Path A fix: new GPU kernel `cql_raw_norm_compute` reads + `cql_grad_scratch` directly and writes a 3-float `[mag, dir, trunk]` + L2-norm tuple to `grad_decomp_result_pinned[3..6]` — replacing the + always-zero `cql` slot with a real raw-CQL norm. Same 256-thread + single-block shared-memory tree-reduction shape as `grad_decomp_kernel` + (no atomicAdd per `feedback_no_atomicadd`), reusing the trainer's + `grad_decomp_trunk_start/_dir_start/_mag_start` slice indices for + layout parity. Wired AFTER `apply_cql_gradient` populates the scratch + and BEFORE `apply_cql_saxpy` consumes it in `fused_training.rs:2391-2470`, + unconditionally (active/skip/error/no-CQL paths) so the slot stays + populated every step (was the prior contract via the now-removed + `grad_decomp_launch_cql` call). SP7 launcher + (`launch_loss_balance_controller`) updated to read element offset 3 + (`cql_dev`); `loss_balance_controller_kernel.cu` docstring + + `cql_decomp` arg comment updated for the new contract. HEALTH_DIAG + `grad_decomp_pinned` line label renamed `cql_sx` → `cql_raw` and + switched from component index 2 → 1 in the cached + `grad_component_norms_*` arrays. The historical + `grad_decomp_launch_cql()` call is removed (would otherwise overwrite + the slot with 0 after `cql_raw_norm_compute` fires); the paired + `grad_decomp_snapshot_cql` snapshot is left in place to keep the + diff scoped to Path A — buffer cleanup belongs in a follow-up commit. + + Files: `cuda_pipeline/cql_raw_norm_kernel.cu` (new, 91 LOC), + `build.rs` (+5 LOC), `cuda_pipeline/gpu_dqn_trainer.rs` (struct field + + cubin static + load + launcher; +60 LOC), `cuda_pipeline/loss_balance_controller_kernel.cu` + (docstring + arg comment, +10 LOC), `trainers/dqn/fused_training.rs` + (3 new `launch_cql_raw_norm` call sites in active/skip/error/no-CQL + branches + comment block, +30 LOC, −3 LOC removing the dead + `grad_decomp_launch_cql` call), `trainers/dqn/trainer/training_loop.rs` + (HEALTH_DIAG label + index, +6 LOC, −5 LOC), this audit entry. Contract test — every RegistryEntry has dispatch arm (2026-05-03): added `every_fold_and_soft_reset_entry_has_dispatch_arm` unit test to the existing `#[cfg(test)] mod tests` block in `state_reset_registry.rs`. Source-introspection design: `include_str!` both `state_reset_registry.rs` and `trainer/training_loop.rs`; brace-balance walk extracts the `match name {` body; for each FoldReset/SoftReset entry (97 total: 95 FoldReset + 2 SoftReset) asserts the literal `""` appears in that body. No production-code change. No new dependency (manual brace walk instead of regex). Catches the recurring "add RegistryEntry, forget dispatch arm → fold-boundary panic" bug (occurred twice: SP5 Layer A #281, SP7 T7 commit 6e479c55c) at `cargo test -p ml --lib` rather than mid-training. Test currently passes (all 97 dispatch arms present post SP7 T7 fix). Touched: `trainers/dqn/state_reset_registry.rs` (+59 LOC in cfg-test block).