feat(sp22): H6 Phase 3b — α infrastructure loaded (NOT yet wired)
Phase 3b builds on Phase 3a (2e4c7ebf6). Adds the α infrastructure to the trainer: 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. Phase 3b is functionally equivalent to Phase 3a at runtime — the loaded kernels are dead code until the captured-graph integration (Phase 3c) lands. Architectural finding deferred to Phase 3c ────────────────────────────────────────── The trainer's 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 (V + A dueling combine: Q[a] = V + (A[a] - mean(A)) per atom), then immediately concatenates the result with h_s2 in one fused pass. There's no intermediate buffer between "Q_dir computed" and "Q_dir consumed" where α could inject as a parallel skip connection. Two options for α integration (Phase 3c will pick): 1. Modify mag_concat_qdir to take W_aux + state_121 args and apply the α bias to its internal Q_dir computation before concat. Invasive — changes a load-bearing kernel. 2. Add a NEW α-precompute kernel: write Q_dir into a dedicated buffer (V + A combine), then mag_concat_qdir reads from that buffer instead of doing the combine inline. Refactors the dueling head's forward — cleaner separation, bigger change. Phase 3b commits the α infrastructure so Phase 3c can focus solely on the captured-graph integration design choice without also needing to allocate buffers + load kernels. Files ───── - crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs: + pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN + pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN + 7 struct fields: 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 + new() block: alloc 4 zero-init f32 buffers (b0_size=4) for W + Adam moments + grad accumulator; load 2 cubins, 3 function handles. + Struct construction list extended with 7 new fields. - docs/dqn-wire-up-audit.md: Phase 3b entry documenting the architectural finding + Phase 3c scope. Verification ──────────── - cargo check -p ml --features cuda: 0 errors, 21 pre-existing warnings (Phase 2/3a baseline parity). - nvcc cubins unchanged (kernels built in Phase A). - Runtime equivalent to Phase 3a: α kernels never launched. Phase 3c scope (remaining for full α activation) ──────────────────────────────────────────────── - Pick option 1 or 2 for mag_concat_qdir integration. - Wire α forward in training captured forward graph (Step 7). - Wire α backward kernels in captured backward graph (Step 8). - Wire α Adam-step update (Step 9). - C1: α forward in collector's rollout-time captured graph. - D1-D7: A2 eval-side aux trunk + α + state-gather wiring. - B6: SP11 controller extension for non-zero scale_β. - B7/B10/B11: HEALTH_DIAG telemetry extensions. - E + F: verification gates + atomic Phase F commit + smoke + verdict. Estimated remaining: ~20-30 hr engineering + ~37 min smoke wall-clock. Refs ──── - docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md (spec) - docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md (runbook) -464bc5f7a(Phase A foundation) -2e4c7ebf6(Phase 3a — 7-component contract + β producer) - pearl_no_partial_refactor (Phase 3b is additive struct fields) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<f32>,
|
||||
/// 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<f32>,
|
||||
/// SP22 H6 Phase 3 α — Adam second moment (v) for `w_aux_to_q_dir`.
|
||||
adam_v_w_aux: cudarc::driver::CudaSlice<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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::<f32>(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::<f32>(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::<f32>(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::<f32>(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,
|
||||
|
||||
Reference in New Issue
Block a user