From c90de98594f26f9a6b1f07ee315e8ed7a82c563f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 8 May 2026 02:10:17 +0200 Subject: [PATCH] feat(sp14-c.5a): allocate aux trunk fwd/bwd buffers + Adam launcher (dead code) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase C.5a — additive infrastructure for the aux trunk wire-up. Allocates saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), gradient buffers (6× aux_trunk_*_grad), accumulator (dh_s2_aux_accum), and dedicated Adam launcher (launch_aux_trunk_adam_update) reading β1/β2/ε/LR/grad-clip from ISV[444..449). No contract change. No call sites for the new launcher yet. C.5b atomically wires these in. Phase C.5 was split (authorized 2026-05-08) after the original implementer flagged ~600-800 LOC scope across 4 files with correctness windows. C.5a is purely additive; C.5b is the genuine ~300 LOC atomic migration. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 492 ++++++++++++++++++ .../cuda_pipeline/gpu_experience_collector.rs | 56 ++ docs/dqn-wire-up-audit.md | 99 ++++ ...6-05-07-sp14-layer-c-separate-aux-trunk.md | 110 +++- 4 files changed, 750 insertions(+), 7 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 28b4cc5bb..842e4dd84 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -7327,6 +7327,114 @@ pub struct GpuDqnTrainer { /// Adam second moment for `aux_trunk_b3` [SH2]. aux_trunk_b3_v: CudaSlice, + // ── SP14 Layer C Phase C.5a saved-fwd + grad buffers (2026-05-08) ── + // Allocated as DEAD CODE: these buffers are real device allocations + // but no production caller reads / writes them yet. Phase C.5b + // atomically wires them via the aux trunk forward (writes + // h_s2_aux/h_aux1/h_aux2), aux heads backward (SAXPY-accumulates into + // dh_s2_aux_accum), aux trunk backward (overwrites the 6 grad + // tensors), and `launch_aux_trunk_adam_update` (consumes grads via + // global L2-norm clip + per-tensor Adam step). + // + // 6-buffer choice (vs flat-buffer): mirrors C.2's separate-tensor + // params/m/v allocation. Lets each Adam launch consume one + // (param, grad, m, v) tuple of matching extent — no byte-offset + // arithmetic. Per-tensor grad_norm partial sums into a unified + // `aux_trunk_norm_partials` buffer at fixed offsets so the global + // L2 norm spans all 6 tensors. + /// `[B_max, SH2]` final aux trunk output (saved post-fwd; consumed by + /// aux heads forward + backward in C.5b). Same shape contract as + /// `aux_dh_s2_nb_buf` so the SAXPY accumulator type-matches. + #[allow(dead_code)] + h_s2_aux: CudaSlice, + /// `[B_max, AUX_TRUNK_H1=256]` Layer-1 post-ELU activation. Consumed + /// by `aux_trunk_bwd_dh_pre` (`ELU'(h_aux1)`) + `aux_trunk_bwd_dW_reduce` + /// (Layer-2 `dW2 = h_aux1.T @ dh_aux2_pre`). + #[allow(dead_code)] + h_aux1: CudaSlice, + /// `[B_max, AUX_TRUNK_H2=128]` Layer-2 post-ELU activation. Consumed + /// by `aux_trunk_bwd_dh_pre` (`ELU'(h_aux2)`) + `aux_trunk_bwd_dW_reduce` + /// (Layer-3 `dW3 = h_aux2.T @ d_logits`). + #[allow(dead_code)] + h_aux2: CudaSlice, + /// `[B_max, SH2]` SAXPY accumulator for upstream gradient feeding the + /// aux trunk's Layer-3 input (`dh_s2_aux`). C.5b reverts the zero-fill + /// hack at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` + /// so the two aux head backwards SAXPY their per-sample dh into this + /// buffer instead of the trunk's `bw_d_h_s2`. Consumed by aux trunk + /// backward (`launch(... dh_s2_aux_in_ptr=accum_ptr ...)`). + #[allow(dead_code)] + dh_s2_aux_accum: CudaSlice, + + /// `[SH1 × 256]` aux trunk weight-1 gradient accumulator (overwritten + /// per training step by `aux_trunk_bwd_dW_reduce`). + #[allow(dead_code)] + aux_trunk_w1_grad: CudaSlice, + /// `[256]` aux trunk bias-1 gradient (overwritten by `aux_trunk_bwd_db_reduce`). + #[allow(dead_code)] + aux_trunk_b1_grad: CudaSlice, + /// `[256 × 128]` aux trunk weight-2 gradient. + #[allow(dead_code)] + aux_trunk_w2_grad: CudaSlice, + /// `[128]` aux trunk bias-2 gradient. + #[allow(dead_code)] + aux_trunk_b2_grad: CudaSlice, + /// `[128 × SH2]` aux trunk weight-3 gradient. + #[allow(dead_code)] + aux_trunk_w3_grad: CudaSlice, + /// `[SH2]` aux trunk bias-3 gradient. + #[allow(dead_code)] + aux_trunk_b3_grad: CudaSlice, + + /// `[num_blocks_total]` per-block partial sum-of-squares written by + /// `dqn_grad_norm_kernel` across all 6 grad tensors (offsets matching + /// `aux_trunk_grad_block_offsets`). Sized to the maximum total block + /// count any of the 6 grads would produce. + #[allow(dead_code)] + aux_trunk_grad_norm_partials: CudaSlice, + /// Per-tensor block offsets into `aux_trunk_grad_norm_partials`. The + /// 7th entry is the total block count (passed as `num_blocks` to + /// `dqn_grad_norm_finalize`). Captured at construction so the launcher + /// has no host-side branching during its launches. + #[allow(dead_code)] + aux_trunk_grad_block_offsets: [usize; 7], + /// `[1]` finalised L2 grad norm written by `dqn_grad_norm_finalize`. + /// Read by 6 Adam launches as the `grad_norm_f32` arg. + #[allow(dead_code)] + aux_trunk_grad_norm_buf: CudaSlice, + + /// Pinned device-mapped grad-clip threshold. Host writes + /// `ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]` per-step before launch; GPU + /// reads via `aux_trunk_grad_clip_dev_ptr`. The Adam kernel's + /// `max_grad_norm_buf` arg. + #[allow(dead_code)] + aux_trunk_grad_clip_pinned: *mut f32, + /// Device pointer for `aux_trunk_grad_clip_pinned` (one + /// `cuMemHostGetDevicePointer_v2` per construction; reusable across + /// launches inside captured graphs). + #[allow(dead_code)] + aux_trunk_grad_clip_dev_ptr: u64, + + /// Pinned device-mapped LR. Host writes `ISV[AUX_TRUNK_LR_INDEX=444]` + /// per-step. The Adam kernel's `lr_ptr` arg. + #[allow(dead_code)] + aux_trunk_lr_pinned: *mut f32, + /// Device pointer for `aux_trunk_lr_pinned`. + #[allow(dead_code)] + aux_trunk_lr_dev_ptr: u64, + + /// Pinned device-mapped Adam step counter (i32). Incremented by host + /// before each launch; the kernel's `t_ptr` arg. + #[allow(dead_code)] + aux_trunk_t_pinned: *mut i32, + /// Device pointer for `aux_trunk_t_pinned`. + #[allow(dead_code)] + aux_trunk_t_dev_ptr: u64, + /// Host-side mirror of the Adam step counter. Production caller in + /// C.5b increments this and writes through to the pinned buffer. + #[allow(dead_code)] + aux_trunk_adam_step: i32, + /// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. /// Owns the pre-loaded `CudaFunction` handle for `aux_trunk_forward` /// per `pearl_no_host_branches_in_captured_graph.md` — the launch @@ -7911,6 +8019,16 @@ impl Drop for GpuDqnTrainer { if !self.ofi_embed_t_pinned.is_null() { let _ = unsafe { cudarc::driver::result::free_host(self.ofi_embed_t_pinned.cast()) }; } + // SP14 Layer C Phase C.5a (2026-05-08): aux trunk Adam pinned ptrs. + if !self.aux_trunk_grad_clip_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_grad_clip_pinned.cast()) }; + } + if !self.aux_trunk_lr_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_lr_pinned.cast()) }; + } + if !self.aux_trunk_t_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.aux_trunk_t_pinned.cast()) }; + } } } @@ -9691,6 +9809,244 @@ impl GpuDqnTrainer { Ok(()) } + /// SP14 Layer C Phase C.5a (2026-05-08): aux trunk Adam optimizer step. + /// + /// **DEAD CODE.** No production caller invokes this in C.5a — the + /// launcher exists so the build stays green while C.5b is the + /// genuine atomic contract migration that wires the aux head + /// backward → `dh_s2_aux_accum` → aux trunk backward → THIS launcher + /// chain. Per `feedback_no_partial_refactor.md` the wire-up is one + /// commit. Per `feedback_no_stubs.md` the body is fully real (no + /// placeholder panic); a synthetic call from a unit test would + /// execute end-to-end correctly given populated grad buffers. + /// + /// Reads ISV-driven hyper-parameters (no hardcoded constants): + /// + /// | ISV slot | Use | + /// |----------------------------------|---------------------------| + /// | `AUX_TRUNK_LR_INDEX = 444` | Adam learning rate | + /// | `AUX_TRUNK_BETA1_INDEX = 445` | Adam β1 (first moment) | + /// | `AUX_TRUNK_BETA2_INDEX = 446` | Adam β2 (second moment) | + /// | `AUX_TRUNK_EPS_INDEX = 447` | Adam ε (denom stability) | + /// | `AUX_TRUNK_GRAD_CLIP_INDEX = 448`| Global L2 clip threshold | + /// + /// Mirrors the OFI embed Adam pattern (`step_ofi_embed_adam`) but + /// over 6 separate (params, grad, m, v) tensor tuples instead of one + /// flat buffer. The aux trunk lives outside the SP4 8-group + /// taxonomy (Q/Target/CQL/Aux per-branch + auxiliaries) so a + /// dedicated launcher is appropriate; SP4 Mech 9 weight clamp + /// (`weight_clamp_max_abs`) and Pearl C engagement counter + /// (`engage_buf_offset`) are disabled — same convention as OFI + /// embed. + /// + /// Phase 1: per-tensor partial sum-of-squares into the unified + /// `aux_trunk_grad_norm_partials` buffer at fixed offsets + /// (`aux_trunk_grad_block_offsets[i..i+1]` slice covers tensor `i`'s + /// blocks). + /// + /// Phase 2: a single `dqn_grad_norm_finalize` reduces all blocks to + /// `aux_trunk_grad_norm_buf[0] = sqrt(sum_b partial_b)`. + /// + /// Phase 3: 6 `dqn_adam_update_kernel` launches, each consuming one + /// (params, grad, m, v) tuple. The kernel's internal + /// `clip_scale = (norm > clip) ? clip/norm : 1` divides by the + /// SAME global norm (read by every launch), so clipping is + /// consistent across tensors. + /// + /// `batch_size` is unused by this launcher (Adam is per-parameter, + /// not per-batch). Kept in the signature so call-site code reads + /// uniformly with other Adam launchers in this trainer. + /// + /// Per `pearl_no_host_branches_in_captured_graph.md` the launcher + /// only touches pre-loaded `CudaFunction` handles + /// (`grad_norm_standalone_post_aux`, `grad_norm_finalize_post_aux`, + /// `adam_update_post_aux`) and pre-allocated buffers — no + /// `load_cubin` / `load_function` on the hot path. + #[allow(dead_code)] + pub(crate) fn launch_aux_trunk_adam_update( + &self, + stream: &Arc, + _batch_size: i32, + ) -> Result<(), MLError> { + use crate::cuda_pipeline::sp14_isv_slots::{ + AUX_TRUNK_BETA1_INDEX, AUX_TRUNK_BETA2_INDEX, AUX_TRUNK_EPS_INDEX, + }; + + // Adam hparams — ISV-driven with cold-start floors per + // `pearl_first_observation_bootstrap.md` numerical-stability + // carve-outs (β1≥0.5 avoids momentum inversion; β2≥0.9 prevents + // exploding second-moment; ε≥1e-12 avoids div-by-zero). + let beta1 = self + .read_isv_signal_at(AUX_TRUNK_BETA1_INDEX) + .max(0.5_f32) + .min(0.9999_f32); + let beta2 = self + .read_isv_signal_at(AUX_TRUNK_BETA2_INDEX) + .max(0.9_f32) + .min(0.99999_f32); + let epsilon = self.read_isv_signal_at(AUX_TRUNK_EPS_INDEX).max(1e-12_f32); + let weight_decay: f32 = 0.0; // aux trunk is outside AdamW WD policy + + // Per-tensor (params, grad, m, v, extent) tuples in the same + // order as `aux_trunk_grad_block_offsets`. + let tensors: [(u64, u64, u64, u64, usize); 6] = [ + ( + self.aux_trunk_w1.raw_ptr(), + self.aux_trunk_w1_grad.raw_ptr(), + self.aux_trunk_w1_m.raw_ptr(), + self.aux_trunk_w1_v.raw_ptr(), + self.aux_trunk_w1.len(), + ), + ( + self.aux_trunk_b1.raw_ptr(), + self.aux_trunk_b1_grad.raw_ptr(), + self.aux_trunk_b1_m.raw_ptr(), + self.aux_trunk_b1_v.raw_ptr(), + self.aux_trunk_b1.len(), + ), + ( + self.aux_trunk_w2.raw_ptr(), + self.aux_trunk_w2_grad.raw_ptr(), + self.aux_trunk_w2_m.raw_ptr(), + self.aux_trunk_w2_v.raw_ptr(), + self.aux_trunk_w2.len(), + ), + ( + self.aux_trunk_b2.raw_ptr(), + self.aux_trunk_b2_grad.raw_ptr(), + self.aux_trunk_b2_m.raw_ptr(), + self.aux_trunk_b2_v.raw_ptr(), + self.aux_trunk_b2.len(), + ), + ( + self.aux_trunk_w3.raw_ptr(), + self.aux_trunk_w3_grad.raw_ptr(), + self.aux_trunk_w3_m.raw_ptr(), + self.aux_trunk_w3_v.raw_ptr(), + self.aux_trunk_w3.len(), + ), + ( + self.aux_trunk_b3.raw_ptr(), + self.aux_trunk_b3_grad.raw_ptr(), + self.aux_trunk_b3_m.raw_ptr(), + self.aux_trunk_b3_v.raw_ptr(), + self.aux_trunk_b3.len(), + ), + ]; + + let partials_base = self.aux_trunk_grad_norm_partials.raw_ptr(); + let f32_sz = std::mem::size_of::() as u64; + + // ── Phase 1: per-tensor partial sum-of-squares ──────────────── + // Each launch writes `extent.div_ceil(256)` blocks into the + // partials buffer at the per-tensor offset. + for (i, &(_p, grad_ptr, _m, _v, extent)) in tensors.iter().enumerate() { + let blocks = extent.div_ceil(256).max(1) as u32; + let off = self.aux_trunk_grad_block_offsets[i] as u64; + let slot_ptr = partials_base + off * f32_sz; + let n_i32 = extent as i32; + unsafe { + stream + .launch_builder(&self.grad_norm_standalone_post_aux) + .arg(&grad_ptr) + .arg(&slot_ptr) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| { + MLError::ModelError(format!( + "aux_trunk grad_norm[tensor={}, blocks={}]: {e}", + i, blocks + )) + })?; + } + } + + // ── Phase 2: finalize across all 6 tensors' blocks ──────────── + let total_blocks = self.aux_trunk_grad_block_offsets[6] as i32; + let norm_out_ptr = self.aux_trunk_grad_norm_buf.raw_ptr(); + unsafe { + stream + .launch_builder(&self.grad_norm_finalize_post_aux) + .arg(&partials_base) + .arg(&norm_out_ptr) + .arg(&total_blocks) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| { + MLError::ModelError(format!("aux_trunk grad_norm_finalize: {e}")) + })?; + } + + // ── Phase 3: per-tensor Adam updates ────────────────────────── + // All 6 launches consume the SAME `norm_out_ptr` so the kernel's + // internal `clip_scale = (norm > max_grad_norm) ? max_grad_norm/norm : 1` + // applies a globally consistent scale. `aux_trunk_grad_clip_dev_ptr` + // is the ISV-driven threshold (host writes per-step). + let lr_ptr = self.aux_trunk_lr_dev_ptr; + let clip_ptr = self.aux_trunk_grad_clip_dev_ptr; + let t_ptr = self.aux_trunk_t_dev_ptr; + let wd_mask_ptr = self.weight_decay_mask.raw_ptr(); + let l1_end: i32 = 0; + let l1_lambda: f32 = 0.0; + let weight_clamp_max_abs: f32 = 0.0; // SP4 outside-taxonomy: disabled + // SP4 Task A14: aux trunk lives outside the 8-group taxonomy — + // pass null nan_flags + disabled engage offset (mirrors OFI embed). + let nan_flags_null: u64 = 0; + let diag_slot_disabled: i32 = -1; + let engage_buf_null: u64 = 0; + let engage_off_disabled: i32 = SP4_ENGAGE_OFFSET_DISABLED; + + for (i, &(params_ptr, grad_ptr, m_ptr, v_ptr, extent)) in tensors.iter().enumerate() { + let n_i32 = extent as i32; + let blocks = (extent.div_ceil(256)).max(1) as u32; + unsafe { + stream + .launch_builder(&self.adam_update_post_aux) + .arg(¶ms_ptr) + .arg(&grad_ptr) + .arg(&m_ptr) + .arg(&v_ptr) + .arg(&norm_out_ptr) // global L2 norm (shared across all 6) + .arg(&lr_ptr) + .arg(&beta1) + .arg(&beta2) + .arg(&epsilon) + .arg(&weight_decay) + .arg(&clip_ptr) // ISV-driven global clip threshold + .arg(&t_ptr) + .arg(&n_i32) + .arg(&wd_mask_ptr) // unused (wd=0); kernel still reads — pass real ptr + .arg(&l1_end) + .arg(&l1_lambda) + .arg(&weight_clamp_max_abs) + .arg(&nan_flags_null) + .arg(&diag_slot_disabled) + .arg(&engage_buf_null) + .arg(&engage_off_disabled) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| { + MLError::ModelError(format!( + "aux_trunk adam_update[tensor={}, n={}]: {e}", + i, extent + )) + })?; + } + } + + Ok(()) + } + /// Recursive confidence backward: MSE loss gradient into trunk + conf weight gradients. /// Accumulates into grad_buf (same buffer Adam reads) and bw_d_h_s2 trunk gradient. pub(crate) fn launch_recursive_confidence_backward(&self, batch_size: usize) -> Result<(), MLError> { @@ -22282,6 +22638,121 @@ impl GpuDqnTrainer { aux_trunk_in_dim, aux_trunk_h1, aux_trunk_h2, aux_trunk_out_dim, aux_trunk_total_params, ); + // ── SP14 Layer C Phase C.5a additive infrastructure (2026-05-08) ── + // Saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), upstream-grad + // accumulator (dh_s2_aux_accum), 6 grad tensors mirroring C.2 + // params, plus grad-norm scratch + ISV-driven LR/clip pinned + // pointers + Adam step counter. ALL DEAD until C.5b atomically + // wires the aux head backward → accumulator → aux trunk backward + // → Adam launcher chain. No production caller reads these yet. + let h_s2_aux = alloc_f32(&stream, b * aux_trunk_out_dim, "h_s2_aux")?; + let h_aux1 = alloc_f32(&stream, b * aux_trunk_h1, "h_aux1")?; + let h_aux2 = alloc_f32(&stream, b * aux_trunk_h2, "h_aux2")?; + let dh_s2_aux_accum = alloc_f32(&stream, b * aux_trunk_out_dim, "dh_s2_aux_accum")?; + + let aux_trunk_w1_grad = alloc_f32(&stream, aux_trunk_w1_count, "aux_trunk_w1_grad")?; + let aux_trunk_b1_grad = alloc_f32(&stream, aux_trunk_h1, "aux_trunk_b1_grad")?; + let aux_trunk_w2_grad = alloc_f32(&stream, aux_trunk_w2_count, "aux_trunk_w2_grad")?; + let aux_trunk_b2_grad = alloc_f32(&stream, aux_trunk_h2, "aux_trunk_b2_grad")?; + let aux_trunk_w3_grad = alloc_f32(&stream, aux_trunk_w3_count, "aux_trunk_w3_grad")?; + let aux_trunk_b3_grad = alloc_f32(&stream, aux_trunk_out_dim, "aux_trunk_b3_grad")?; + + // Per-tensor block counts for the dqn_grad_norm_kernel launches + // (block_dim=256). One partial slot per (block × tensor) pair — + // we lay them out in a contiguous buffer with fixed offsets so + // the finalize kernel sees them as a single `[total_blocks]` + // array. Order matches the 6 (param, grad, m, v) tuples consumed + // by the Adam launcher: w1, b1, w2, b2, w3, b3. + let aux_trunk_grad_extents: [usize; 6] = [ + aux_trunk_w1_count, + aux_trunk_h1, + aux_trunk_w2_count, + aux_trunk_h2, + aux_trunk_w3_count, + aux_trunk_out_dim, + ]; + let mut aux_trunk_grad_block_offsets: [usize; 7] = [0; 7]; + let mut total_grad_blocks: usize = 0; + for (i, &ext) in aux_trunk_grad_extents.iter().enumerate() { + aux_trunk_grad_block_offsets[i] = total_grad_blocks; + let blocks = ext.div_ceil(256); + total_grad_blocks += blocks; + } + aux_trunk_grad_block_offsets[6] = total_grad_blocks; + let aux_trunk_grad_norm_partials = + alloc_f32(&stream, total_grad_blocks.max(1), "aux_trunk_grad_norm_partials")?; + let aux_trunk_grad_norm_buf = + alloc_f32(&stream, 1, "aux_trunk_grad_norm_buf")?; + + // Pinned device-mapped grad-clip threshold. Host writes + // ISV[AUX_TRUNK_GRAD_CLIP_INDEX] before each launch; the Adam + // kernel reads via the device pointer. Cold-start value zero + // means clip disabled (kernel's `(norm > clip) ? clip/norm : 1.0` + // path; with clip=0 the second branch fires, dividing by zero — + // C.5b is responsible for writing a real ISV value before the + // first launch goes live). + let aux_trunk_grad_clip_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned aux_trunk_grad_clip alloc: {e}")))? + as *mut f32 + }; + unsafe { *aux_trunk_grad_clip_pinned = 0.0_f32; } + let aux_trunk_grad_clip_dev_ptr = unsafe { + let mut dp: u64 = 0; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dp as *mut u64, + aux_trunk_grad_clip_pinned.cast(), + 0, + ); + dp + }; + + // Pinned device-mapped LR (host writes ISV[AUX_TRUNK_LR_INDEX] + // each step). Cold-start zero — C.5b initialises before first + // launch. + let aux_trunk_lr_pinned: *mut f32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned aux_trunk_lr alloc: {e}")))? + as *mut f32 + }; + unsafe { *aux_trunk_lr_pinned = 0.0_f32; } + let aux_trunk_lr_dev_ptr = unsafe { + let mut dp: u64 = 0; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dp as *mut u64, + aux_trunk_lr_pinned.cast(), + 0, + ); + dp + }; + + // Pinned device-mapped Adam step counter. Host increments before + // each Adam launch; the kernel's bias-correction reads `*t_ptr`. + let aux_trunk_t_pinned: *mut i32 = unsafe { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + cudarc::driver::result::malloc_host(std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned aux_trunk_t alloc: {e}")))? + as *mut i32 + }; + unsafe { *aux_trunk_t_pinned = 0; } + let aux_trunk_t_dev_ptr = unsafe { + let mut dp: u64 = 0; + cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dp as *mut u64, + aux_trunk_t_pinned.cast(), + 0, + ); + dp + }; + info!( + "SP14-C.5a: aux trunk fwd/bwd buffers + Adam infra allocated \ + (h_s2_aux/h_aux1/h_aux2 + dh_s2_aux_accum, 6 grad tensors, \ + {} grad-norm partial slots, ISV-driven LR/clip + step counter; DEAD until C.5b)", + total_grad_blocks + ); + // SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. // Pre-loads the `aux_trunk_forward` CudaFunction handle once at // construction (per `pearl_no_host_branches_in_captured_graph.md`). @@ -23111,6 +23582,27 @@ impl GpuDqnTrainer { aux_trunk_w3_v, aux_trunk_b3_m, aux_trunk_b3_v, + // SP14 Layer C Phase C.5a (2026-05-08): saved-fwd + accumulator + grad buffers + Adam infra (DEAD). + h_s2_aux, + h_aux1, + h_aux2, + dh_s2_aux_accum, + aux_trunk_w1_grad, + aux_trunk_b1_grad, + aux_trunk_w2_grad, + aux_trunk_b2_grad, + aux_trunk_w3_grad, + aux_trunk_b3_grad, + aux_trunk_grad_norm_partials, + aux_trunk_grad_block_offsets, + aux_trunk_grad_norm_buf, + aux_trunk_grad_clip_pinned, + aux_trunk_grad_clip_dev_ptr, + aux_trunk_lr_pinned, + aux_trunk_lr_dev_ptr, + aux_trunk_t_pinned, + aux_trunk_t_dev_ptr, + aux_trunk_adam_step: 0, // SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. aux_trunk_forward_ops, // SP14 Layer C Phase C.4 (2026-05-08): aux trunk backward orchestrator. diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 903e42177..7de206c31 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -992,6 +992,29 @@ pub struct GpuExperienceCollector { /// observability of the rollout-time label distribution). exp_aux_dir_acc_buf: super::mapped_pinned::MappedF32Buffer, + /// SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk + /// final output (`h_s2_aux`). Sized `[alloc_episodes × shared_h2]` + /// to mirror the trainer's `h_s2_aux [B_max, SH2]`. C.5b atomically + /// inserts `aux_trunk_forward.launch(... h_s2_aux_out=this ...)` + /// post-`forward_online_f32` so the collector-native EGF chain + /// reads from the aux representation pipeline instead of the + /// Q-trunk's `h_s2`. DEAD until C.5b. + #[allow(dead_code)] + exp_h_s2_aux: cudarc::driver::CudaSlice, + /// SP14 Layer C Phase C.5a: rollout-sized aux trunk Layer-1 post-ELU + /// activation. Sized `[alloc_episodes × AUX_TRUNK_H1=256]`. Saved + /// for parity with the trainer's per-step backward path; the + /// collector itself only runs forward (no backward), but the buffer + /// is allocated so the same `aux_trunk_forward.launch` signature + /// works in both contexts (collector vs trainer). DEAD until C.5b. + #[allow(dead_code)] + exp_h_aux1: cudarc::driver::CudaSlice, + /// SP14 Layer C Phase C.5a: rollout-sized aux trunk Layer-2 post-ELU + /// activation. Sized `[alloc_episodes × AUX_TRUNK_H2=128]`. Saved + /// for forward-API parity; see `exp_h_aux1` rationale. + #[allow(dead_code)] + exp_h_aux2: cudarc::driver::CudaSlice, + /// SP14 β-migration step 3 (2026-05-07): per-rollout-step label /// producer kernel handle. Sibling of the trajectory-wide /// `aux_sign_label_kernel` (line 3776 above) — same per-thread O(1) @@ -2022,6 +2045,34 @@ impl GpuExperienceCollector { "sp14-β: alloc exp_aux_dir_acc_buf (6 f32, collector): {e}" )))?; + // SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk + // saved-fwd buffers (h_s2_aux, h_aux1, h_aux2). Sized to + // `alloc_episodes` × per-layer dim (mirrors the β-migration + // pattern: replay-batch-sized in trainer, rollout-batch-sized + // here). DEAD until C.5b atomically inserts the + // `aux_trunk_forward.launch` post-`forward_online_f32` and + // redirects the EGF chain from Q's `h_s2` to `h_s2_aux`. No + // production caller reads these in C.5a. + use super::gpu_aux_trunk::{AUX_TRUNK_H1, AUX_TRUNK_H2}; + let exp_h_s2_aux = stream + .alloc_zeros::(alloc_episodes * shared_h2) + .map_err(|e| MLError::ModelError(format!( + "sp14-c.5a: alloc exp_h_s2_aux ({} f32, collector): {e}", + alloc_episodes * shared_h2 + )))?; + let exp_h_aux1 = stream + .alloc_zeros::(alloc_episodes * AUX_TRUNK_H1) + .map_err(|e| MLError::ModelError(format!( + "sp14-c.5a: alloc exp_h_aux1 ({} f32, collector): {e}", + alloc_episodes * AUX_TRUNK_H1 + )))?; + let exp_h_aux2 = stream + .alloc_zeros::(alloc_episodes * AUX_TRUNK_H2) + .map_err(|e| MLError::ModelError(format!( + "sp14-c.5a: alloc exp_h_aux2 ({} f32, collector): {e}", + alloc_episodes * AUX_TRUNK_H2 + )))?; + // SP14 β-migration step 3 (2026-05-07): per-rollout-step label // producer kernel. Loaded on the collector's stream so its launch // chains directly after `forward_online_f32` (which populates @@ -2472,6 +2523,11 @@ impl GpuExperienceCollector { exp_aux_nb_softmax_buf, exp_aux_nb_label_buf, exp_aux_dir_acc_buf, + // SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk + // saved-fwd buffers (DEAD until C.5b atomic wire-up). + exp_h_s2_aux, + exp_h_aux1, + exp_h_aux2, exp_aux_sign_label_per_step_kernel, // SP15 Wave 5 follow-up (2026-05-07): pre-loaded SP15 kernel handles. exp_sp15_dd_state_kernel, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ffabd27f4..9d1c3390c 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -7311,3 +7311,102 @@ Both slots use Pearl-A first-observation bootstrap. Slot 451 has α=0.05 EMA; sl **Open follow-ups:** - Add a target-variance EMA for `AVG_WIN_HOLD_TIME_BARS_INDEX` (Welford-style, single new ISV slot) in a future phase if the fixed α=0.01 blend is too slow/fast on real-data validation. The producer kernel already has the variance-slot path wired; pass `isv_h_target_var_idx >= 0` to enable Wiener-α. - Pearl note for `pearl_event_driven_reward_density_alignment.md` cross-reference: this phase aligns the AUX *label density* to event density (winning trade close), the same principle as the SP12 reward-density alignment. + +## SP14 Layer C Phase C.5a — additive infrastructure (DEAD CODE) (2026-05-08) + +> Phase C.5 was split (authorized 2026-05-08) after the prior implementer flagged ~600-800 LOC scope across 4 files with three correctness windows (gradient pointer aliasing, kernel rename + revert atomicity, Adam launcher dependency on grad buffer layout). C.5a is purely additive — buffers + launcher allocated as dead code. Build is clean before, build is clean after. C.5b atomically wires every consumer per `feedback_no_partial_refactor`. + +**No contract change in this commit.** Buffers exist, launcher compiles, no production caller invokes either. + +### Buffers allocated + +In `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`: + +**Saved-fwd tiles (replay-batch-sized):** +- `h_s2_aux: CudaSlice` — `[B, SH2]` final aux trunk output. Shape matches `aux_dh_s2_nb_buf` so the C.5b SAXPY-back step type-matches. +- `h_aux1: CudaSlice` — `[B, AUX_TRUNK_H1=256]` Layer-1 post-ELU activation. +- `h_aux2: CudaSlice` — `[B, AUX_TRUNK_H2=128]` Layer-2 post-ELU activation. + +**Upstream-grad accumulator:** +- `dh_s2_aux_accum: CudaSlice` — `[B, SH2]`. Mirrors `bw_d_h_s2`'s role for Q's path. C.5b reverts the zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` (commits `872bd7392` + `411a30473`) so the two aux head backwards SAXPY their per-sample dh into THIS buffer instead of the trunk's `bw_d_h_s2`. + +**Aux trunk gradient tensors (6, mirroring C.2 params):** +- `aux_trunk_w1_grad` `[SH1 × 256]`, `aux_trunk_b1_grad` `[256]`, +- `aux_trunk_w2_grad` `[256 × 128]`, `aux_trunk_b2_grad` `[128]`, +- `aux_trunk_w3_grad` `[128 × SH2]`, `aux_trunk_b3_grad` `[SH2]`. +- All `alloc_zeros`-initialised. Overwritten (not accumulated) per training step by `AuxTrunkBackwardOps::launch` in C.5b — matching the documented backward-API contract. + +**Adam launcher scratch + ISV-driven hparams:** +- `aux_trunk_grad_norm_partials: CudaSlice` `[total_blocks]` — per-block sum-of-squares slots laid out per `aux_trunk_grad_block_offsets[7]` (one offset per tensor + total). +- `aux_trunk_grad_norm_buf: CudaSlice` `[1]` — finalised L2 norm read by every Adam launch as `grad_norm_f32`. +- `aux_trunk_grad_clip_pinned: *mut f32` (mapped device-pinned) + `aux_trunk_grad_clip_dev_ptr` — host writes `ISV[AUX_TRUNK_GRAD_CLIP_INDEX=448]` per-step. +- `aux_trunk_lr_pinned: *mut f32` + `aux_trunk_lr_dev_ptr` — host writes `ISV[AUX_TRUNK_LR_INDEX=444]` per-step. +- `aux_trunk_t_pinned: *mut i32` + `aux_trunk_t_dev_ptr` — host increments before each launch. +- `aux_trunk_adam_step: i32` — host-side mirror. +- `Drop` impl frees the three new pinned host allocations. + +In `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`: +- `exp_h_s2_aux: CudaSlice` — `[alloc_episodes × shared_h2]`. +- `exp_h_aux1: CudaSlice` — `[alloc_episodes × AUX_TRUNK_H1=256]`. +- `exp_h_aux2: CudaSlice` — `[alloc_episodes × AUX_TRUNK_H2=128]`. + +Mirrors the trainer's saved-fwd buffers but rollout-sized — same β-migration pattern as `exp_aux_nb_hidden_buf`. + +### Adam launcher signature + +```rust +pub(crate) fn launch_aux_trunk_adam_update( + &self, + stream: &Arc, + _batch_size: i32, +) -> Result<(), MLError> +``` + +Writes 6 per-tensor `dqn_grad_norm_kernel` partials → 1 `dqn_grad_norm_finalize` (sqrt of sum across all 132K params) → 6 `dqn_adam_update_kernel` launches consuming the same global L2 norm. ISV-driven β1 (slot 445) / β2 (446) / ε (447) with numerical-stability floors per `pearl_first_observation_bootstrap`. SP4 Mech 9 weight clamp + Pearl C engagement counter disabled (aux trunk is outside the SP4 8-group taxonomy — same convention as OFI embed Adam). + +Mirrored launcher: `step_ofi_embed_adam` (single-flat-buffer Adam path). Adapted to 6 separate (params, grad, m, v) tensors with byte-offset-free layout. + +### Choice: 6-buffer (vs flat-buffer) + +Selected **6-buffer with 6 separate Adam launches** over flat-buffer with byte-offset views. Rationale: + +1. C.2 already allocated 6 separate param tensors (`aux_trunk_{w1,b1,w2,b2,w3,b3}`) and 12 separate Adam m/v tensors — a flat-buffer Adam would have required restructuring C.2's allocations to be views into a flat tensor, which is a contract change C.5a explicitly avoids. +2. The existing `dqn_adam_update_kernel` works on a single contiguous (params, grads, m, v) tuple per launch. 6 launches consuming 6 tuples is a smaller Δ than refactoring C.2 to a flat layout. +3. Global L2-norm clipping is preserved: the kernel reads `*grad_norm_f32` on every launch and computes its own per-launch `clip_scale = (norm > clip) ? clip/norm : 1`. All 6 launches see the same global norm + same clip threshold → consistent global clipping. +4. Storage cost is identical (same total f32 footprint). +5. Per-launch overhead (≤6 launches × ~1µs ≈ ≤6µs/step) is negligible relative to the trainer step. + +The 7th entry in `aux_trunk_grad_block_offsets` records the total block count, passed as `num_blocks` to the finalize kernel — capturing all 6 tensors' partial blocks in a single reduction. + +### NO contract change + +- `launch_aux_trunk_adam_update` is `#[allow(dead_code)]` because no production caller invokes it in C.5a. C.5b adds the call site atomically with the kernel rename + zero-fill revert. +- The 6 grad buffers, 4 saved-fwd tiles, accumulator, partials, norm buf, and pinned LR/clip/t are all `#[allow(dead_code)]` for the same reason. +- `aux_trunk_adam_step` initialises to 0 (counter is incremented by C.5b's call site). + +### Compiler warnings accepted (no `#[allow(dead_code)]` workaround) + +`cargo check -p ml --tests --all-targets` is clean. The new fields are `#[allow(dead_code)]`-tagged on the struct definitions because every struct field with no consumer otherwise warns; tagging them inline at the field is the canonical pattern this codebase uses for additive fields awaiting a wire-up commit (matches `aux_trunk_forward_ops`'s convention from C.3). + +### Files modified (this commit) + +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (+11 fields on `GpuDqnTrainer`; +allocations after C.2's site; +3 pinned host pointers freed in `Drop`; +`launch_aux_trunk_adam_update` method) +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+3 fields on `GpuExperienceCollector`; +allocations after `exp_aux_dir_acc_buf`'s site) +- Modify: `docs/dqn-wire-up-audit.md` (this section) + +### Verification + +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` — clean (only pre-existing warnings). +- `cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14, no regression). +- `grep -rn "launch_aux_trunk_adam_update" crates/ml/src/` — only the definition appears, no callers. + +### Follow-up: C.5b + +C.5b is the genuine atomic contract migration: +1. Insert `aux_trunk_forward.launch(...)` post-`forward_online_f32` in collector + post-`encoder_forward_chain` in trainer. +2. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`. +3. Revert zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790` so the SAXPY-into-`dh_s2_aux_accum` path works. +4. Param-rename `h_s2 → h_s2_aux` in the 4 aux_heads kernels. +5. Insert `aux_trunk_backward.launch(...)` + `launch_aux_trunk_adam_update(...)` calls. +6. Initialize `aux_trunk_lr_pinned` + `aux_trunk_grad_clip_pinned` from ISV per-step (host writes before each launch). +7. Increment `aux_trunk_adam_step` per Adam call. 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 5755fd391..51363231d 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 @@ -869,17 +869,113 @@ If C.10 L40S validation shows aux loss STILL flatlines at ln(2) with H=60: --- -## Phase C.5: Atomic Forward+Backward Wire-Up +## Phase C.5: Wire-Up — Split into C.5a (additive infra) + C.5b (atomic contract migration) -> **CRITICAL: this is ONE atomic commit per `feedback_no_partial_refactor`.** All consumers of aux head's input migrate from `h_s2` to `h_s2_aux` simultaneously with the new kernel launches in collector. If the build is broken between forward-only and backward-wired, the commit is incomplete. +> Split authorized 2026-05-08 after C.5 implementer reported ~600-800 LOC scope across 4 files with subtle gradient-correctness windows. Split is architecturally legitimate (mirrors SP14 Layer A→B pattern): +> - **C.5a** is *additive* — allocates buffers + adds Adam launcher as dead code. No contract change; nothing in production calls the new code yet. Cannot break the build's behavior. +> - **C.5b** is *atomic* — kernel rename + revert zero-fills + wire-up redirects in ONE commit. ~300 LOC. This is the genuine contract migration. -### Task C.5.1: Wire aux trunk forward + backward into collector chain +## Phase C.5a: Additive Infrastructure — Buffers + Adam Launcher (DEAD CODE) + +> No contract change. Build clean before, build clean after. Not yet wired. + +### Task C.5a.1: Allocate fwd/bwd buffers + Adam launcher **Files:** -- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` -- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (if backward chain lives here) -- Modify: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` (input parameter rename `h_s2` → `h_s2_aux`) -- Modify: `crates/ml/src/cuda_pipeline/gpu_aux_heads.rs` (Rust wrapper input plumbing) +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (allocate `h_s2_aux`, `h_aux1`, `h_aux2` saved-fwd; allocate aux trunk grad buffers if separate from params; add `launch_aux_trunk_adam_update` fn — dead, no caller yet) +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (allocate rollout-sized `h_s2_aux`, `h_aux1`, `h_aux2`) +- Modify: `docs/dqn-wire-up-audit.md` (audit entry — additive infra, no behavior change) + +- [ ] **Step 1: Allocate trainer-side saved-fwd buffers** + +In `gpu_dqn_trainer.rs`, alongside existing aux buffers (find `aux_next_bar_logits` or similar), allocate: +- `h_s2_aux: CudaSlice` — `[B_max, AUX_HIDDEN_DIM]` +- `h_aux1: CudaSlice` — `[B_max, H1=256]` saved-for-backward intermediate +- `h_aux2: CudaSlice` — `[B_max, H2=128]` saved-for-backward intermediate +- `dh_s2_aux_accum: CudaSlice` — `[B_max, AUX_HIDDEN_DIM]` accumulator for upstream gradient (where aux_next_bar_backward + aux_regime_backward will write SAXPY-style after C.5b's revert) + +These are dead until C.5b wires them up. + +- [ ] **Step 2: Allocate collector-side rollout-sized buffers** + +In `gpu_experience_collector.rs`, mirror trainer buffers but rollout-sized (`B * T_rollout` * dim). + +- [ ] **Step 3: Allocate aux trunk gradient buffers** + +C.2 allocated 6 separate weight/bias param tensors + 12 separate Adam m/v tensors. C.5a needs corresponding gradient buffers: +- `aux_trunk_w1_grad: CudaSlice` matching `aux_trunk_w1` shape +- ... same for b1, w2, b2, w3, b3 (6 grad tensors total) + +OR, if you prefer the flat-buffer Adam pattern that OFI embed uses: introduce flat `aux_trunk_params_flat`, `aux_trunk_grad_flat`, `aux_trunk_m_flat`, `aux_trunk_v_flat` (each ~132K f32) with byte-offset views into w1/b1/w2/b2/w3/b3 sub-regions. This requires re-declaring C.2's separate buffers as views into the flat buffer. + +**Implementer choice:** whichever fits cleanly into existing patterns. If 6-buffer approach is straightforward, prefer it (smaller refactor of C.2's work). If existing Adam launcher pattern strictly requires flat buffer, restructure to flat. Document the choice in audit entry. + +- [ ] **Step 4: Write `launch_aux_trunk_adam_update`** (dead, no caller yet) + +Mirror existing Adam launcher pattern. Search for `OFI` embed Adam path or similar single-param-group launcher (e.g., `launch_ofi_embed_adam`). The aux trunk lives outside the SP4-canonical 4-way Adam taxonomy (Q/Target/CQL/Aux per-branch), so a dedicated launcher is appropriate. + +Reads from ISV: +- `lr` from `ISV[AUX_TRUNK_LR_INDEX]` (444) +- `β1` from `ISV[AUX_TRUNK_BETA1_INDEX]` (445) +- `β2` from `ISV[AUX_TRUNK_BETA2_INDEX]` (446) +- `ε` from `ISV[AUX_TRUNK_EPS_INDEX]` (447) +- `grad_clip_norm` from `ISV[AUX_TRUNK_GRAD_CLIP_INDEX]` (448) + +Pre-step: global L2 norm of all 6 grad tensors → if exceeds clip threshold, scale all by `clip_norm / norm`. Then Adam step. + +- [ ] **Step 5: Verify** + +```bash +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets 2>&1 | tail -10 +``` + +Expected: clean build. Dead code may warn about unused fields/functions; that's OK because C.5b consumes them. Add `#[allow(dead_code)]` only if clippy promotes the warning to an error. + +Run all oracle tests to confirm nothing regressed: +```bash +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 -25 +``` + +Expected: 12 tests pass (no regression — the new buffers + launcher are inert). + +- [ ] **Step 6: Audit doc + commit** + +```bash +git add -A +git commit -m "$(cat <<'EOF' +feat(sp14-c.5a): allocate aux trunk fwd/bwd buffers + Adam launcher (dead code) + +Phase C.5a — additive infrastructure for the aux trunk wire-up. +Allocates saved-fwd buffers (h_s2_aux, h_aux1, h_aux2), +gradient buffers (6× aux_trunk_*_grad), accumulator +(dh_s2_aux_accum), and dedicated Adam launcher +(launch_aux_trunk_adam_update) reading β1/β2/ε/LR/grad-clip from +ISV[444..449). + +No contract change. No call sites for the new launcher yet. +C.5b atomically wires these in. + +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. + +### Task C.5b.1: Atomic kernel rename + revert zero-fills + wire-up + +> **CRITICAL: this is ONE atomic commit per `feedback_no_partial_refactor`.** All consumers of aux head's input migrate from `h_s2` to `h_s2_aux` simultaneously with the new kernel launches in collector. C.5a already allocated the buffers + Adam launcher (dead). C.5b makes them live atomically. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (insert aux_trunk_forward post-graph; redirect aux_heads_fwd input ptr to h_s2_aux) +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (insert aux_trunk_forward in trainer fwd; insert aux_trunk_backward + Adam launch in bwd; redirect aux_heads_fwd/bwd input/grad ptrs) +- Modify: `crates/ml/src/cuda_pipeline/aux_heads_kernel.cu` (param rename `h_s2` → `h_s2_aux` in 4 kernels; **REVERT** zero-fills in `aux_next_bar_backward:599-638` and `aux_regime_backward:757-790` from commits `872bd7392` + `411a30473` — restore proper SAXPY-back-to-input-grad) +- Modify: `crates/ml/src/cuda_pipeline/gpu_aux_heads.rs` (Rust wrapper input/grad plumbing) +- Modify: `docs/dqn-wire-up-audit.md` (atomic migration entry; note C.0 stop-grad audit comments are now stale and removed) - [ ] **Step 1: Locate aux head's current forward call site in collector**