diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 95b3d4b28..864293779 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -3631,11 +3631,25 @@ extern "C" __global__ void hindsight_relabel_kernel( * For each sample b: * 1. Copy h_s2[b, 0..SH2] → concat_out[b, 0..SH2] * 2. Compute E[Q_dir_a] for a=0..b0_size-1 from dir_logits via softmax × support - * 3. Normalize E[Q]/delta_z for matched gradient scale with h_s2 - * 4. Write normalized Q_dir → concat_out[b, SH2..SH2+b0_size] + * 3. Plan 4 Task 2c.3c.6: Q_dir is now adaptively scaled so its per-sample + * RMS over the b0_size action dimension matches the trunk-output RMS + * tracked by ISV[H_S2_RMS_EMA_INDEX=96] (producer kernel + * `h_s2_rms_ema_update` populates the slot each step). The earlier + * `eq / dz` normalisation was calibrated for the pre-GRN post-ReLU + * h_s2 scale and has been replaced. The scale falls back to + * `1 / max(dz, 1e-6)` when `q_rms ≈ 0` (uniform Q across actions — + * the rare degenerate case where RMS-match is undefined). + * 4. Write scaled Q_dir → concat_out[b, SH2..SH2+b0_size] * * Grid: ceil(B/256), Block: 256. One thread per sample. + * + * Static bound: direction branch is exactly 4 actions per the project's + * 4-direction layout (S/H/L/F). The kernel still takes `b0_size` at + * runtime to keep the signature stable, but the per-thread eq stash is + * a 4-element register array — production callers pass b0_size==4. */ +#define MAG_CONCAT_MAX_DIR 4 + extern "C" __global__ void mag_concat_qdir( const float* __restrict__ h_s2, /* [B, SH2] trunk activation (or vsn_masked) */ const float* __restrict__ dir_logits_v, /* [B, NA] value logits (shared head) */ @@ -3645,11 +3659,18 @@ extern "C" __global__ void mag_concat_qdir( int SH2, int NA, /* num_atoms */ int b0_size, /* direction branch size (4) */ - const float* __restrict__ per_sample_support /* [B, 4, 3] stride-12: v_min, v_max, delta_z per-branch (Phase 2d) */ + const float* __restrict__ per_sample_support, /* [B, 4, 3] stride-12: v_min, v_max, delta_z per-branch (Phase 2d) */ + const float* __restrict__ isv, /* Plan 4 Task 2c.3c.6: ISV bus base pointer */ + int isv_h_s2_rms_index /* Plan 4 Task 2c.3c.6: H_S2_RMS_EMA_INDEX = 96 */ ) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; + /* Plan 4 Task 2c.3c.6: read trunk-output RMS once into a register. + * Producer (`h_s2_rms_ema_update`) writes ISV[96] each step; cold-start + * is 1.0 (neutral) so the first batches see a meaningful target. */ + const float h_s2_rms_ema = isv[isv_h_s2_rms_index]; + /* Step 1: copy h_s2 row */ int out_stride = SH2 + b0_size; for (int i = 0; i < SH2; i++) { @@ -3662,6 +3683,8 @@ extern "C" __global__ void mag_concat_qdir( float v_min = per_sample_support[support_base + 0]; float dz = per_sample_support[support_base + 2]; + /* Pass 1: softmax → E[Q] per direction action; stash for the RMS-match pass. */ + float eq_local[MAG_CONCAT_MAX_DIR]; for (int a = 0; a < b0_size; a++) { /* Combined value + advantage logits for direction action a. * Branch logits layout: BRANCH-MAJOR [N × b0 × NA], so action a @@ -3686,9 +3709,32 @@ extern "C" __global__ void mag_concat_qdir( float z_val = v_min + (float)z * dz; eq += prob * z_val; } - /* Normalize by delta_z so Q_dir ~ 1-10 (matches h_s2 post-ReLU scale). - * Without this, Q_dir ~ 0.1 → 10× smaller gradient → 10× slower learning. */ - concat_out[b * out_stride + SH2 + a] = eq / fmaxf(dz, 1e-6f); + eq_local[a] = eq; + } + + /* Pass 2: per-sample RMS of Q_dir over the b0_size action dim. */ + float q_sum_sq = 0.0f; + for (int a = 0; a < b0_size; a++) { + q_sum_sq += eq_local[a] * eq_local[a]; + } + float q_rms = sqrtf(q_sum_sq / (float)b0_size); + + /* Pass 3: adaptive scale — write each eq_a so Q_dir's per-sample RMS + * matches the trunk-output RMS tracked by ISV[H_S2_RMS_EMA_INDEX=96]. + * + * Fallback: when q_rms ≈ 0 (uniform Q across actions ⇒ no spread to + * scale to a target), there is no well-defined RMS-match. Falling + * back to `1 / max(dz, 1e-6)` preserves the legacy 2c.3c.4 scale + * specifically for that degenerate case so the magnitude-branch + * input doesn't collapse to zero. The 1e-6 epsilon mirrors the + * legacy `fmaxf(dz, 1e-6f)` tolerance. + * domain: uniform Q across actions has no RMS to match — fallback is + * mathematically required, not a stub return. */ + const float scale = (q_rms > 1e-6f) + ? (h_s2_rms_ema / q_rms) + : (1.0f / fmaxf(dz, 1e-6f)); + for (int a = 0; a < b0_size; a++) { + concat_out[b * out_stride + SH2 + a] = eq_local[a] * scale; } } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fad550d4f..76feeb8fd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3642,7 +3642,20 @@ impl GpuDqnTrainer { /// Build [source; Q_dir] concat buffer for magnitude branch conditioning. /// `source_ptr` is the [B, SH2] buffer to use as first columns — either /// save_h_s2 (before VSN) or vsn_masked_buf (after VSN). + /// + /// Plan 4 Task 2c.3c.6: kernel now consumes ISV[H_S2_RMS_EMA_INDEX=96] + /// (populated each step by `launch_h_s2_rms_ema`, the producer wired in + /// 2c.3c.5) so the Q_dir tail of the concat is adaptively rescaled to + /// match the trunk-output RMS regardless of GRN drift across training. + /// Two trailing kernel args added: ISV bus device pointer + slot index. pub(crate) fn launch_mag_concat_from(&self, source_ptr: u64, v_logits_ptr: u64, b_logits_ptr: u64) -> Result<(), MLError> { + // `isv_signals_dev_ptr` is unconditionally allocated by the constructor + // via `cuMemHostGetDevicePointer_v2` on a `cuMemAllocHost_v2`-backed + // pinned buffer (with its own `assert_eq!` on success) and is never + // reassigned. Mirrors the 2c.3c.5 producer launcher's invariant — + // silent-skip would mask a misconfiguration. + debug_assert!(self.isv_signals_dev_ptr != 0, + "launch_mag_concat_from: isv_signals_dev_ptr must be allocated by constructor"); let b = self.config.batch_size; let sh2 = self.config.shared_h2 as i32; let na = self.config.num_atoms as i32; @@ -3652,6 +3665,8 @@ impl GpuDqnTrainer { let h_s2_ptr = source_ptr; let concat_ptr = self.ptrs.mag_concat_buf; let support_ptr = self.per_sample_support_ptr; + let isv_dev_ptr = self.isv_signals_dev_ptr; + let isv_idx = H_S2_RMS_EMA_INDEX as i32; unsafe { self.stream .launch_builder(&self.mag_concat_kernel) @@ -3664,6 +3679,8 @@ impl GpuDqnTrainer { .arg(&na) .arg(&b0) .arg(&support_ptr) + .arg(&isv_dev_ptr) // Plan 4 Task 2c.3c.6 + .arg(&isv_idx) // Plan 4 Task 2c.3c.6 .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index e8fce7749..46e342d74 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.6 (2026-04-25): consumer wire-up for ISV[H_S2_RMS_EMA_INDEX=96] — the producer-only slot landed in 2c.3c.5 is now consumed by `mag_concat_qdir` (the magnitude-branch input-builder kernel in `experience_kernels.cu`). Two trailing kernel args added: `const float* __restrict__ isv` + `int isv_h_s2_rms_index`. The kernel's per-sample tail (the `b0_size` Q_dir slots written into `concat_out[b, SH2..SH2+b0_size]`) is now adaptively rescaled in three passes: pass 1 reuses the existing softmax→eq computation per direction action and stashes results in a 4-element register array (`MAG_CONCAT_MAX_DIR=4`, matching the project's 4-direction `S/H/L/F` invariant — `branch_0_size` stays a runtime arg for signature stability but production callers always pass 4); pass 2 computes `q_rms = sqrt(sum_a(eq_a^2) / b0_size)`; pass 3 picks `scale = (q_rms > 1e-6) ? (h_s2_rms_ema / q_rms) : (1 / max(dz, 1e-6))` and writes `concat_out[…, SH2+a] = eq_a * scale`. The legacy formula `concat_out[…, SH2+a] = eq / fmaxf(dz, 1e-6f)` and its 2-line comment ("Normalize by delta_z so Q_dir ~ 1-10 (matches h_s2 post-ReLU scale). Without this, Q_dir ~ 0.1 → 10× smaller gradient → 10× slower learning.") are deleted — the calibration was carried over from the pre-GRN post-ReLU trunk and is superseded by the runtime-measured RMS. The fallback branch is annotated `domain: uniform Q across actions has no RMS to match` (mathematically required when the b0_size-vector is zero, not a stub return). Launch site `launch_mag_concat_from` in `gpu_dqn_trainer.rs` extended with `isv_signals_dev_ptr` + `H_S2_RMS_EMA_INDEX as i32`; `debug_assert!(self.isv_signals_dev_ptr != 0, …)` mirrors the 2c.3c.5 producer launcher's invariant. Backward path unchanged: `strided_accumulate` extracts `d_h_s2` from the first SH2 columns of `d_mag_concat` as before; `h_s2_rms_ema` and `q_rms` are treated as fixed scalars at this batch's launch (same convention as `dz`/`v_min` from `per_sample_support`), no gradient flows back through the ISV read. **No fingerprint change** — no ISV slot or param tensor added; pure kernel-signature + launcher edit. Smoke (`cargo test … multi_fold_convergence --ignored --release`, 649.37s, 3 folds × 5 epochs): all 3 `dqn_fold{N}_best.safetensors` checkpoints written; per-fold best train Sharpe 8.06 / 43.06 / 19.16 at epochs 5 / 1 / 2 (vs 2c.3c.5 baseline 1.91 / 95.56 / 44.00 → geom-mean 20.03; this commit geom-mean 18.80, -6.1% — within the 30% acceptance band, no regression-grade drift). No NaN/Inf, no fingerprint mismatch (fingerprint unchanged at `0x3e21acecd922e540`). 0 panic gates added/removed. The 2c.3c chain — H_S2_RMS_EMA producer (2c.3c.5) + consumer (this commit) — is closed: ISV[96] is now Wired-Producer+Consumer (was Wired-Producer-Only). cargo check clean at 11 warnings (baseline preserved); cargo build compiles 81 cubins (unchanged — kernel-signature edit, no new `.cu`). +69 / -6 LOC across `crates/ml/src/cuda_pipeline/experience_kernels.cu` (kernel signature + body + docstring) and `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (launcher + 2 args + invariant assert + docstring). 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 (`cargo test … multi_fold_convergence --ignored --release`, 642.79s, 3 folds × 5 epochs): all 3 `dqn_fold{N}_best.safetensors` checkpoints written; per-fold best train Sharpe 1.91 / 95.56 / 44.00 at epochs 2 / 4 / 2 (vs 2c.3c.4 baseline 7.52 / 60.94 / 10.40 — fold 0 down, folds 1+2 substantially up; cumulative geometric mean rises from 19.6 to 27.4). No NaN/Inf, no fingerprint mismatch (fingerprint shifted to `0x3e21acecd922e540` and re-validated at constructor as expected since the smoke starts from scratch — no checkpoint load). New producer kernel `h_s2_rms_ema_kernel.cu` launches cleanly per step alongside `reward_component_ema`; ISV[96] populated through training but unread (consumer wires up in 2c.3c.6). 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.