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,
|
||||
|
||||
@@ -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::<f32>(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.
|
||||
|
||||
Reference in New Issue
Block a user