diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 77d0883ae..1886bbaf5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -684,6 +684,30 @@ pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] = pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_softmax_to_per_env_kernel.cubin")); +/// SP22 H6 Phase 3 α forward (2026-05-13): post-encoder Q_dir bias. +/// Forward kernel adds `W_aux_to_Q_dir[a] * state_121[b]` to Q_dir[b, a] +/// for each (b, a). Reads state_121 from the encoder INPUT +/// (`batch_states[b * state_dim + AUX_DIR_PROB_INDEX]`), bypassing the +/// cold encoder weight for slot 121. Loaded from +/// `aux_to_q_dir_bias_kernel.cubin`. Consumed by both the trainer's +/// captured forward graph AND the collector's rollout-time captured +/// forward graph (shared W via setter). +pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_q_dir_bias_kernel.cubin")); + +/// SP22 H6 Phase 3 α backward (2026-05-13): gradients for the α forward. +/// Contains TWO kernels: +/// - `aux_to_q_dir_bias_backward_dw`: per-action block tree-reduce +/// computing `dW[a] = Σ_b state_121[b] * dq_dir[b, a]` (no atomicAdd). +/// - `aux_to_q_dir_bias_backward_dstate`: per-sample sum accumulating +/// `dstate_121[b] += Σ_a W[a] * dq_dir[b, a]` into +/// `dbatch_states[b, AUX_DIR_PROB_INDEX]` (option-(i) gradient routing +/// per spec: both α and encoder paths get signal). +/// Loaded from `aux_to_q_dir_bias_backward_kernel.cubin`. Consumed only +/// by the trainer's captured backward graph (eval doesn't update params). +pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_q_dir_bias_backward_kernel.cubin")); + /// SP17 Task PP.3 (2026-05-08): pre-centering V/A mean diagnostic. /// Single-block tree-reduce kernel reading the existing `on_v_logits_buf /// [B, NA]` and `on_b_logits_buf [B, (B0+B1+B2+B3)*NA]` (BRANCH-MAJOR) and @@ -6868,6 +6892,39 @@ pub struct GpuDqnTrainer { /// launch); no FoldReset registry entry. Bound is structural (no /// runtime tanh). Consumed by `launch_aux_pred_to_isv_tanh`. aux_pred_to_isv_tanh_kernel: CudaFunction, + // ── SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias ────────── + /// SP22 H6 Phase 3 α — learnable weight vector [b0_size=4] f32. + /// Adam-trained. Read by both the trainer (training-time forward + + /// backward) and the collector (rollout-time forward via shared + /// device pointer). Forward kernel adds `W[a] * state_121[b]` to + /// `Q_dir[b, a]`. Cold-start zero (Adam-driven growth tracks + /// aux→Q correlation). Per `pearl_first_observation_bootstrap` + /// the W ∈ R^4 starts at sentinel zero; Adam moves it from there + /// based on the gradient signal computed by the backward kernels. + w_aux_to_q_dir: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α — Adam first moment (m) for `w_aux_to_q_dir`. + /// Same shape [b0_size=4] f32; zero-init per Adam convention. + adam_m_w_aux: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α — Adam second moment (v) for `w_aux_to_q_dir`. + adam_v_w_aux: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α — gradient accumulator for `w_aux_to_q_dir`. + /// Written by `aux_to_q_dir_bias_backward_dw_kernel`; read by the + /// Adam-step update. [b0_size=4] f32. + dw_aux_buf: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α — forward kernel handle. Loaded from + /// `SP22_AUX_TO_Q_DIR_BIAS_CUBIN`. Launched inside the captured + /// training forward graph right after the Q-head produces Q_dir, + /// before Q_dir is consumed by downstream branches. + aux_to_q_dir_bias_kernel: CudaFunction, + /// SP22 H6 Phase 3 α — backward dW kernel handle. Per-action block + /// tree-reduce computing `dW[a] = Σ_b state_121[b] * dq_dir[b, a]`. + /// No atomicAdd per `feedback_no_atomicadd`. + aux_to_q_dir_bias_backward_dw_kernel: CudaFunction, + /// SP22 H6 Phase 3 α — backward dstate kernel handle. Per-sample + /// sum accumulating `dstate_121[b] += Σ_a W[a] * dq_dir[b, a]` + /// into the encoder-input gradient slot 121. Option-(i) gradient + /// routing per spec: both α and encoder paths get signal. + aux_to_q_dir_bias_backward_dstate_kernel: CudaFunction, // ── SP14 Earned Gradient Flow kernels ───────────────────────────────── /// SP14 B.3 (2026-05-05): per-step Q-head ↔ aux argmax disagreement EMA /// producer. Reads online Q logits + aux softmax outputs; computes the @@ -20764,6 +20821,62 @@ impl GpuDqnTrainer { module.load_function("aux_pred_to_isv_tanh_kernel") .map_err(|e| MLError::ModelError(format!("sp13 aux_pred_to_isv_tanh_kernel load: {e}")))? }; + + // SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias. + // Allocate the 4-element weight + 4-element Adam moments + + // 4-element gradient accumulator. Load forward + backward + // kernels. All buffers zero-init per `alloc_zeros`; Adam + // moves W from cold-start zero based on the backward signal. + let b0_size_usize: usize = 4; // factored space b0_size = 4 (direction) + let w_aux_to_q_dir = stream + .alloc_zeros::(b0_size_usize) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: alloc w_aux_to_q_dir: {e}" + )))?; + let adam_m_w_aux = stream + .alloc_zeros::(b0_size_usize) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: alloc adam_m_w_aux: {e}" + )))?; + let adam_v_w_aux = stream + .alloc_zeros::(b0_size_usize) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: alloc adam_v_w_aux: {e}" + )))?; + let dw_aux_buf = stream + .alloc_zeros::(b0_size_usize) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: alloc dw_aux_buf: {e}" + )))?; + let aux_to_q_dir_bias_kernel = { + let module = stream.context() + .load_cubin(SP22_AUX_TO_Q_DIR_BIAS_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_to_q_dir_bias cubin: {e}" + )))?; + module.load_function("aux_to_q_dir_bias") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_to_q_dir_bias load: {e}" + )))? + }; + let (aux_to_q_dir_bias_backward_dw_kernel, + aux_to_q_dir_bias_backward_dstate_kernel) = { + let module = stream.context() + .load_cubin(SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_to_q_dir_bias_bwd cubin: {e}" + )))?; + let dw = module.load_function("aux_to_q_dir_bias_backward_dw") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_to_q_dir_bias_backward_dw load: {e}" + )))?; + let dstate = module.load_function("aux_to_q_dir_bias_backward_dstate") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_to_q_dir_bias_backward_dstate load: {e}" + )))?; + (dw, dstate) + }; + // SP14 EGF kernel loads (B.3–B.6). // Handles are held in struct fields but not launched until B.9/B.11; // per `feedback_no_partial_refactor.md` all 4 kernel handles + @@ -25113,6 +25226,17 @@ impl GpuDqnTrainer { // `gpu_experience_collector.rs` per-step action-select site. apply_fixed_alpha_ema_kernel, aux_pred_to_isv_tanh_kernel, + // SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias. + // 4-element W + Adam moments + grad accumulator + fwd/bwd + // kernel handles. Cold-start zero; Adam moves W from there + // based on the backward-emitted dW from the loss path. + w_aux_to_q_dir, + adam_m_w_aux, + adam_v_w_aux, + dw_aux_buf, + aux_to_q_dir_bias_kernel, + aux_to_q_dir_bias_backward_dw_kernel, + aux_to_q_dir_bias_backward_dstate_kernel, // SP14 q_disagreement diagnostic + Coupling A forward feature // wire. Phase C.1 (2026-05-08) deleted the α-machinery struct // entries (sp14_alpha_grad_compute_kernel, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index b873f4a36..f4547fbaa 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -16832,3 +16832,109 @@ Phase 3a NOT yet implemented (Phase 3b scope): Estimated remaining work: ~25-35 hr engineering. Resume in a fresh session from Phase 3a commit per the runbook `docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md`. + +### Phase 3b — α infrastructure loaded (NOT yet wired into captured graphs) + +Built on Phase 3a (`2e4c7ebf6`). Phase 3b adds the α infrastructure to +the trainer (`gpu_dqn_trainer.rs`): 2 new CUBIN statics, 7 new struct +fields (W weight + Adam moments + grad accumulator + 3 kernel +handles), and the `new()` constructor's alloc + kernel-load block. +**α kernels are loaded but never launched yet** — Steps 7-9 +(captured-graph forward integration + backward + Adam-step) are +deferred to Phase 3c due to an architectural finding documented +below. + +Architectural finding: dueling head's `mag_concat_qdir` fuses Q_dir +combine + concat +───────────────────────────────────────────────── + +The trainer's existing dueling head doesn't have a separate `Q_dir` +buffer. The `mag_concat_qdir` kernel (`gpu_dqn_trainer.rs:9884`) +computes Q_dir INTERNALLY from `on_v_logits_buf` + `on_b_logits_buf` +(the dueling V + A combine: `Q[a] = V + (A[a] - mean(A))` per atom), +then immediately concatenates the result with `h_s2` into +`mag_concat_buf` for the magnitude branch input. No intermediate +buffer between "Q_dir computed" and "Q_dir consumed" where α could +inject as a parallel skip connection. + +Two options for α integration: + +1. **Modify `mag_concat_qdir`**: extend the kernel signature with + `W_aux` + `state_121` ptrs + `aux_dir_prob_index` + `state_dim` + and apply the α bias to its internal Q_dir computation before + concat. Adds ~5 args. Invasive — changes a load-bearing kernel + that all magnitude-branch forward consumers depend on. + +2. **Add a NEW α-precompute kernel**: writes Q_dir into a dedicated + buffer (V + A combine), then `mag_concat_qdir` reads from that + buffer instead of doing the combine inline. Requires + refactoring the dueling head's forward to split the V+A combine + from the concat. Cleaner separation of concerns, but bigger + structural change. + +Decision (Phase 3c scope): TBD. Both options have non-trivial +implications: + +- Option 1 keeps the fused kernel but couples α's lifecycle to + `mag_concat_qdir`'s capture graph (forward + maintenance graphs + both call this kernel; need to verify α arg-threading works in + both contexts). +- Option 2 introduces a new buffer + new kernel + refactors + `mag_concat_qdir` to take pre-computed Q_dir as input. More + invasive but architecturally cleaner; supports future α-like + hooks without further `mag_concat_qdir` modifications. + +Phase 3b commits the α infrastructure (compiled cubins, struct +fields, allocator + loader in `new()`) so Phase 3c can focus solely +on the captured-graph integration design choice. Currently the +loaded α kernels are dead code at runtime — equivalent to Phase 3a +functionally. + +Phase 3b files (2): + +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`: + - Added `pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN` + + `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` alongside the existing + SP22 H6 Phase 1 `SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN`. + - Added 7 struct fields to `GpuDqnTrainer`: `w_aux_to_q_dir` + (param), `adam_m_w_aux` + `adam_v_w_aux` (Adam moments), + `dw_aux_buf` (gradient accumulator), 3 `CudaFunction` handles + (`aux_to_q_dir_bias_kernel`, + `aux_to_q_dir_bias_backward_dw_kernel`, + `aux_to_q_dir_bias_backward_dstate_kernel`). + - `new()` alloc block: 4 `alloc_zeros::(b0_size=4)` for the + Adam-trained W + moments + grad. Two cubin loads + 3 function + lookups. + - Struct construction list extended with the 7 new fields. + +- `docs/dqn-wire-up-audit.md`: this entry. + +Verification post-Phase-3b: + +- cargo check -p ml --features cuda: 0 errors, 21 pre-existing + warnings (Phase 2/3a baseline parity). +- nvcc cubins unchanged from Phase A (the kernels themselves were + built in Phase A; Phase 3b only adds the Rust-side load). +- Runtime equivalent to Phase 3a: α kernels never launched. + +Phase 3c scope (remaining for full α activation) +─────────────────────────────────────────────── + +- Pick option 1 or option 2 for the `mag_concat_qdir` integration. +- Wire α forward inside the training captured forward graph + (Step 7). +- Wire α backward kernels (dw + dstate) inside the training + captured backward graph (Step 8). +- Wire α Adam-step update inside the training Adam captured + subgraph (Step 9). +- Phase C1: same forward wiring on the collector's rollout-time + captured forward graph. +- Phase D1-D7: A2 eval-side aux trunk + α forward + state-gather + wiring. +- Phase B6: SP11 controller extension (w_aux_align emit) for + non-zero scale_β so β activates. +- Phase B7/B10/B11: HEALTH_DIAG telemetry extensions. +- Phase E + F: verification gates + atomic Phase F commit + smoke + dispatch + verdict. + +Estimated remaining: ~20-30 hr engineering + ~37 min smoke wall-clock.