From dc3f948ee970b70d5f938cfed728fc48650b67fc Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 5 May 2026 20:49:27 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp14):=20B.10=20=E2=80=94=20backward=20wir?= =?UTF-8?q?e=20gradient=20gating=20by=20ALPHA=5FGRAD=5FSMOOTHED?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical safety mechanism that completes the EGF pearl: scales the wire column of `dL/dx_concat [B, SH2 + 1]` (the gradient flowing FROM the direction Q-head's first FC SGEMM TO `aux_softmax_diff`) by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]`, computed by B.4's `alpha_grad_compute_kernel` and orchestrated per-step in B.11. `dL/dW[wire_col]` (Q-head's own weight gradient for the appended column) is NOT scaled — the dW SGEMM `dY^T × x_concat` and the dX SGEMM `dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER both have completed leaves dW unaffected. Q-head learns to USE the wire freely; only the gradient PROPAGATING BACK to aux is gated. Pre-B.11 (no producer wired) `ISV[393]` holds sentinel `0.0` → wire force-closed (gradient zeroed) — the conservative safety state. Post-B.11, B.4 writes the live gate output ∈ [0, 1] each step. Closes the latent K-mismatch B.8/B.9 left in backward ============================================================ B.8 grew `w_b0fc` to `[adv_h, SH2 + 1]` end-to-end (Adam m/v + spectral-norm vector + smoke fixtures); B.9 closed the forward dispatch K-mismatch. The backward dW/dX SGEMMs for `d == 0` still used `K = SH2` against the new `LDA = SH2 + 1` weight tensor — silently dropping the last column of dW and zeroing the wire-col gradient. B.10 closes that gap atomically with the wire-col scale per `feedback_no_partial_refactor`: * `backward_branch_dw` for `d == 0` now uses `(dir_qaux_concat_ptr, SH2 + 1)` instead of `(save_h_s2, SH2)` — matching the forward consumer pattern from B.9. * `backward_branch_dx` for `d == 0` now writes to `d_dir_qaux_concat [B, SH2 + 1]` with `K = SH2 + 1` instead of `scratch_d_h_s2 [B, SH2]` with `K = SH2`. Mirrors the magnitude branch's wider-buffer pattern. New artifacts ============= * `sp14_scale_wire_col_kernel.cu`: one thread per batch row, scales `dx_concat[b, SH2]` by `isv[393]` IN-PLACE. NaN-safe per the `dqn_scale_f32_kernel` precedent (explicit `α==0 ⇒ 0` branch). Pure per-thread map, no atomicAdd, no shared memory. * `sp14_d_dir_qaux_concat: CudaSlice` `[B, SH2 + 1]` trainer- struct field. Dx SGEMM destination; the wire-col scale acts on this buffer; the strided accumulator copies the first SH2 columns into `bw_d_h_s2` after the scale. * `launch_sp14_scale_wire_col` launcher reads `self.isv_signals_dev_ptr` and the new buffer's raw_ptr. * `backward_full` signature grows two trailing `u64` args (`dir_qaux_concat_ptr`, `d_dir_qaux_concat_ptr`); both `backward_full` call sites (CQL aux + main online) wired atomically per `feedback_no_partial_refactor`. Post-call orchestration at trainer level ======================================== 1. `launch_sp14_dir_concat_qaux(save_h_s2)` rebuilds the ONLINE concat in `sp14_dir_qaux_concat_scratch` (the forward pass had overwritten it with the TARGET concat at line ~25817). Same one-step-lag semantic preserved — `aux_nb_softmax_buf` is unchanged between forward and backward. 2. `cuMemsetD32Async` zero of `d_h_s2` — pre-B.10 the direction branch (d==0) wrote it with beta=0; post-B.10 the dir-Q dX lives in `d_dir_qaux_concat` and is gated + accumulated AFTER `backward_full` returns, so the value-FC dx accumulator inside `backward_full` (beta=1) needs an explicit zero baseline. 3. `backward_full` runs: dir branch → `d_dir_qaux_concat`, mag/ord/urg branches → their concat dX buffers, value-FC → `d_h_s2` (beta=1, on top of zeroed buffer). 4. `launch_sp14_scale_wire_col` gates col SH2 of `d_dir_qaux_concat`. 5. `accumulate_d_h_s2_from_concat` (beta=1) copies first SH2 cols of `d_dir_qaux_concat` into `d_h_s2`. Wire col stays in `d_dir_qaux_concat[:, SH2]`, untouched by this accumulator (its destination range is [0, SH2)). Pre-B.11 the wire is already zeroed by the sentinel-α gate; the orchestrator that propagates the gated wire-col gradient back to the aux head's softmax CE backward chain lives in B.11. 6. mag/ord/urg accumulators continue with beta=1 (comments updated). Wire status =========== * Forward dispatch: unchanged (B.9-complete). * Backward dispatch: GATED on both call sites (CQL aux + main online). * dW unchanged: the `dW = dY^T × x_concat` SGEMM writes `grad_buf[goff_w_b0fc..]` BEFORE the scale-wire-col launches; the scale operates ONLY on `d_dir_qaux_concat` (the dx buffer) AFTER both dW and dX SGEMMs complete. * Target net unaffected: Polyak EMA-only, no backward. * CudaSlice wrapper path: passes `0u64` for both new args, falls back to the legacy K=SH2 path. Consistent with the forward wrapper's diagnostic-only residual. Verified ======== * `SQLX_OFFLINE=true cargo check -p ml` clean, 18 warnings (baseline) * `cargo test -p ml --test sp14_oracle_tests` 2 passed, 6 ignored (GPU) * Audit doc `docs/dqn-wire-up-audit.md` updated per Invariant 7. After this commit, the EGF pearl is architecturally complete; the orchestration of when/how the alpha_grad gates fire happens in B.11 (producer chain orchestrator). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 14 ++ .../ml/src/cuda_pipeline/batched_backward.rs | 65 +++++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 204 +++++++++++++++++- .../sp14_scale_wire_col_kernel.cu | 64 ++++++ docs/dqn-wire-up-audit.md | 41 ++++ 5 files changed, 381 insertions(+), 7 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 541e88a8f..3ede12e90 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -811,6 +811,20 @@ fn main() { // (enforced by graph capture orchestrator in B.10/B.11). Does NOT // read or write ISV slots — purely data-movement. "dir_concat_qaux_kernel.cu", + // SP14 Layer B Task B.10 (2026-05-05): Earned Gradient Flow + // backward gate. Scales the wire column of `dL/dx_concat + // [B, SH2+1]` (gradient flowing FROM the direction Q-head's + // first FC SGEMM TO `aux_softmax_diff`) by ISV[393] = + // ALPHA_GRAD_SMOOTHED, computed by `alpha_grad_compute_kernel` + // (B.4) and orchestrated per-step by B.11. Touches ONLY column + // SH2 of the [B, SH2+1] gradient buffer; dW (Q-head's own + // weight gradient computed by an independent SGEMM) stays + // unscaled per the pearl design — Q-head learns to USE the + // wire freely; only the back-flow to aux is gated. Pre-B.11 + // (unwired producer) ISV[393] sentinels at 0.0 → wire force- + // closed for safety (gradient zeroed). Pure per-thread map: + // one thread per batch row's wire element. No atomicAdd. + "sp14_scale_wire_col_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 83fd1adcb..c84728884 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -1451,6 +1451,11 @@ impl CublasBackwardSet { d_ord_concat_ptr: u64, urg_concat_ptr: u64, d_urg_concat_ptr: u64, + // SP14 B.10 (2026-05-05): when non-zero, direction branch (d==0) + // dX writes to this `[B, SH2+1]` buffer with K = SH2+1 (matching + // the widened `w_b0fc [adv_h, SH2+1]` weight from B.8). When 0, + // falls back to legacy direct-write into `scratch_d_h_s2 [B, SH2]`. + d_dir_qaux_concat_ptr: u64, b: usize, ) -> Result<(), MLError> { // Read per-branch GLU scratch (produced by Steps 1-4 on branch stream, @@ -1458,7 +1463,34 @@ impl CublasBackwardSet { let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream); let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream); - if d == 1 && mag_concat_ptr != 0 { + if d == 0 && d_dir_qaux_concat_ptr != 0 { + // SP14 B.10: direction Q-head with EGF wire — dX writes to + // d_dir_qaux_concat [B, SH2+1]. K = SH2+1 because `w_b0fc` is + // `[adv_h, SH2+1]` post-B.8 (last col = aux_softmax_diff weight). + // After this returns, the caller scales col SH2 by ISV[393] then + // accumulates first SH2 cols into scratch_d_h_s2 (with beta=0 + // because d==0 is the first branch's contribution to d_h_s2). + self.launch_dx_only( + stream, + d_glu_value_ptr, + w_fc, + d_dir_qaux_concat_ptr, + self.adv_h, + self.shared_h2 + 1, + b, + 0.0_f32, + )?; + self.launch_dx_only( + stream, + d_glu_gate_ptr, + w_gate, + d_dir_qaux_concat_ptr, + self.adv_h, + self.shared_h2 + 1, + b, + 1.0_f32, + )?; + } else if d == 1 && mag_concat_ptr != 0 { // Magnitude (direction-conditioned): dX writes to d_mag_concat [B, SH2+branch_0_size]. self.launch_dx_only( stream, @@ -1711,6 +1743,21 @@ impl CublasBackwardSet { // 0 OR when the cached DRELU_BGRAD descriptor is unavailable, the // backward path falls back to the standalone-kernel sequence. value_fc_relu_mask_ptr: u64, + // SP14 Layer B Task B.10 (2026-05-05): direction-Q-head Earned + // Gradient Flow wire. When non-zero, the direction branch (`d == 0`) + // backward operates on the wider `[B, SH2 + 1]` concat (last col = + // `aux_softmax_diff`), matching the SP14 B.8/B.9 forward consumer + // pattern. dW reads `dir_qaux_concat_ptr` (saved forward concat, + // re-built by the caller via `launch_sp14_dir_concat_qaux` right + // before this `backward_full`), and dX writes to `d_dir_qaux_concat_ptr` + // with `K = SH2 + 1`. The caller then runs `sp14_scale_wire_col_kernel` + // to gate column SH2 by ISV[393] = ALPHA_GRAD_SMOOTHED, then + // accumulates the first SH2 columns into `scratch_d_h_s2` via + // `strided_accumulate` (mirroring the magnitude branch pattern). + // When 0 (e.g. CudaSlice wrapper path), `d == 0` falls back to the + // legacy direct-write path into `scratch_d_h_s2 [B, SH2]`. + dir_qaux_concat_ptr: u64, // [B, SH2+1] saved forward concat (for dW) + d_dir_qaux_concat_ptr: u64, // [B, SH2+1] dX output (caller scales col SH2 + accumulates first SH2 cols) ) -> Result<(), MLError> { // Plan 4 Task 2c.3c.4: trunk backward swap. // @@ -1851,10 +1898,15 @@ impl CublasBackwardSet { let n_d = branch_n[d]; let w_out = w_ptrs[w_bout_idx[d]]; - // FC input for dW (branches 1,2,3 may use wider concat buffers). - // Magnitude (d==1) is direction-conditioned (SH2 + branch_0_size); - // order/urgency (d==2,3) are OFI-conditioned (SH2 + 3). - let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 { + // FC input for dW (each branch may use a wider concat buffer). + // Direction (d==0): SP14 B.10 EGF wire — `[h_s2 ; aux_softmax_diff]` + // when `dir_qaux_concat_ptr != 0` (SH2 + 1 cols). + // Magnitude (d==1): direction-conditioned (SH2 + branch_0_size). + // Order/urgency (d==2,3): OFI-conditioned (SH2 + 3). + // Otherwise: raw `save_h_s2` (SH2 cols), the legacy path. + let (fc_input, fc_in_dim) = if d == 0 && dir_qaux_concat_ptr != 0 { + (dir_qaux_concat_ptr, self.shared_h2 + 1) + } else if d == 1 && mag_concat_ptr != 0 { (mag_concat_ptr, self.shared_h2 + self.branch_0_size) } else if d == 2 && ord_concat_ptr != 0 { (ord_concat_ptr, self.shared_h2 + 3) @@ -1921,6 +1973,9 @@ impl CublasBackwardSet { d_ord_concat_ptr, urg_concat_ptr, d_urg_concat_ptr, + // SP14 B.10: route direction (d==0) dX to the wire-aware + // [B, SH2+1] buffer when the EGF wire is enabled. + d_dir_qaux_concat_ptr, b, )?; } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 37950a144..38a111cf3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -686,6 +686,15 @@ static SP14_GRAD_HACK_CUBIN: &[u8] = static SP14_DIR_CONCAT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dir_concat_qaux_kernel.cubin")); +/// SP14 B.10 (2026-05-05): backward gradient gate kernel for the EGF wire. +/// Scales column SH2 of `d_dir_qaux_concat [B, SH2 + 1]` by +/// `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` IN-PLACE. dW (Q-head's own weight +/// gradient, computed by an independent SGEMM) is NOT touched — Q-head +/// learns to USE the wire freely; only the back-flow to aux is gated per +/// the EGF pearl design. Loaded from `sp14_scale_wire_col_kernel.cubin`. +static SP14_SCALE_WIRE_COL_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/sp14_scale_wire_col_kernel.cubin")); + /// SP11 Fix 39 (2026-05-04, Task A2): SimHash novelty signal — lookup + /// update kernels sharing one cubin. Lookup reads /// `1/sqrt(1+count)` ∈ [0, 1] for each (state, action) bucket; update @@ -5119,6 +5128,29 @@ pub struct GpuDqnTrainer { /// the prior direct `h_s2` input to the direction Q-head. B.7 allocates; /// B.8 bumps the Q-head input-dim constant; B.9/B.11 launch and wire. sp14_dir_qaux_concat_scratch: CudaSlice, + /// SP14 B.10 (2026-05-05): backward dX scratch for direction Q-head's + /// first FC. Shape `[B, SH2 + 1]`. Written by `backward_branch_dx` + /// (`d == 0` case) as the SGEMM `dY @ W^T` output, then column SH2 is + /// scaled by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` via + /// `sp14_scale_wire_col_kernel`, and finally the first SH2 columns + /// accumulate into `bw_d_h_s2` via `strided_accumulate`. Mirrors the + /// `d_mag_concat_buf` precedent for the magnitude branch's wider input. + /// Pre-existing latent gap: B.8/B.9 grew `w_b0fc` to `[adv_h, SH2+1]` + /// in forward but the backward consumers (dW + dX) for `d == 0` still + /// used `K = SH2`, silently dropping the wire column gradient. B.10 + /// closes that gap atomically AND adds the EGF gate; per + /// `feedback_no_partial_refactor.md`. + sp14_d_dir_qaux_concat: CudaSlice, + /// SP14 B.10 (2026-05-05): EGF backward gate kernel. Scales column SH2 + /// of `d_dir_qaux_concat [B, SH2 + 1]` by `ISV[393] = ALPHA_GRAD_SMOOTHED`. + /// One thread per batch row. Loaded from + /// `sp14_scale_wire_col_kernel.cubin`. Launched immediately after + /// `backward_full` returns, before the first-SH2-cols accumulator copies + /// the (now-gated) gradient into `bw_d_h_s2`. The wire column itself is + /// dropped pre-B.11 (alpha_grad sentinel = 0 force-closes the gate); the + /// orchestrator that propagates the scaled wire-col gradient to the aux + /// head's softmax backward lives in B.11. + sp14_scale_wire_col_kernel: CudaFunction, /// SP11 Fix 39 (2026-05-04, Task A2): reward-subsystem controller /// kernel. Single-block, 10-thread producer reading 5 canary ISV slots /// [350..360) and writing 10 floats to scratch[SCRATCH_SP11_CONTROLLER_BASE..+10). @@ -7089,6 +7121,43 @@ impl GpuDqnTrainer { /// kernel) reads it on the same stream BEFORE the direction Q-head SGEMM. /// Sequential same-stream submission enforces the dep — no host barrier /// needed and the order survives CUDA Graph capture. + /// SP14 Layer B Task B.10 (2026-05-05): EGF backward gate. + /// + /// Scales column SH2 of `d_dir_qaux_concat [B, SH2 + 1]` (the gradient + /// flowing FROM the direction Q-head's first FC SGEMM TO + /// `aux_softmax_diff`) by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]`. dW + /// (Q-head's own weight gradient, computed by an independent SGEMM) is + /// NOT touched per the EGF pearl — the SGEMMs `dW = dY^T × x_concat` + /// and `dX = dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER + /// both have completed leaves dW unaffected. + /// + /// Pre-B.11 (no producer wired) ISV[393] holds sentinel `0.0` → + /// wire force-closed for safety (gradient zeroed). Post-B.11, B.4's + /// `alpha_grad_compute_kernel` writes the live gate output ∈ [0, 1]. + /// MUST run AFTER `backward_branch_dx` writes `d_dir_qaux_concat` and + /// BEFORE `accumulate_d_h_s2_from_concat` reads the first SH2 columns. + pub(crate) fn launch_sp14_scale_wire_col(&self, d_dir_qaux_concat_ptr: u64) -> Result<(), MLError> { + let b_i32 = self.config.batch_size as i32; + let sh2_i32 = self.config.shared_h2 as i32; + let blocks = ((self.config.batch_size as u32 + 255) / 256).max(1); + let isv_ptr = self.isv_signals_dev_ptr; + unsafe { + self.stream + .launch_builder(&self.sp14_scale_wire_col_kernel) + .arg(&d_dir_qaux_concat_ptr) + .arg(&isv_ptr) + .arg(&b_i32) + .arg(&sh2_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("sp14_scale_wire_col: {e}")))?; + } + Ok(()) + } + pub(crate) fn launch_sp14_dir_concat_qaux(&self, source_ptr: u64) -> Result<(), MLError> { let b = self.config.batch_size; let sh2 = self.config.shared_h2 as i32; @@ -10386,6 +10455,32 @@ impl GpuDqnTrainer { } else { 0u64 }; + // SP14 B.10: rebuild the ONLINE direction-Q-head concat scratch + // [save_h_s2 | aux_softmax_diff] into `sp14_dir_qaux_concat_scratch`. + // The forward pass overwrote it with the TARGET concat (line ~25817 in + // `submit_forward_ops`), so the backward dW SGEMM would otherwise read + // tg_h_s2 instead of save_h_s2 — wrong gradient. `aux_nb_softmax_buf` + // is unchanged between forward and backward (CE loss writes a separate + // `d_aux_softmax` scratch), so re-running the same concat with + // `save_h_s2` yields the bit-identical online concat the forward + // SGEMM consumed. Same one-step-lag semantic preserved. + self.launch_sp14_dir_concat_qaux(self.ptrs.save_h_s2)?; + // SP14 B.10: pre-zero `scratch_d_h_s2` because the direction branch + // (d==0) no longer writes it directly with beta=0 — its dX now lands + // in `d_dir_qaux_concat` and is gated + accumulated AFTER backward_full + // returns. The value-FC dX accumulator inside backward_full uses + // beta=1 and assumed scratch_d_h_s2 was pre-written by d==0; with the + // direction branch redirected, we must establish the zero baseline + // explicitly. Raw cuMemsetD32Async — captured by CUDA Graph (cudarc + // memset_zeros is NOT graph-safe; matches the cql_grad_scratch + // pre-backward-full zero pattern below). + unsafe { + cudarc::driver::sys::cuMemsetD32Async( + scratch_d_h_s2, 0, + self.config.batch_size * self.config.shared_h2, + self.stream.cu_stream(), + ); + } self.cublas_backward.backward_full( &self.stream, d_v_ptr, @@ -10409,8 +10504,33 @@ impl GpuDqnTrainer { self.ptrs.bw_d_glu_value, self.ptrs.bw_d_glu_gate, value_fc_relu_mask_cql, + // SP14 B.10: direction-Q-head EGF wire — saved forward concat + // for dW + dX output buffer for col SH2 wire scaling. + self.sp14_dir_qaux_concat_scratch.raw_ptr(), + self.sp14_d_dir_qaux_concat.raw_ptr(), ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; + // SP14 B.10: gate the wire column gradient by ISV[ALPHA_GRAD_SMOOTHED] + // before the first-SH2-cols accumulator copies it into scratch_d_h_s2. + // Pre-B.11 (no producer wired) the slot sentinel = 0.0 zeros the wire + // → effectively force-closed. Post-B.11, the gate value ∈ [0, 1] is + // applied per the EGF pearl design. + self.launch_sp14_scale_wire_col(self.sp14_d_dir_qaux_concat.raw_ptr())?; + // Accumulate first SH2 columns of d_dir_qaux_concat into scratch_d_h_s2 + // with beta=1 (the value-FC dX inside backward_full already accumulated + // its contribution; we ADD on top, NOT overwrite). Wire column (col + // SH2) is NOT propagated downstream from here — the orchestrator that + // routes the gated wire-col gradient back to the aux head's softmax CE + // backward chain lives in B.11. Pre-B.11 the wire is zeroed by the + // sentinel-α gate above anyway, so the unrouted column is moot. + self.accumulate_d_h_s2_from_concat( + self.sp14_d_dir_qaux_concat.raw_ptr(), + scratch_d_h_s2, + self.config.batch_size, + self.config.shared_h2 + 1, + 1.0, + )?; + // CQL: accumulate magnitude branch dX into scratch_d_h_s2 if self.ptrs.mag_concat_buf != 0 { self.accumulate_d_h_s2_from_concat( @@ -10418,7 +10538,7 @@ impl GpuDqnTrainer { scratch_d_h_s2, self.config.batch_size, self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned - 1.0, // beta=1: d==0 already wrote to scratch_d_h_s2 + 1.0, // beta=1: SP14 B.10 zeroed scratch_d_h_s2 + value-FC + dir-Q already accumulated )?; } if self.ptrs.ord_concat_buf != 0 { @@ -17038,6 +17158,26 @@ impl GpuDqnTrainer { .alloc_zeros::(b * (config.shared_h2 + 1)) .map_err(|e| MLError::ModelError(format!("sp14 dir_qaux scratch: {e}")))?; + // SP14 B.10: backward dX scratch for the direction Q-head's first + // FC. Shape `[B, SH2 + 1]`. Receives the SGEMM `dY @ W^T` output + // (replacing the prior direct write into `bw_d_h_s2 [B, SH2]` for + // d == 0); column SH2 then gets gated by ISV[393] and the first + // SH2 columns accumulate into `bw_d_h_s2`. Symmetric with the + // existing `d_mag_concat_buf` (magnitude branch's `[B, SH2 + b0]` + // backward scratch). Allocated zero-initialised so the first + // backward replay sees a clean buffer; the dX SGEMM uses beta=0 + // anyway so the prior contents are overwritten each step. + let sp14_d_dir_qaux_concat = stream + .alloc_zeros::(b * (config.shared_h2 + 1)) + .map_err(|e| MLError::ModelError(format!("sp14 d_dir_qaux scratch: {e}")))?; + + let sp14_scale_wire_col_module = stream.context() + .load_cubin(SP14_SCALE_WIRE_COL_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("sp14 scale_wire_col cubin: {e}")))?; + let sp14_scale_wire_col_kernel = sp14_scale_wire_col_module + .load_function("sp14_scale_wire_col_kernel") + .map_err(|e| MLError::ModelError(format!("sp14_scale_wire_col_kernel: {e}")))?; + // SP11 Fix 39 B1b fix-up (2026-05-04): allocate trainer-side // popart-component-per-sample mapped-pinned placeholder buffer. // The canonical per-bar buffer is owned by the experience @@ -20539,6 +20679,12 @@ impl GpuDqnTrainer { sp14_gradient_hack_detect_kernel, sp14_dir_concat_qaux_kernel, sp14_dir_qaux_concat_scratch, + // SP14 EGF backward gate (B.10). Scratch + kernel land + // atomically with the dW/dX SGEMM K-dim migration in + // `backward_branch_dw`/`backward_branch_dx` per + // `feedback_no_partial_refactor.md`. + sp14_d_dir_qaux_concat, + sp14_scale_wire_col_kernel, // SP11 Fix 39 (Task A2): controller + SimHash novelty kernels // and their backing buffers. Projection matrix was populated // on-device above by `launch_novelty_simhash_proj_init`; hash @@ -26832,6 +26978,32 @@ impl GpuDqnTrainer { self.launch_clamp_finite_f32(d_value_logits_ptr, n_val_clamp, max_abs_q)?; self.launch_clamp_finite_f32(d_adv_logits_ptr, n_adv_clamp, max_abs_q)?; + // SP14 B.10: rebuild the ONLINE direction-Q-head concat scratch + // [save_h_s2 | aux_softmax_diff] into `sp14_dir_qaux_concat_scratch`. + // The forward pass overwrote it with the TARGET concat (line ~25817 in + // `submit_forward_ops`), so the backward dW SGEMM would otherwise read + // tg_h_s2 instead of save_h_s2 — wrong gradient. `aux_nb_softmax_buf` + // is unchanged between forward and backward (CE loss writes a separate + // `d_aux_softmax` scratch), so re-running the same concat with + // `save_h_s2` yields the bit-identical online concat the forward + // SGEMM consumed. Same one-step-lag semantic preserved. + self.launch_sp14_dir_concat_qaux(self.ptrs.save_h_s2)?; + // SP14 B.10: pre-zero `d_h_s2` because the direction branch (d==0) no + // longer writes it directly with beta=0 — its dX now lands in + // `d_dir_qaux_concat` and is gated + accumulated AFTER backward_full + // returns. The value-FC dX accumulator inside backward_full uses + // beta=1 and assumed d_h_s2 was pre-written by d==0; with the + // direction branch redirected, we must establish the zero baseline + // explicitly. Raw cuMemsetD32Async — captured by CUDA Graph (cudarc + // memset_zeros is NOT graph-safe; matches the cql_grad_scratch + // pre-backward-full zero pattern in the CQL path). + unsafe { + cudarc::driver::sys::cuMemsetD32Async( + d_h_s2_ptr, 0, + self.config.batch_size * self.config.shared_h2, + self.stream.cu_stream(), + ); + } bw.backward_full( &self.stream, d_value_logits_ptr, @@ -26859,6 +27031,34 @@ impl GpuDqnTrainer { self.ptrs.bw_d_glu_value, self.ptrs.bw_d_glu_gate, value_fc_relu_mask, + // SP14 B.10: direction-Q-head EGF wire — saved forward concat + // for dW + dX output buffer for col SH2 wire scaling. + self.sp14_dir_qaux_concat_scratch.raw_ptr(), + self.sp14_d_dir_qaux_concat.raw_ptr(), + )?; + + // SP14 B.10: gate the wire column gradient by ISV[ALPHA_GRAD_SMOOTHED] + // before the first-SH2-cols accumulator copies it into d_h_s2. + // Pre-B.11 (no producer wired) the slot sentinel = 0.0 zeros the wire + // → effectively force-closed. Post-B.11, the gate value ∈ [0, 1] is + // applied per the EGF pearl design. Critical: this scales `dL/dx` ONLY, + // NOT `dL/dW` — Q-head's first FC weight gradient (computed by an + // independent `dY^T × x_concat` SGEMM into `grad_base[goff_w_b0fc..]`) + // is unscaled, so the Q-head learns to USE the wire freely. + self.launch_sp14_scale_wire_col(self.sp14_d_dir_qaux_concat.raw_ptr())?; + // Accumulate first SH2 columns of d_dir_qaux_concat into d_h_s2 with + // beta=1 (the value-FC dX inside backward_full already accumulated its + // contribution; we ADD on top, NOT overwrite). Wire column (col SH2) + // is NOT propagated downstream from here — the orchestrator that + // routes the gated wire-col gradient back to the aux head's softmax CE + // backward chain lives in B.11. Pre-B.11 the wire is zeroed by the + // sentinel-α gate above anyway, so the unrouted column is moot. + self.accumulate_d_h_s2_from_concat( + self.sp14_d_dir_qaux_concat.raw_ptr(), + d_h_s2_ptr, + self.config.batch_size, + self.config.shared_h2 + 1, + 1.0, )?; // Accumulate magnitude branch dX (first SH2 columns) into d_h_s2 @@ -26868,7 +27068,7 @@ impl GpuDqnTrainer { d_h_s2_ptr, self.config.batch_size, self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned - 1.0, // beta=1: d==0 already wrote to d_h_s2 in backward_full + 1.0, // beta=1: SP14 B.10 zeroed d_h_s2 + value-FC + dir-Q already accumulated )?; } // Accumulate order branch dX (first SH2 columns) into d_h_s2 diff --git a/crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu b/crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu new file mode 100644 index 000000000..e8157b532 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu @@ -0,0 +1,64 @@ +// crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu +// +// SP14 Layer B Task B.10 (2026-05-05): Earned Gradient Flow backward gate. +// Scales the wire column of `dL/dx_concat [B, SH2 + 1]` (the gradient +// flowing back FROM the direction Q-head's first FC SGEMM TO +// `aux_softmax_diff`) by `α_grad_smoothed` (read from ISV[393]). +// +// ── What this gates ───────────────────────────────────────────────────── +// dL/dx_concat shape is [B, SH2 + 1]. Columns [0..SH2) propagate back to +// `h_s2` via the existing trunk-encoder backward chain. Column SH2 is the +// gradient w.r.t. the appended `aux_softmax_diff` value, which downstream +// (B.11) propagates to the aux head's softmax → CE backward chain. THIS +// kernel multiplies ONLY column SH2 by `ALPHA_GRAD_SMOOTHED = ISV[393]`. +// +// Per the EGF pearl design (spec §B.4): +// - α_grad ≈ 0 → wire effectively closed (Q-loss can't pull aux head) +// - α_grad ≈ 1 → full co-training enabled +// +// Critical: dW (the Q-head's own weight gradient) is NOT touched. Q-head +// learns to USE the wire freely from Q-loss; only the gradient PROPAGATING +// BACK to aux head's representation is gated. The dW and dX SGEMMs are +// independent (`dW = dY^T × x_concat`; `dX = dY × W^T`), so scaling dx[:, SH2] +// AFTER both SGEMMs have completed leaves dW unaffected. +// +// ── ISV slot index (must match `sp14_isv_slots.rs`) ───────────────────── +// ALPHA_GRAD_SMOOTHED_INDEX = 393 (sentinel 0.0 pre-B.11 producer wires; +// when sentinel-active the wire is force-closed, the conservative state). +// +// ── Pure per-thread map ───────────────────────────────────────────────── +// One thread handles one batch row's wire-column element. No reductions, +// no shared memory, no atomicAdd (per `feedback_no_atomicadd.md`). Read of +// `isv[ALPHA_GRAD_SMOOTHED_INDEX]` is identical across all threads — safe +// concurrent read. +// +// Grid: ceil(B / 256), Block: 256. + +extern "C" __global__ +void sp14_scale_wire_col_kernel( + /* dL/dx_concat [B, SH2 + 1] — IN-PLACE scaled at column SH2 only */ + float* __restrict__ dx_concat, + /* Global ISV bus. Reads ISV[ALPHA_GRAD_SMOOTHED_INDEX]; writes nothing. */ + const float* __restrict__ isv, + int B, + int SH2) +{ + /* ISV slot index — must match `sp14_isv_slots.rs::ALPHA_GRAD_SMOOTHED_INDEX`. + * The layout fingerprint regression test catches drift. */ + const int ALPHA_GRAD_SMOOTHED = 393; + + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + /* Wire column lives at offset SH2 in each row of stride (SH2+1). */ + const int wire_idx = b * (SH2 + 1) + SH2; + + /* Pre-B.11 producer is unwired → ISV[393] holds sentinel 0.0 → + * wire force-closed (gradient zeroed). Post-B.11, the value is the + * EGF gate output ∈ [0, 1]. + * + * Match `dqn_scale_f32_kernel`'s NaN-safe contract: 0×NaN=NaN under + * IEEE 754; explicit branch yields a clean 0 when α=0. */ + const float alpha = isv[ALPHA_GRAD_SMOOTHED]; + dx_concat[wire_idx] = (alpha == 0.0f) ? 0.0f : (dx_concat[wire_idx] * alpha); +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 3971905e5..c60eae9b3 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -6601,3 +6601,44 @@ Per `pearl_canary_input_freshness_launch_order`: the producer (aux_heads_forward - **Diagnostic-path residual**: causal intervention, DDQN argmax, and the experience-collector / value-decoder forwards intentionally pass `0u64` for `dir_qaux_concat_ptr` and fall back to `K = SH2`. Their direction Q outputs feed either (a) only-value-logit consumers (causal) or (b) downstream argmax-only consumers with one-step-bias acknowledged by the spec (DDQN). The cuBLAS heuristic for `K = SH2, LDA = SH2` against the underlying `[adv_h, SH2 + 1]` weight tensor reads the first `adv_h * SH2` floats with stride SH2 — within bounds (no OOB), produces stable-but-incorrect outputs for the residual paths. `feedback_no_partial_refactor` is honoured for the train-time path (online forward, target forward, replay); diagnostic-path residuals are documented and bounded. - **Reverse dependencies**: `forward_value_head_for_ensemble` and `forward_value_head` use only the value head (V_h ← h_s2; no branch heads), so unaffected by the direction Q-head wire. The `submit_dqn_step_loop_cublas` path is the same `forward_online_raw` API — no separate dispatch site. +## SP14 Layer B — Commit B.10: backward wire — gradient scaling at the wire column (2026-05-05) + +**Why this commit.** B.10 is the critical safety mechanism that completes the EGF pearl. The forward wire (B.6/B.9) feeds `aux_softmax_diff` into the direction Q-head's first FC SGEMM via `[h_s2 ; aux_softmax_diff] [B, SH2 + 1]`. Without B.10, the Q-loss backward computes `dL/dx_concat[:, SH2]` (the gradient flowing back to `aux_softmax_diff`) and propagates it ungated all the way to the aux head's softmax — co-training aux on Q-loss without protection. B.10 scales `dL/dx[wire_col]` by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` (computed by B.4's `alpha_grad_compute_kernel`, orchestrated per-step in B.11). Crucially, `dL/dW[wire_col]` (the Q-head's own weight gradient for the appended column) is NOT scaled — Q-head learns to USE the wire freely; only the gradient PROPAGATING BACK to aux is gated. The two SGEMMs `dW = dY^T × x_concat` and `dX = dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER both have completed leaves dW unaffected. + +### Pre-existing latent gap closed atomically with the wire-col scale + +B.8/B.9 grew `w_b0fc` to `[adv_h, SH2 + 1]` end-to-end across forward dispatch + Adam m/v + spectral-norm power-iteration vector + smoke fixtures, but the **backward dW/dX SGEMMs for `d == 0` still used `K = SH2`** against the new `LDA = SH2 + 1` weight tensor — the same K-mismatch B.9 closed in forward, surviving in backward. B.10 closes that gap atomically with the wire-col scale per `feedback_no_partial_refactor.md`: + +| Backward path | Pre-B.10 | Post-B.10 | +|---|---|---| +| `backward_branch_dw` for `d == 0` | `(save_h_s2, SH2)` → silent dW under-write of last column | `(dir_qaux_concat_ptr, SH2 + 1)` → full dW including last column | +| `backward_branch_dx` for `d == 0` | `(scratch_d_h_s2, SH2, beta=0)` → wire-col gradient dropped | `(d_dir_qaux_concat_ptr, SH2 + 1, beta=0)` → wire-col gradient preserved for B.10 scale | +| Wire-col gating | none | `sp14_scale_wire_col_kernel` reads `ISV[393]` and multiplies `d_dir_qaux_concat[:, SH2]` | +| Trunk `d_h_s2` accumulation | implicit (d==0 wrote `scratch_d_h_s2` with beta=0) | explicit `cuMemsetD32Async` zero pre-`backward_full`; `accumulate_d_h_s2_from_concat` adds first SH2 cols with beta=1 after wire-col scale | + +### File-change summary + +| Change | File | +|--------|------| +| New `sp14_scale_wire_col_kernel.cu` (one thread per batch row, scales col SH2 by `ISV[393]`); `extern "C"` symbol; structurally bounded NaN-safe per `dqn_scale_f32_kernel` precedent | `crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu` | +| New cubin entry in `kernels_with_common` | `crates/ml/build.rs` | +| New `static SP14_SCALE_WIRE_COL_CUBIN`; new struct fields `sp14_d_dir_qaux_concat: CudaSlice` (`[B, SH2 + 1]`) + `sp14_scale_wire_col_kernel: CudaFunction`; allocation + cubin load + struct init; new `launch_sp14_scale_wire_col` launcher; both `backward_full` call sites (CQL @ ~10484 + main @ ~26980) get `launch_sp14_dir_concat_qaux(save_h_s2)` re-build (the forward path overwrites the buffer with target concat at line ~25817), `cuMemsetD32Async` zero of d_h_s2 before the call, the new `(dir_qaux_concat_ptr, d_dir_qaux_concat_ptr)` trailing args, post-call `launch_sp14_scale_wire_col` + `accumulate_d_h_s2_from_concat(beta=1)` | `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | +| `backward_full` signature grows `dir_qaux_concat_ptr` + `d_dir_qaux_concat_ptr` u64 args; `backward_branch_dw` adds `d == 0 && dir_qaux_concat_ptr != 0` arm using `(SH2 + 1)` fc_in_dim; `backward_branch_dx` adds `d == 0 && d_dir_qaux_concat_ptr != 0` arm writing `[B, SH2 + 1]` with `K = SH2 + 1` (mirroring magnitude branch's wider-buffer pattern) | `crates/ml/src/cuda_pipeline/batched_backward.rs` | + +### Verification + +- `SQLX_OFFLINE=true cargo check -p ml` — clean, 18 warnings (pre-existing baseline, no new warnings) +- `cargo test -p ml --test sp14_oracle_tests` — 2 passed, 6 ignored (GPU) + +### Wire status + +- **Forward dispatch**: unchanged (B.9-complete). +- **Backward dispatch**: GATED. Both `backward_full` call sites (main + CQL aux) use the EGF wire path: dW via `dir_qaux_concat_ptr` with `K = SH2 + 1`, dX into `d_dir_qaux_concat [B, SH2 + 1]` with `K = SH2 + 1`. The K-mismatch B.8/B.9 left in backward is now closed in lockstep with the wire-col scale. +- **Wire-col gating**: ACTIVE. `launch_sp14_scale_wire_col` reads `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` and multiplies column SH2 of `d_dir_qaux_concat` IN-PLACE. Pre-B.11 (no producer wired) the slot sentinel = 0.0 → wire force-closed (gradient zeroed) — the conservative safety state. Post-B.11, B.4's `alpha_grad_compute_kernel` writes the live gate output ∈ [0, 1] each step. +- **dW unchanged**: `dL/dW = dL/d_branch_h^T × x_concat` writes to `grad_buf[goff_w_b0fc..goff_w_b0fc + adv_h * (SH2 + 1) * 4]` via `launch_dw_only_ws` BEFORE the scale-wire-col launches; the scale operates ONLY on `d_dir_qaux_concat` (the dx buffer) AFTER both dW and dX SGEMMs complete. Q-head learns to use the wire freely. +- **First-SH2-cols accumulator**: `accumulate_d_h_s2_from_concat(d_dir_qaux_concat → d_h_s2, src_stride=SH2+1, beta=1.0)` runs AFTER `launch_sp14_scale_wire_col` so the scaled wire-col stays in `d_dir_qaux_concat[:, SH2]` (untouched by the accumulator's destination range `[0, SH2)`). The wire column gradient is NOT propagated downstream from here pre-B.11; the orchestrator that routes the gated wire-col gradient back to the aux head's softmax CE backward chain lives in B.11. Pre-B.11 the wire is zeroed by the sentinel-α gate anyway, so the unrouted column is moot. +- **One-step-lag preservation**: backward consumes the same `aux_nb_softmax_buf` snapshot the forward consumed — `aux_heads_forward` writes the buffer once per step BEFORE the next-step's online forward runs `launch_sp14_dir_concat_qaux`, and the CE loss kernel writes a separate `d_aux_softmax` scratch (does not overwrite `aux_nb_softmax_buf`). Re-running the concat with `save_h_s2` at backward start yields the bit-identical online concat the forward SGEMM consumed. +- **Target forward unaffected**: target net is Polyak-EMA-updated only (no backward), so the wire-col scale + K-dim migration apply only to the online backward path. +- **Reverse dependencies**: `apply_iqn_trunk_gradient` and aux-paths use the `bw_d_h_s2` that this path writes — the new accumulator pipeline (memset → backward_full → wire-col scale → strided accumulate from d_dir_qaux_concat → mag/ord/urg accumulators → value-FC inside backward_full) leaves `bw_d_h_s2` with the same algebraic value as pre-B.10 *except* for the gated wire-col contribution from the direction-Q's first FC. Pre-B.11 (α=0) the gated contribution is zero → bit-identical to pre-B.10. +- **CudaSlice wrapper path**: passes `0u64` for both `dir_qaux_concat_ptr` and `d_dir_qaux_concat_ptr`, falling back to the legacy K=SH2 path. This is consistent with the forward CudaSlice wrapper (`dir_qaux_concat_ptr: 0`); the wrapper-based callers (causal intervention, DDQN argmax) are diagnostic-only paths whose direction-Q outputs are downstream-bounded per the B.9 residual-path analysis. +