diff --git a/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu index 85d509d7a..19ba02ba0 100644 --- a/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu +++ b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu @@ -106,6 +106,16 @@ __device__ __forceinline__ float aux_elu_bwd_from_post(float y) { * - aux_dir_acc_reduce_kernel (argmax-vs-label compare) * - aux_pred_to_isv_tanh_kernel (mean(softmax[1] - softmax[0]) → ISV[375]) * + * SP14 Layer C Phase C.5b (2026-05-08): the input parameter was renamed + * `h_s2` → `h_s2_aux`. The aux heads now read the SEPARATE aux trunk's + * output rather than Q's GRN trunk output. The Rust caller passes the + * trainer's `h_s2_aux` buffer (populated by `aux_trunk_forward` per + * `gpu_aux_trunk.rs`). Q's `save_h_s2` is no longer wired here. Symmetric + * rename in `aux_next_bar_backward` (`dh_s2_out` → `dh_s2_aux_out`) so + * the gradient flows back through the aux trunk to its own params, NOT + * into Q's encoder trunk. The encoder boundary stop-gradient is enforced + * structurally by `aux_trunk_backward`'s missing `dx_in` output param. + * * Mirroring `aux_regime_forward`'s K-fanout structure — thread 0 computes * all K logits serially (K=2 / K=5; the matmul is small enough that a * tree reduce wastes more cycles in shfl bookkeeping than the serial @@ -115,7 +125,7 @@ __device__ __forceinline__ float aux_elu_bwd_from_post(float y) { * (h cache for Linear_2 + temporary K-vector for stable softmax). * ===================================================================== */ extern "C" __global__ void aux_next_bar_forward( - const float* __restrict__ h_s2, /* [B, SH2] row-major */ + const float* __restrict__ h_s2_aux, /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */ const float* __restrict__ w1, /* [H, SH2] row-major */ const float* __restrict__ b1, /* [H] */ const float* __restrict__ w2, /* [K, H] row-major */ @@ -143,7 +153,7 @@ extern "C" __global__ void aux_next_bar_forward( for (int k = tid; k < H; k += blockDim.x) { float acc = b1[k]; const float* w_row = w1 + (size_t)k * SH2; - const float* h_row = h_s2 + (size_t)b * SH2; + const float* h_row = h_s2_aux + (size_t)b * SH2; for (int j = 0; j < SH2; ++j) { acc += w_row[j] * h_row[j]; } @@ -200,7 +210,7 @@ extern "C" __global__ void aux_next_bar_forward( * Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H floats. * ===================================================================== */ extern "C" __global__ void aux_regime_forward( - const float* __restrict__ h_s2, /* [B, SH2] row-major */ + const float* __restrict__ h_s2_aux, /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */ const float* __restrict__ w1, /* [H, SH2] row-major */ const float* __restrict__ b1, /* [H] */ const float* __restrict__ w2, /* [K, H] row-major (K = 5) */ @@ -221,7 +231,7 @@ extern "C" __global__ void aux_regime_forward( for (int k = tid; k < H; k += blockDim.x) { float acc = b1[k]; const float* w_row = w1 + (size_t)k * SH2; - const float* h_row = h_s2 + (size_t)b * SH2; + const float* h_row = h_s2_aux + (size_t)b * SH2; for (int j = 0; j < SH2; ++j) { acc += w_row[j] * h_row[j]; } @@ -480,7 +490,7 @@ extern "C" __global__ void aux_regime_label_from_states( * ===================================================================== */ extern "C" __global__ void aux_next_bar_backward( /* Forward inputs (re-read for dW). */ - const float* __restrict__ h_s2, /* [B, SH2] */ + const float* __restrict__ h_s2_aux, /* [B, SH2] (SP14-C.5b: aux trunk output) */ const float* __restrict__ w1, /* [H, SH2] */ const float* __restrict__ w2, /* [K, H] row-major */ const float* __restrict__ hidden_post, /* [B, H] saved h_post */ @@ -496,9 +506,12 @@ extern "C" __global__ void aux_next_bar_backward( float* __restrict__ db1_partial, /* [B, H] */ float* __restrict__ dW2_partial, /* [B, K, H] flat */ float* __restrict__ db2_partial, /* [B, K] */ - /* Trunk gradient (contiguous in B, no batch-reduce needed — caller - * accumulates into trunk's main-path d_h_s2 via SAXPY). */ - float* __restrict__ dh_s2_out /* [B, SH2] */ + /* Aux trunk-output gradient: SP14-C.5b accumulator + * `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). Caller SAXPYs both heads' + * dh_s2_aux into this buffer, then aux_trunk_backward reads it as the + * upstream gradient feeding `dh_s2_aux_in_ptr`. Encoder stop-grad is + * enforced structurally by aux_trunk_backward (no dx_in output). */ + float* __restrict__ dh_s2_aux_out /* [B, SH2] */ ) { const int b = blockIdx.x; if (b >= B) return; @@ -586,7 +599,7 @@ extern "C" __global__ void aux_next_bar_backward( /* dW1 partial — per-(k, j) write. Stride loop over (H * SH2). */ { - const float* h_row = h_s2 + (size_t)b * SH2; + const float* h_row = h_s2_aux + (size_t)b * SH2; float* dW1_b = dW1_partial + (size_t)b * H * SH2; const int total = H * SH2; for (int idx = tid; idx < total; idx += blockDim.x) { @@ -596,44 +609,31 @@ extern "C" __global__ void aux_next_bar_backward( } } - /* Step 3: SP14 stop-gradient (2026-05-07). + /* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j]. * - * Diagnostic from L40S `train-v8ztm` 9-epoch HEALTH_DIAG aux next_bar_mse: - * Ep 0: 0.352 (learnable signal — below random baseline ln(2)≈0.693) - * Ep 9: 0.717 (above random baseline — aux is now WORSE than random) - * aux_dir_acc_long stuck at 0.19 (anti-correlated with truth) + * SP14 Layer C Phase C.5b (2026-05-08): the prior commit `872bd7392` + * zero-filled this buffer as a stop-grad band-aid because the aux + * head's gradient was leaking into Q's shared `h_s2`. Phase C.5b + * REVERTS that band-aid: aux heads now read the SEPARATE aux trunk's + * output `h_s2_aux`, and `dh_s2_aux_out` accumulates into the aux + * trunk's own input gradient (`dh_s2_aux_accum`), NOT into Q's + * `bw_d_h_s2`. The aux trunk's backward kernel terminates at the + * encoder boundary (no `dx_in` output param), so Q's encoder is + * structurally protected — no leak possible. With that structural + * guarantee the genuine SAXPY-back-to-input gradient is restored. * - * Root cause: the aux head's backward gradient was flowing back into the - * shared trunk activation `h_s2` via the SAXPY of `dh_s2_out` into the - * trunk's main-path d_h_s2. Q-loss gradient on h_s2 dominates (larger - * magnitude, structurally different objective: cumulative discounted - * reward vs next-bar direction). h_s2 evolves to support Q's task; aux's - * CE loss climbs as h_s2 features become anti-aligned with direction - * prediction. - * - * Fix: stop-gradient on aux's input. We zero `dh_s2_out` rather than - * computing the sum-of-w1·dh_pre. The aux head still trains its own - * params (w1/b1/w2/b2) from the CE loss via the dW1/db1/dW2/db2 - * partials written above — those are unaffected. But h_s2 is no longer - * pulled by aux's gradient. Q-loss is the sole shaping force on h_s2; - * aux must adapt to whatever h_s2 happens to be. - * - * If aux still cannot learn after this fix, the deeper diagnosis is that - * h_s2 carries no direction signal at all — separate-trunk Option 2 - * (deferred). The trunk SAXPY adds zero either way (same arithmetic the - * masked-row branch already produced for skipped samples), so no other - * consumer is affected — `dh_s2_out` is read only by the trunk - * accumulation step. - * - * NOTE: `aux_regime_backward` (the K=5 regime CE head) has the same - * architecture and received the SAME fix in the immediately-following - * commit (mirrors this kernel's Step 3 zero-write). Both auxiliary - * heads now train their own params from CE loss without pulling on - * shared h_s2 — Q-loss is the sole trunk-shaping force. */ + * Per-sample, per-feature serial reduction over H lanes — H is small + * (32) and SH2 is moderate; stride loop over j. Masked rows propagate + * zero d_h_pre through this sum, so dh_s2_aux row is zero for skipped + * samples (the trunk SAXPY adds zero — no spurious gradient). */ { - float* dh_row = dh_s2_out + (size_t)b * SH2; + float* dh_row = dh_s2_aux_out + (size_t)b * SH2; for (int j = tid; j < SH2; j += blockDim.x) { - dh_row[j] = 0.0f; + float acc = 0.0f; + for (int k = 0; k < H; ++k) { + acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j]; + } + dh_row[j] = acc; } } } @@ -658,7 +658,7 @@ extern "C" __global__ void aux_next_bar_backward( * ===================================================================== */ extern "C" __global__ void aux_regime_backward( /* Forward inputs. */ - const float* __restrict__ h_s2, /* [B, SH2] */ + const float* __restrict__ h_s2_aux, /* [B, SH2] (SP14-C.5b: aux trunk output) */ const float* __restrict__ w1, /* [H, SH2] */ const float* __restrict__ w2, /* [K, H] row-major */ const float* __restrict__ hidden_post, /* [B, H] saved h_post */ @@ -673,8 +673,9 @@ extern "C" __global__ void aux_regime_backward( float* __restrict__ db1_partial, /* [B, H] */ float* __restrict__ dW2_partial, /* [B, K, H] */ float* __restrict__ db2_partial, /* [B, K] */ - /* Trunk gradient. */ - float* __restrict__ dh_s2_out /* [B, SH2] */ + /* Aux trunk-output gradient (SAXPYed into `dh_s2_aux_accum`; encoder + * stop-grad enforced structurally by `aux_trunk_backward`). */ + float* __restrict__ dh_s2_aux_out /* [B, SH2] */ ) { const int b = blockIdx.x; if (b >= B) return; @@ -742,9 +743,9 @@ extern "C" __global__ void aux_regime_backward( db1_partial[(size_t)b * H + k] = sh_dh_pre[k]; } - /* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2[b, j]. */ + /* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2_aux[b, j]. */ { - const float* h_row = h_s2 + (size_t)b * SH2; + const float* h_row = h_s2_aux + (size_t)b * SH2; float* dW1_b = dW1_partial + (size_t)b * H * SH2; const int total = H * SH2; for (int idx = tid; idx < total; idx += blockDim.x) { @@ -754,38 +755,24 @@ extern "C" __global__ void aux_regime_backward( } } - /* Step 3: SP14 stop-gradient (2026-05-07, mirrors next-bar fix at commit - * 872bd7392). + /* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j]. * - * The 5-class regime CE head's backward gradient was flowing back into - * the shared trunk activation `h_s2` via the SAXPY of `dh_s2_out` into - * the trunk's main-path d_h_s2 — IDENTICAL architectural conflict to - * the next-bar head closed earlier the same day. Q-loss gradient on - * h_s2 dominates (larger magnitude, structurally different objective: - * cumulative discounted reward vs 5-class regime classification). - * h_s2 evolves to support Q's task; regime's CE loss can't shape - * h_s2 without conflicting with Q's directional pressure. - * - * Fix: stop-gradient. Regime reads h_s2 via forward, trains its own - * w1/b1/w2/b2 from CE loss via the dW1_partial / db1_partial / - * dW2_partial / db2_partial writes above (those are unaffected — they - * read sh_dh_pre, sh_dlogits, sh_h_post, h_row, all computed - * pre-write); only the trunk-bound gradient is severed. Q-loss is the - * sole shaping force on h_s2; regime adapts to whatever h_s2 happens - * to be. - * - * If regime can't learn even with this fix, the deeper bug is that - * h_s2 has no regime-relevant signal — would need separate aux trunk - * (Option 2, deferred). No other consumer reads `dh_s2_out` — only - * the trunk accumulation step. - * - * Today's fix has been validated by passing oracle tests; we're not - * waiting for L40S validation to apply the same fix to regime since - * the fix is mechanical and identical. */ + * SP14 Layer C Phase C.5b (2026-05-08): the prior commit `411a30473` + * zero-filled this buffer as a stop-grad band-aid because the regime + * head's gradient was leaking into Q's shared `h_s2`. Phase C.5b + * REVERTS that band-aid: regime now reads `h_s2_aux` (aux trunk + * output), and `dh_s2_aux_out` SAXPYs into the aux trunk's own input + * gradient `dh_s2_aux_accum`. The aux trunk's backward kernel has + * NO `dx_in` output param so Q's encoder is structurally protected. + * The genuine SAXPY-back-to-input gradient is restored. */ { - float* dh_row = dh_s2_out + (size_t)b * SH2; + float* dh_row = dh_s2_aux_out + (size_t)b * SH2; for (int j = tid; j < SH2; j += blockDim.x) { - dh_row[j] = 0.0f; + float acc = 0.0f; + for (int k = 0; k < H; ++k) { + acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j]; + } + dh_row[j] = acc; } } } diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs index 050b7b07b..a87fb7027 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -142,7 +142,8 @@ impl AuxHeadsForwardOps { /// /// Caller-owned buffers (raw `u64` device pointers for graph-capture /// safety; mirrors `gpu_grn.rs::forward_raw`): - /// * `h_s2_ptr` — `[B, SH2]` row-major. + /// * `h_s2_aux_ptr` — `[B, SH2]` row-major (SP14-C.5b: aux trunk + /// output, NOT Q's GRN trunk `save_h_s2`). /// * `w1_ptr`/`b1_ptr` — `[H, SH2]` / `[H]`. Slice from `params_buf[119]`/`[120]`. /// * `w2_ptr`/`b2_ptr` — `[K, H]` / `[K]`. Slice from `params_buf[121]`/`[122]`. /// * `hidden_out_ptr` — `[B, H]` SAVED post-ELU buffer (consumed by backward). @@ -161,7 +162,7 @@ impl AuxHeadsForwardOps { pub(crate) fn forward_next_bar( &self, stream: &Arc, - h_s2_ptr: u64, + h_s2_aux_ptr: u64, w1_ptr: u64, b1_ptr: u64, w2_ptr: u64, @@ -180,7 +181,7 @@ impl AuxHeadsForwardOps { unsafe { stream .launch_builder(&self.next_bar_forward_kernel) - .arg(&h_s2_ptr) + .arg(&h_s2_aux_ptr) .arg(&w1_ptr) .arg(&b1_ptr) .arg(&w2_ptr) @@ -209,7 +210,7 @@ impl AuxHeadsForwardOps { pub(crate) fn forward_regime( &self, stream: &Arc, - h_s2_ptr: u64, + h_s2_aux_ptr: u64, w1_ptr: u64, b1_ptr: u64, w2_ptr: u64, @@ -227,7 +228,7 @@ impl AuxHeadsForwardOps { unsafe { stream .launch_builder(&self.regime_forward_kernel) - .arg(&h_s2_ptr) + .arg(&h_s2_aux_ptr) .arg(&w1_ptr) .arg(&b1_ptr) .arg(&w2_ptr) @@ -469,7 +470,7 @@ impl AuxHeadsBackwardOps { pub(crate) fn backward_next_bar( &self, stream: &Arc, - h_s2_ptr: u64, + h_s2_aux_ptr: u64, w1_ptr: u64, w2_ptr: u64, hidden_post_ptr: u64, @@ -483,7 +484,7 @@ impl AuxHeadsBackwardOps { db1_partial_ptr: u64, dw2_partial_ptr: u64, db2_partial_ptr: u64, - dh_s2_out_ptr: u64, + dh_s2_aux_out_ptr: u64, ) -> Result<(), MLError> { let b_i32 = b as i32; let sh2_i32 = sh2 as i32; @@ -494,7 +495,7 @@ impl AuxHeadsBackwardOps { unsafe { stream .launch_builder(&self.next_bar_backward_kernel) - .arg(&h_s2_ptr) + .arg(&h_s2_aux_ptr) .arg(&w1_ptr) .arg(&w2_ptr) .arg(&hidden_post_ptr) @@ -508,7 +509,7 @@ impl AuxHeadsBackwardOps { .arg(&db1_partial_ptr) .arg(&dw2_partial_ptr) .arg(&db2_partial_ptr) - .arg(&dh_s2_out_ptr) + .arg(&dh_s2_aux_out_ptr) .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (AUX_BLOCK, 1, 1), @@ -526,7 +527,7 @@ impl AuxHeadsBackwardOps { pub(crate) fn backward_regime( &self, stream: &Arc, - h_s2_ptr: u64, + h_s2_aux_ptr: u64, w1_ptr: u64, w2_ptr: u64, hidden_post_ptr: u64, @@ -539,7 +540,7 @@ impl AuxHeadsBackwardOps { db1_partial_ptr: u64, dw2_partial_ptr: u64, db2_partial_ptr: u64, - dh_s2_out_ptr: u64, + dh_s2_aux_out_ptr: u64, ) -> Result<(), MLError> { let b_i32 = b as i32; let sh2_i32 = sh2 as i32; @@ -549,7 +550,7 @@ impl AuxHeadsBackwardOps { unsafe { stream .launch_builder(&self.regime_backward_kernel) - .arg(&h_s2_ptr) + .arg(&h_s2_aux_ptr) .arg(&w1_ptr) .arg(&w2_ptr) .arg(&hidden_post_ptr) @@ -562,7 +563,7 @@ impl AuxHeadsBackwardOps { .arg(&db1_partial_ptr) .arg(&dw2_partial_ptr) .arg(&db2_partial_ptr) - .arg(&dh_s2_out_ptr) + .arg(&dh_s2_aux_out_ptr) .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (AUX_BLOCK, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 9e14814d2..d00222e3f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -17593,8 +17593,18 @@ impl GpuDqnTrainer { pub(crate) fn aux_heads_forward(&self) -> Result<(), MLError> { let b = self.config.batch_size; let sh2 = self.config.shared_h2; - let h_s2_ptr = self.ptrs.save_h_s2; + // SP14 Layer C Phase C.5b (2026-05-08): aux heads now read the + // SEPARATE aux trunk's output `h_s2_aux` rather than Q's GRN trunk + // output `save_h_s2`. The aux trunk forward (below) consumes + // `save_h_s1` (encoder output) and writes `h_s2_aux`. Symmetric + // contract change in `aux_heads_backward` + the kernel signature + // rename `h_s2 → h_s2_aux` lock-in step. + let h_s2_aux_ptr = self.h_s2_aux.raw_ptr(); let states_ptr = self.ptrs.states_buf; + let encoder_out_ptr = self.ptrs.save_h_s1; + let encoder_out_dim = self.config.shared_h1; + let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1; + let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2; let param_sizes = compute_param_sizes(&self.config); let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); @@ -17608,6 +17618,32 @@ impl GpuDqnTrainer { let rg_w2 = on_w_ptrs[125]; let rg_b2 = on_w_ptrs[126]; + // Step 0: aux trunk forward — Linear→ELU→Linear→ELU→Linear MLP that + // produces `h_s2_aux [B, SH2]` from the encoder's layer-1 output + // `save_h_s1 [B, SH1]`. Runs BEFORE the regime label builder + per- + // head forwards so both heads see the freshly-written `h_s2_aux`. + // Aux trunk weights live in dedicated `aux_trunk_w{1,2,3}` / + // `aux_trunk_b{1,2,3}` slabs (NOT in the flat `params_buf`); they + // are trained by `launch_aux_trunk_adam_update` in the backward. + self.aux_trunk_forward_ops.launch( + &self.stream, + encoder_out_ptr, + self.aux_trunk_w1.raw_ptr(), + self.aux_trunk_b1.raw_ptr(), + self.aux_trunk_w2.raw_ptr(), + self.aux_trunk_b2.raw_ptr(), + self.aux_trunk_w3.raw_ptr(), + self.aux_trunk_b3.raw_ptr(), + self.h_aux1.raw_ptr(), + self.h_aux2.raw_ptr(), + h_s2_aux_ptr, + b, + encoder_out_dim, + aux_h1, + aux_h2, + sh2, + )?; + let state_dim_padded = ml_core::state_layout::STATE_DIM_PADDED; let regime_score_idx = ml_core::state_layout::OFI_START + 29; @@ -17639,9 +17675,10 @@ impl GpuDqnTrainer { // Step 3: next-bar direction-classification forward (SP13 B1.1a: // K=2 softmax). Writes hidden tile + logits tile + softmax tile. + // SP14-C.5b: input is `h_s2_aux` (aux trunk output), NOT Q's `h_s2`. self.aux_heads_fwd.forward_next_bar( &self.stream, - h_s2_ptr, + h_s2_aux_ptr, nb_w1, nb_b1, nb_w2, nb_b2, b, sh2, AUX_NEXT_BAR_K, self.aux_nb_hidden_buf.raw_ptr(), @@ -17650,9 +17687,10 @@ impl GpuDqnTrainer { )?; // Step 4: regime classification forward. + // SP14-C.5b: input is `h_s2_aux` (aux trunk output), NOT Q's `h_s2`. self.aux_heads_fwd.forward_regime( &self.stream, - h_s2_ptr, + h_s2_aux_ptr, rg_w1, rg_b1, rg_w2, rg_b2, b, sh2, AUX_REGIME_K, self.aux_rg_hidden_buf.raw_ptr(), @@ -17685,28 +17723,50 @@ impl GpuDqnTrainer { Ok(()) } - /// Plan 4 Task 6 Commit B: aux-heads backward orchestrator. + /// Plan 4 Task 6 Commit B + SP14 Layer C Phase C.5b (2026-05-08): + /// aux-heads + aux-trunk backward orchestrator. /// - /// Runs both backward kernels (next-bar + regime) over the saved forward - /// state, then for each of the 8 aux param tensors: + /// Runs both head backward kernels (next-bar + regime) over the saved + /// forward state, then for each of the 8 aux head param tensors: /// 1. `aux_param_grad_reduce` collapses `partials [B, P]` into `final [P]`. /// 2. SAXPY `grad_buf[tensor_idx] += aux_weight * final` via the /// existing `dqn_saxpy_f32_kernel`. - /// Finally, the per-sample `dh_s2` from both heads is SAXPY-accumulated - /// into `bw_d_h_s2 [B, SH2]` with `alpha = aux_weight`. + /// Then SP14-C.5b: the per-sample `dh_s2_aux` from both heads is + /// SAXPY-accumulated into `dh_s2_aux_accum [B, SH2]` (the AUX TRUNK's + /// upstream-grad accumulator, NOT Q's `bw_d_h_s2`). Aux trunk backward + /// then reads the accumulator and propagates gradient through w3/w2/w1 + /// + b3/b2/b1, terminating at the encoder boundary (structural stop- + /// grad — kernel set has NO `dx_in` output, so Q's `bw_d_h_s2` is + /// unaffected by aux loss). Finally, `launch_aux_trunk_adam_update` + /// applies global L2-norm clip + per-tensor Adam updates over the 6 + /// aux trunk grad buffers. /// /// MUST run AFTER `backward_full` + concat-accumulators fill `bw_d_h_s2` - /// AND BEFORE `encoder_backward_chain` consumes it. + /// AND BEFORE `encoder_backward_chain` consumes `bw_d_h_s2` (so the + /// trunk dW/dB chain rule sees the un-augmented Q gradient). /// /// `grad_base` is the same value passed to `launch_cublas_backward_to` — /// usually `self.ptrs.grad_buf`, but the vaccine path (g_val computation) - /// passes a scratch buffer here. The aux SAXPY targets `grad_base[119..127)` - /// at the same `padded_byte_offset` math the rest of the backward uses. + /// passes a scratch buffer here. The aux head SAXPY targets + /// `grad_base[119..127)` at the same `padded_byte_offset` math the rest + /// of the backward uses; the aux trunk Adam launcher writes its 6 grad + /// tensors directly (separate slabs, NOT in `params_buf`/`grad_buf`). pub(crate) fn aux_heads_backward(&self, grad_base: u64) -> Result<(), MLError> { let b = self.config.batch_size; let sh2 = self.config.shared_h2; - let h_s2_ptr = self.ptrs.save_h_s2; - let bw_d_h_s2_ptr = self.bw_d_h_s2.raw_ptr(); + // SP14 Layer C Phase C.5b (2026-05-08): aux heads now consume the + // SEPARATE aux trunk's output `h_s2_aux`, and their per-sample + // dh_s2 SAXPYs into `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). The + // aux trunk's backward kernel then reads `dh_s2_aux_accum` as + // `dh_s2_aux_in_ptr` and propagates grad through the aux trunk's + // own w1/w2/w3/b1/b2/b3 (terminating at the encoder boundary; + // structural stop-grad enforced by missing `dx_in` output). + let h_s2_aux_ptr = self.h_s2_aux.raw_ptr(); + let dh_s2_aux_accum_ptr = self.dh_s2_aux_accum.raw_ptr(); + let encoder_out_ptr = self.ptrs.save_h_s1; + let encoder_out_dim = self.config.shared_h1; + let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1; + let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2; let aux_w = self.aux_weight; let aux_h = AUX_HIDDEN_DIM; let aux_kr = AUX_REGIME_K; @@ -17725,9 +17785,12 @@ impl GpuDqnTrainer { // SP13 B1.1a (2026-05-05): softmax CE backward. Reads the saved // softmax tile + i32 labels + B_valid scalar (written by loss // reduce) so loss + gradient share the same `1/B_valid` divisor. + // SP14-C.5b: input is `h_s2_aux` (aux trunk output). The dh_s2_aux + // outputs are written to per-head buffers (overwrite, not + // accumulate); the SAXPY-into-`dh_s2_aux_accum` happens below. self.aux_heads_bwd.backward_next_bar( &self.stream, - h_s2_ptr, + h_s2_aux_ptr, nb_w1, nb_w2, self.aux_nb_hidden_buf.raw_ptr(), self.aux_nb_softmax_buf.raw_ptr(), @@ -17742,7 +17805,7 @@ impl GpuDqnTrainer { )?; self.aux_heads_bwd.backward_regime( &self.stream, - h_s2_ptr, + h_s2_aux_ptr, rg_w1, rg_w2, self.aux_rg_hidden_buf.raw_ptr(), self.aux_rg_logits_buf.raw_ptr(), @@ -17816,18 +17879,32 @@ impl GpuDqnTrainer { } } - // SAXPY both per-head dh_s2 buffers into the trunk accumulator - // `bw_d_h_s2`. Both source and destination already share the [B, SH2] - // shape + unit stride. encoder_backward_chain consumes the augmented - // accumulator transparently through the same pointer. + // SP14 Layer C Phase C.5b (2026-05-08): SAXPY both per-head + // `aux_dh_s2_*_buf` buffers into the AUX TRUNK upstream-grad + // accumulator `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). The aux + // trunk's backward kernel reads `dh_s2_aux_accum` as its + // `dh_s2_aux_in_ptr` and propagates gradient through w3/w2/w1 + + // b3/b2/b1, terminating at the encoder boundary (no `dx_in` + // output param). Pre-zero the accumulator each step so we don't + // carry forward last step's gradient — graph-safe via + // `cuMemsetD32Async` (matches the d_h_s2 zero pattern in the + // backward chain). let n_dh = (b * sh2) as i32; let blocks_dh = ((n_dh as u32 + 255) / 256).max(1); let dh_nb_ptr = self.aux_dh_s2_nb_buf.raw_ptr(); let dh_rg_ptr = self.aux_dh_s2_rg_buf.raw_ptr(); unsafe { + // Pre-zero `dh_s2_aux_accum` so the SAXPY below establishes + // the per-step value rather than accumulating across steps. + cudarc::driver::sys::cuMemsetD32Async( + dh_s2_aux_accum_ptr, + 0, + (b * sh2), + self.stream.cu_stream(), + ); self.stream .launch_builder(&self.saxpy_f32_kernel) - .arg(&bw_d_h_s2_ptr) + .arg(&dh_s2_aux_accum_ptr) .arg(&dh_nb_ptr) .arg(&aux_w) .arg(&n_dh) @@ -17837,11 +17914,11 @@ impl GpuDqnTrainer { shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!( - "aux_heads_backward saxpy dh_s2 next_bar: {e}" + "aux_heads_backward saxpy dh_s2_aux next_bar: {e}" )))?; self.stream .launch_builder(&self.saxpy_f32_kernel) - .arg(&bw_d_h_s2_ptr) + .arg(&dh_s2_aux_accum_ptr) .arg(&dh_rg_ptr) .arg(&aux_w) .arg(&n_dh) @@ -17851,10 +17928,53 @@ impl GpuDqnTrainer { shared_mem_bytes: 0, }) .map_err(|e| MLError::ModelError(format!( - "aux_heads_backward saxpy dh_s2 regime: {e}" + "aux_heads_backward saxpy dh_s2_aux regime: {e}" )))?; } + // SP14 Layer C Phase C.5b: aux trunk backward — propagates the + // accumulated upstream gradient `dh_s2_aux_accum` through the + // 3-layer aux trunk MLP, writing per-tensor gradients into the + // 6 grad buffers (overwrite, not accumulate). Encoder boundary is + // structural stop-grad: the kernel set has NO `dx_in` output, so + // Q's encoder gradient is unaffected. Pearl + // `pearl_no_host_branches_in_captured_graph.md` is preserved — + // all kernels pre-loaded at construction; no host dispatch in + // the captured region. + self.aux_trunk_backward_ops.launch( + &self.stream, + dh_s2_aux_accum_ptr, + encoder_out_ptr, + self.h_aux1.raw_ptr(), + self.h_aux2.raw_ptr(), + self.aux_trunk_w2.raw_ptr(), + self.aux_trunk_w3.raw_ptr(), + self.dh_aux1_pre_scratch.raw_ptr(), + self.dh_aux2_pre_scratch.raw_ptr(), + self.aux_trunk_w1_grad.raw_ptr(), + self.aux_trunk_b1_grad.raw_ptr(), + self.aux_trunk_w2_grad.raw_ptr(), + self.aux_trunk_b2_grad.raw_ptr(), + self.aux_trunk_w3_grad.raw_ptr(), + self.aux_trunk_b3_grad.raw_ptr(), + b, + encoder_out_dim, + aux_h1, + aux_h2, + sh2, + )?; + + // SP14 Layer C Phase C.5b: aux trunk Adam optimiser step — + // consumes the 6 grad tensors via global L2-norm clip (single + // norm spans all 6 tensors) + per-tensor Adam update. ISV-driven + // hyperparams (LR, β1, β2, ε, grad_clip) read inside the launcher; + // pinned LR/clip buffers are populated by the caller pre-capture + // (host-write-through-mapped-pinned). The aux trunk step counter + // `aux_trunk_t_dev_ptr` is incremented by `submit_counter_increments` + // before this launch, ensuring Adam bias correction tracks step + // count consistently across captured graph replays. + self.launch_aux_trunk_adam_update(&self.stream, b as i32)?; + Ok(()) } @@ -29909,13 +30029,49 @@ impl GpuDqnTrainer { self.check_nan_f32(self.ptrs.bw_d_h_s2, b_post * sh2_post, 33)?; } + // ── SP14 Layer C Phase C.5b (2026-05-08): aux trunk Adam pre-capture ── + // Host-write through mapped-pinned to populate Adam hyperparams + + // step counter for the captured aux trunk Adam launch (inside + // `aux_heads_backward` below). The captured graph reads pointers, + // not values — these host writes are picked up at replay time. + // + // Per `pearl_no_host_branches_in_captured_graph.md` we cannot do + // these host writes from inside the captured region, but + // `launch_cublas_backward_to` is `&mut self` and runs OUTSIDE the + // captured graph (the captured chain is each `cuGraphLaunch` + // replay; this Rust function runs each per-step host driver). + // + // ISV[AUX_TRUNK_LR_INDEX=444] / ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448] + // are populated by the StateResetRegistry at fold boundaries + // (`AUX_TRUNK_LR_DEFAULT=1e-4` / `AUX_TRUNK_GRAD_CLIP_DEFAULT=1.0`) + // and unchanged within a fold; reading once per step lets future + // adaptive controllers drive these slots without re-wiring. + { + use crate::cuda_pipeline::sp14_isv_slots::{ + AUX_TRUNK_GRAD_CLIP_INDEX, AUX_TRUNK_LR_INDEX, + }; + let aux_lr = self.read_isv_signal_at(AUX_TRUNK_LR_INDEX); + let aux_clip = self.read_isv_signal_at(AUX_TRUNK_GRAD_CLIP_INDEX); + // Step counter is host-tracked (mirrors `step_ofi_embed_adam` + // pattern): increment then write through to the pinned buffer. + // The captured Adam kernel reads `*t_dev_ptr` at replay time. + self.aux_trunk_adam_step += 1; + unsafe { + *self.aux_trunk_lr_pinned = aux_lr; + *self.aux_trunk_grad_clip_pinned = aux_clip; + *self.aux_trunk_t_pinned = self.aux_trunk_adam_step; + } + } + // ── Plan 4 Task 6 Commit B: aux-heads backward ── // Runs the per-head backward kernels, reduces 8 per-tensor partials // into final dW/dB, SAXPYs them into `grad_base[119..127)` with - // alpha = `aux_weight`, and SAXPYs both per-head dh_s2 buffers into - // the trunk accumulator `bw_d_h_s2` (= `d_h_s2_ptr`). MUST run BEFORE - // `encoder_backward_chain` so the trunk dW/dB chain rule sees the - // augmented dh_s2 (= main + aux contributions). + // alpha = `aux_weight`. SP14-C.5b: SAXPYs both per-head dh_s2_aux + // buffers into `dh_s2_aux_accum` (NOT `bw_d_h_s2`); aux trunk + // backward propagates that gradient to the aux trunk's own + // params; aux trunk Adam launches at the end of the chain. + // Q's `bw_d_h_s2` is unaffected by aux loss — encoder boundary + // structural stop-grad enforced by `aux_trunk_backward`. self.aux_heads_backward(grad_base)?; // SP1 Phase B (slot 34): bw_d_h_s2 snapshot AFTER aux_heads_backward diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index d37648aad..31b7e1a20 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -4737,19 +4737,54 @@ impl GpuExperienceCollector { } } - // Step B: aux next-bar forward. Reads `exp_h_s2_f32` - // (n_episodes × shared_h2) populated by `forward_online_f32` - // above; writes hidden + logits + softmax tiles. The - // softmax tile (`exp_aux_nb_softmax_buf [n_episodes, 2]`) - // is the production output consumed by step 4's SP14 EGF - // chain. Per `pearl_no_host_branches_in_captured_graph` - // this launch is OUTSIDE the captured graph (post + // Step B-1 (SP14-C.5b, 2026-05-08): aux trunk forward. + // Reads encoder layer-1 output `exp_h_s1_f32` + // (n_episodes × shared_h1) and writes the aux trunk + // hidden caches + final `exp_h_s2_aux` (n_episodes × SH2). + // This SEPARATES the aux head input from Q's `exp_h_s2_f32` + // — the aux trunk shapes a representation specifically for + // the next-bar prediction objective without leaking + // gradient back into Q's encoder. Param tensors come from + // the trainer's dedicated aux trunk slabs via + // `aux_trunk_{w,b}{1,2,3}_ptr` (set by + // `set_trainer_aux_trunk_param_ptrs` once the fused + // training context is created). + let enc_out_dim = self.network_dims.0; + let aux_h1 = super::gpu_aux_trunk::AUX_TRUNK_H1; + let aux_h2 = super::gpu_aux_trunk::AUX_TRUNK_H2; + self.exp_aux_trunk_forward_ops.launch( + &self.stream, + self.exp_h_s1_f32.raw_ptr(), + self.aux_trunk_w1_ptr, + self.aux_trunk_b1_ptr, + self.aux_trunk_w2_ptr, + self.aux_trunk_b2_ptr, + self.aux_trunk_w3_ptr, + self.aux_trunk_b3_ptr, + self.exp_h_aux1.raw_ptr(), + self.exp_h_aux2.raw_ptr(), + self.exp_h_s2_aux.raw_ptr(), + n, + enc_out_dim, + aux_h1, + aux_h2, + sh2, + )?; + + // Step B-2: aux next-bar forward. SP14-C.5b: input redirected + // from `exp_h_s2_f32` (Q trunk output) to `exp_h_s2_aux` + // (aux trunk output, written above). Writes hidden + logits + // + softmax tiles. The softmax tile + // (`exp_aux_nb_softmax_buf [n_episodes, 2]`) is the + // production output consumed by step 4's SP14 EGF chain. + // Per `pearl_no_host_branches_in_captured_graph` this + // launch is OUTSIDE the captured graph (post // `cuGraphLaunch`) — same placement as the trainer's // `aux_heads_forward` call relative to the trainer's // captured graph. self.exp_aux_heads_fwd.forward_next_bar( &self.stream, - self.exp_h_s2_f32.raw_ptr(), + self.exp_h_s2_aux.raw_ptr(), nb_w1, nb_b1, nb_w2, nb_b2, n, // batch dim = n_episodes sh2, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 7356ca5e4..6cdb6b224 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -7504,3 +7504,121 @@ With C.5a + C.5a-fixup providing all scaffolding (buffers, scratch, ptrs, access 5. Add aux trunk `t_pinned` to `submit_counters_ops` chain (GPU-side step counter increment). 6. Pre-capture host write of ISV-driven `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned` (cold-path per-epoch). 7. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`. + +## SP14 Layer C Phase C.5b — atomic contract migration COMPLETE (2026-05-08) + +> Atomic contract migration: aux heads consume the SEPARATE aux trunk's +> `h_s2_aux` (NOT Q's GRN trunk `save_h_s2`); aux head backward writes +> `dh_s2_aux` into a NEW accumulator `dh_s2_aux_accum`; aux trunk +> backward propagates that gradient through w3/w2/w1 + b3/b2/b1, and aux +> trunk Adam updates 6 grad tensors via global L2-norm clip. Encoder +> boundary is structural stop-grad (kernel set has NO `dx_in` output). + +### Reverted commits (stop-grad band-aids retired) + +| Commit | What it did | C.5b reversal | +|---------|-----------------------------------------------------------|---------------------------------------------------------| +| `872bd7392` | Zero-fill `dh_s2_out` in `aux_next_bar_backward` (Step 3) | Restored genuine SAXPY: `dh_s2_aux[b,j] = sum_k sh_dh_pre[k] * w1[k,j]` | +| `411a30473` | Same zero-fill in `aux_regime_backward` (Step 3) | Symmetric restoration with same gradient formula | + +The stop-grad audit blocks (lines ~599-638 next-bar + ~757-790 regime) +deleted; replaced by short C.5b comments documenting the atomic +contract flip. + +### Kernel signature renames (4 functions) + +| Kernel | Old param | New param | +|-------------------------|-------------------|-------------------| +| `aux_next_bar_forward` | `const float* h_s2` | `const float* h_s2_aux` | +| `aux_regime_forward` | `const float* h_s2` | `const float* h_s2_aux` | +| `aux_next_bar_backward` | `float* dh_s2_out` | `float* dh_s2_aux_out` | +| `aux_regime_backward` | `float* dh_s2_out` | `float* dh_s2_aux_out` | + +The Rust `gpu_aux_heads.rs` wrapper renamed call-site argument names to +match (`h_s2_ptr → h_s2_aux_ptr`, `dh_s2_out_ptr → dh_s2_aux_out_ptr`) +so the param-name → kernel-arg mapping stays self-documenting at +launch sites. + +### Wire sites lit (4) + +| Site | Function | Insert point | What it does | +|------|----------|--------------|--------------| +| Trainer fwd | `aux_heads_forward` | New Step 0, before regime label builder | `aux_trunk_forward_ops.launch(... encoder_out=save_h_s1, h_s2_aux_out=h_s2_aux)` populates the aux trunk's representation | +| Trainer bwd | `aux_heads_backward` | After per-head SAXPY into `dh_s2_aux_accum` | `aux_trunk_backward_ops.launch(... dh_s2_aux_in=accum, w/b grads...)` propagates grad through aux trunk; then `launch_aux_trunk_adam_update` applies the optimizer step | +| Trainer bwd | `aux_heads_backward` | After both head backwards | Pre-zero `dh_s2_aux_accum` via `cuMemsetD32Async` then SAXPY both `aux_dh_s2_*_buf` into it (REPLACES the old SAXPY into `bw_d_h_s2`) | +| Collector fwd | `submit_dqn_step_loop_cublas` step 2b | Post-`forward_online_f32`, before `forward_next_bar` | `exp_aux_trunk_forward_ops.launch(... encoder_out=exp_h_s1_f32, h_s2_aux_out=exp_h_s2_aux)` then `forward_next_bar` consumes `exp_h_s2_aux` (NOT `exp_h_s2_f32`) | + +### Pre-capture host writes (per-step, `&mut self`) + +In `launch_cublas_backward_to` (BEFORE `aux_heads_backward(grad_base)`): + +```rust +let aux_lr = self.read_isv_signal_at(AUX_TRUNK_LR_INDEX); // ISV[444] +let aux_clip = self.read_isv_signal_at(AUX_TRUNK_GRAD_CLIP_INDEX); // ISV[448] +self.aux_trunk_adam_step += 1; +unsafe { + *self.aux_trunk_lr_pinned = aux_lr; + *self.aux_trunk_grad_clip_pinned = aux_clip; + *self.aux_trunk_t_pinned = self.aux_trunk_adam_step; +} +``` + +These writes are SAFE in graph-capture context: `launch_cublas_backward_to` +is `&mut self` and runs OUTSIDE the captured region (the captured +chain is each `cuGraphLaunch`-replay; this Rust function is the +per-step host driver). The Adam kernel inside the captured graph +reads the pinned device pointer, which resolves to the host's +mapped-pinned memory at replay time — picking up whatever value the +host last wrote. Same pattern as `step_ofi_embed_adam`. + +The step counter is host-tracked (mirrors OFI embed) rather than +extending `increment_step_counters` because the aux trunk lives +outside the SP4 8-group taxonomy and a per-host-call increment is +sufficient (no coupling to other Adam schedulers' bias correction). + +### SAXPY redirect — straightforward + +The previous `aux_heads_backward` SAXPYed `aux_dh_s2_nb_buf` + +`aux_dh_s2_rg_buf` into `bw_d_h_s2_ptr`. C.5b changes the destination +ptr from `bw_d_h_s2_ptr` → `dh_s2_aux_accum_ptr`. The 256-byte SAXPY +launch geometry is identical — only the destination buffer changes. +A `cuMemsetD32Async` pre-zeroes `dh_s2_aux_accum` each step so the +SAXPY establishes per-step value (mirrors the existing `d_h_s2_ptr` +zero pattern in the backward chain). + +### Encoder out pointer + +Both trainer + collector use `save_h_s1` / `exp_h_s1_f32` (config.shared_h1 += 256) as the aux trunk's input — matches the aux trunk's +`ENCODER_OUT_DIM = 256` allocation from C.2. The trainer reads via +`self.ptrs.save_h_s1`; the collector reads via `self.exp_h_s1_f32.raw_ptr()`. +Both buffers are populated by their respective `forward_online` calls +BEFORE the aux trunk forward runs. + +### Verification + +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests` — clean (1m02s, only pre-existing warnings). +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14). The aux_trunk gradient check still validates within tolerance (max_rel_err=1.33e-2 vs 2e-2 tol), and `aux_trunk_backward_does_not_write_dx` still passes (kernel source clean of `dx_in`/`dx_in_out`). + +### Wire status — POST-C.5b (production-ready) + +| Component | Status | +|-----------|--------| +| `aux_trunk_forward_ops` (trainer) | LIVE — called from `aux_heads_forward` before label builder | +| `aux_trunk_backward_ops` (trainer) | LIVE — called from `aux_heads_backward` post-head-SAXPY | +| `launch_aux_trunk_adam_update` (trainer) | LIVE — called from `aux_heads_backward` post-trunk-backward | +| `dh_aux1_pre_scratch` / `dh_aux2_pre_scratch` (trainer) | LIVE — consumed by `aux_trunk_backward.launch` | +| `dh_s2_aux_accum` (trainer) | LIVE — pre-zeroed + SAXPY-target each step | +| `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned` (trainer) | LIVE — host-written from ISV per step | +| `aux_trunk_t_pinned` (trainer) | LIVE — host-incremented per step (OFI embed pattern) | +| `exp_aux_trunk_forward_ops` (collector) | LIVE — called post-`forward_online_f32` before next-bar forward | +| `exp_h_s2_aux` / `exp_h_aux1` / `exp_h_aux2` (collector) | LIVE — populated by `aux_trunk_forward.launch` | + +### Outstanding for C.6+ + +C.5b only wires the trunk forward+backward+Adam mechanics. Phases C.6+ +(separate commits) layer on: +- C.6: `h_s2_aux_rms_ema` producer kernel feeding ISV[449]. +- C.7-C.8: ISV-driven aux trunk Adam β1/β2/ε. +- C.9: synthetic-data smoke + audit close-out + memory pearls. +- C.10: L40S 30-epoch validation. diff --git a/docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md b/docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md index 51363231d..04af6992d 100644 --- a/docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md +++ b/docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md @@ -962,9 +962,115 @@ EOF DO NOT push. +## Phase C.5a-fixup: Add Missing Scaffolding (DEAD CODE) + +> Authorized 2026-05-08 after C.5b implementer flagged 5 specific gaps (most importantly a graph-capture hazard) plus 2 missing buffer allocations from C.5a. C.5a-fixup is **additive-only** — same discipline as C.5a. Goal: by end of fixup, C.5b is a genuine ~300 LOC atomic flip with zero scaffolding work mixed in. + +### Task C.5a-fixup.1: Allocate missing scratch buffers + collector ptr scaffolding + graph-capture audit + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (allocate `dh_aux1_pre_scratch`, `dh_aux2_pre_scratch`; add `aux_trunk_param_ptrs()` accessor) +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (6× `u64` aux_trunk param ptr fields + `set_trainer_aux_trunk_param_ptrs` setter; new `exp_aux_trunk_forward_ops: AuxTrunkForwardOps` field; constructor populates it; **leave wire sites still calling `save_h_s2` — no semantic change**) +- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (call new `set_trainer_aux_trunk_param_ptrs` setter at the existing setup point near `set_trainer_params_ptr` — ~lines 1703 + 2317 per implementer report) +- Modify: `docs/dqn-wire-up-audit.md` (audit entry + graph-capture verification result) + +- [ ] **Step 1: Allocate missing scratch buffers in trainer** + +`aux_trunk_backward.launch` requires `dh_aux1_pre_scratch_ptr [B_max, AUX_TRUNK_H1=256]` and `dh_aux2_pre_scratch_ptr [B_max, AUX_TRUNK_H2=128]` per `gpu_aux_trunk.rs:266-267`. C.5a missed these. Add as `CudaSlice` zero-init alongside existing C.5a buffers. + +- [ ] **Step 2: Collector aux trunk ptr scaffolding** (mirrors existing `set_trainer_params_ptr` pattern — find via grep) + +```rust +// In GpuExperienceCollector struct: +aux_trunk_w1_ptr: u64, +aux_trunk_b1_ptr: u64, +aux_trunk_w2_ptr: u64, +aux_trunk_b2_ptr: u64, +aux_trunk_w3_ptr: u64, +aux_trunk_b3_ptr: u64, +exp_aux_trunk_forward_ops: AuxTrunkForwardOps, + +// Setter: +pub fn set_trainer_aux_trunk_param_ptrs(&mut self, w1: u64, b1: u64, w2: u64, b2: u64, w3: u64, b3: u64) { + self.aux_trunk_w1_ptr = w1; + // ... etc +} +``` + +- [ ] **Step 3: Trainer-side accessor** + +```rust +// In GpuDqnTrainer: +pub fn aux_trunk_param_ptrs(&self) -> (u64, u64, u64, u64, u64, u64) { + use cudarc::driver::DevicePtr; + ( + *self.aux_trunk_w1.device_ptr(), + *self.aux_trunk_b1.device_ptr(), + *self.aux_trunk_w2.device_ptr(), + *self.aux_trunk_b2.device_ptr(), + *self.aux_trunk_w3.device_ptr(), + *self.aux_trunk_b3.device_ptr(), + ) +} +``` + +(Adjust to match the actual cudarc `device_ptr()` API — find an existing accessor like `embed_w_ptr()` and mirror.) + +- [ ] **Step 4: Wire setter call in training_loop** + +Find the existing call to `set_trainer_params_ptr` (~lines 1703 + 2317 per implementer report). Insert call to `set_trainer_aux_trunk_param_ptrs` at the same site, passing the trainer's accessor result. + +- [ ] **Step 5: Graph-capture audit** (the showstopper concern) + +Read `gpu_dqn_trainer.rs` aux_heads_backward orchestrator (around line 29865 per implementer report). Trace: does the call sit INSIDE a `cuStreamBeginCapture`/`end_capture` region? + +If INSIDE: the existing aux backward CANNOT have host-side writes (step counter, pinned buffer). Verify the existing code by inspecting it — the existing aux_heads_backward should have a graph-capture-compatible structure (probably a separate function called pre-capture for any host writes). Document the existing pattern. + +If OUTSIDE: the new aux_trunk Adam launch can sit alongside existing aux_heads_backward; host writes are fine. + +Document the verification result in audit doc — this determines C.5b's wiring strategy. + +- [ ] **Step 6: Verify** + +```bash +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets 2>&1 | tail -10 +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture 2>&1 | tail -15 +``` + +Expected: clean build; 12 tests still pass (no regression — pure scaffolding). + +- [ ] **Step 7: Commit** + +```bash +git add -A +git commit -m "$(cat <<'EOF' +feat(sp14-c.5a-fixup): missing scratch buffers + collector ptr scaffolding (dead code) + +Phase C.5a-fixup — completes C.5a's additive infrastructure: +- Allocate dh_aux1_pre_scratch + dh_aux2_pre_scratch (missed in C.5a) +- Collector struct gains 6× aux_trunk param ptr fields + setter +- Trainer gains aux_trunk_param_ptrs() accessor +- training_loop wires the setter at the existing set_trainer_params_ptr site +- Collector constructor populates AuxTrunkForwardOps + +NO contract change. Wire sites still call save_h_s2. The setter and ptrs +are populated but unused by aux head consumers (still wired through Q's +trunk). + +Graph-capture audit: + +Co-Authored-By: Claude Opus 4.7 (1M context) +EOF +)" +``` + +DO NOT push. + +--- + ## Phase C.5b: Atomic Contract Migration — Aux Head Input Switches h_s2 → h_s2_aux -> ONE atomic commit per `feedback_no_partial_refactor`. ~300 LOC. THIS is the genuine contract migration. C.5a's additive code becomes live here. +> ONE atomic commit per `feedback_no_partial_refactor`. With C.5a + C.5a-fixup providing all scaffolding (buffers, ptrs, accessor, setter, AuxTrunkForwardOps in collector), C.5b is now a genuine ~300 LOC atomic flip: 4 kernel param renames + 2 zero-fill reverts + 4 wire sites + Adam launch + ISV pinned-buffer population. ### Task C.5b.1: Atomic kernel rename + revert zero-fills + wire-up