gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(sp22): H6 Phase 3 α cleanup + runbook revision for atom-shift design

Browse Source 
Same-session cleanup of Phase 3b scalar-bias residue (commit cb80b74ce's
additions that became architecturally redundant under the 2026-05-13
atom-shift revision). The scalar-bias α design is mathematically
ineffective in C51 distributional Q-learning (softmax-shift-invariance
→ W gradient = 0 → never trains). The revised design threads
`aux_atom_shift[b, a] = W[a] * state_121[b]` through compute_expected_q
+ c51_loss_kernel + c51_grad_kernel + mag_concat_qdir; dW + dstate
gradients integrate into c51_grad_kernel's projection backward.

Files
─────
- crates/ml/build.rs:
    Removed `aux_to_q_dir_bias_kernel.cu` + `aux_to_q_dir_bias_backward_kernel.cu`
    from kernels_with_common. Replaced with comment block documenting
    C51 softmax-invariance reason + spec/audit pointers.

- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:
    Removed `pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN` and
    `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` static declarations.
    Removed 3 struct fields (aux_to_q_dir_bias_kernel,
    aux_to_q_dir_bias_backward_dw_kernel,
    aux_to_q_dir_bias_backward_dstate_kernel) and their new()
    loading blocks + struct construction entries.
    KEPT: w_aux_to_q_dir + adam_m_w_aux + adam_v_w_aux + dw_aux_buf
    (still needed for atom-shift Adam-trained W).
    Updated W doc-comment to describe atom-shift threading (was
    scalar-bias) and structural-prior initialization `[-0.5, 0.0,
    +0.5, 0.0]`.

- crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu +
  crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_backward_kernel.cu:
    Source files stay on disk as committed dead code (commit
    464bc5f7a history preserved). nvcc no longer compiles them
    (build.rs unregistered). No `include_bytes!` references remain.

- docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md:
    Tasks A3/A4/A5 marked SUPERSEDED; A5 retains cleanup instructions.
    Task B9 Steps 4-10 rewritten as Steps 4 (cleanup), 5 (W structural
    prior init), 6 (compute_expected_q atom-shift threading),
    7 (c51_loss_kernel), 8 (c51_grad_kernel backward dW + dstate),
    9 (mag_concat_qdir), 10 (quantile_q_select/iqn_dual_head
    investigation), 11 (Adam wireup), 12 (verify + capture).
    Task C1 rewritten: collector passes 4 more args through
    existing compute_expected_q launcher (NO new kernel).

- docs/dqn-wire-up-audit.md:
    Cleanup-commit entry documenting the runbook revision and the
    code-cleanup actions.

Verification
────────────
- cargo check -p ml --features cuda: 0 errors, 21 pre-existing
  warnings (Phase 3b baseline parity).
- Build script no longer compiles the deleted-registration kernels.

Phase 3-final scope (~40-60 hr engineering — unchanged)
───────────────────────────────────────────────────────
Implementation per the revised runbook: atom-shift threading
through 4-5 kernels (compute_expected_q, c51_loss_kernel,
c51_grad_kernel, mag_concat_qdir, possibly quantile_q_select/
iqn_dual_head) + Adam wireup + collector launcher arg passing +
A2 eval-side + SP11 controller extension + HEALTH_DIAG telemetry
+ verification gates + atomic Phase F commit + smoke + verdict.

Refs
────
- docs/plans/2026-05-12-sp22-h6-phase3-alpha-beta.md (spec α section
  revised 2026-05-13 — commit 648078ce2)
- docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md (runbook
  α tasks revised in this commit)
- 464bc5f7a (Phase A — α kernels added; now committed dead code)
- cb80b74ce (Phase 3b — α struct fields + cubin statics added; cleaned
  up in this commit)
- pearl_no_partial_refactor (atom-shift threading is the new atomic
  contract for direction-branch atom positions)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-13 01:24:15 +02:00
parent 648078ce20
commit 08fd5803c4
4 changed files with 168 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
crates/ml/build.rs
View File

@@ -805,19 +805,19 @@ fn main() {
// stream right after `aux_next_bar_forward` in
// `gpu_experience_collector.rs`.
"aux_softmax_to_per_env_kernel.cu",
// SP22 H6 Phase 3 α (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]. Reads state_121 from encoder INPUT (batch_states),
// bypassing the cold encoder weight for slot 121. Backward
// kernel has TWO entry points: `aux_to_q_dir_bias_backward_dw`
// (per-action block tree-reduce; one block per action) and
// `aux_to_q_dir_bias_backward_dstate` (per-sample sum across
// 4 actions, accumulating into encoder-input gradient slot 121).
// No atomicAdd per `feedback_no_atomicadd`; no host branches
// per `pearl_no_host_branches_in_captured_graph` (capture-safe
// for both training and rollout forward graphs).
"aux_to_q_dir_bias_kernel.cu",
"aux_to_q_dir_bias_backward_kernel.cu",
// SP22 H6 Phase 3 α scalar-bias kernels — REGISTRATIONS DELETED
// 2026-05-13. The kernel `.cu` files
// (aux_to_q_dir_bias_kernel.cu, aux_to_q_dir_bias_backward_kernel.cu)
// stay on disk as committed dead code (commit 464bc5f7a) for git
// history. They are mathematically ineffective in C51 distributional
// Q-learning because adding a constant to Q is softmax-shift-
// invariant — W gradient = 0, W never trains. The atom-position
// shift replacement design threads `aux_atom_shift[b, a] = W[a] *
// state_121[b]` through compute_expected_q + c51_loss_kernel +
// c51_grad_kernel + mag_concat_qdir. The dW + dstate_121 gradient
// computations integrate into c51_grad_kernel's existing
// projection backward — no separate α kernel needed. See audit
// doc Phase 3c entry + spec doc α section for the revised design.
// SP14 Layer C Phase C.4b (2026-05-08): adaptive aux prediction
// horizon producer. Single-thread kernel writing
// `ISV[AUX_PRED_HORIZON_BARS_INDEX=450]` from

142
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -684,29 +684,23 @@ 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"));
// SP22 H6 Phase 3 α SCALAR-BIAS DESIGN — DELETED 2026-05-13.
// The scalar-bias `Q_dir[b, a] += W[a] * state_121[b]` approach is
// mathematically ineffective in C51 distributional Q-learning (softmax-
// shift-invariance → W gradient = 0 → never trains). Replaced by the
// per-(b, a) atom-position shift design that threads `aux_atom_shift[b, a]
// = W[a] * state_121[b]` through compute_expected_q + c51_loss_kernel +
// c51_grad_kernel + mag_concat_qdir + (investigate quantile_q_select /
// iqn_dual_head). The dW + dstate gradient computations integrate into
// c51_grad_kernel's projection backward — no separate α forward/backward
// kernel needed. See audit doc Phase 3c entry + spec doc α section
// for the revised design.
//
// The .cu files `aux_to_q_dir_bias_kernel.cu` and
// `aux_to_q_dir_bias_backward_kernel.cu` are removed from build.rs
// registration in the same commit. The Adam-trained W buffer
// (`w_aux_to_q_dir`) + Adam moments + dW accumulator remain — they
// still hold the per-action W vector used by the atom-shift threading.
/// 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
@@ -6892,15 +6886,20 @@ 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 α (2026-05-13): atom-position shift Adam-trained W ──
/// 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.
/// Adam-trained. Read by both the trainer (training-time atom-shift
/// threading through compute_expected_q + c51_loss_kernel +
/// c51_grad_kernel + mag_concat_qdir) and the collector (rollout-time
/// atom-shift in compute_expected_q for action selection, via shared
/// device pointer). Atom-shift formula:
/// `aux_atom_shift[b, a] = W[a] * state_121[b]`
/// applied per atom of direction-branch action a as
/// `z_n_effective[b, a, n] = z_n_base + aux_atom_shift[b, a]`.
/// **Initialization: structural prior** `[-0.5, 0.0, +0.5, 0.0]` for
/// `[Short, Hold, Long, Flat]` — domain belief that aux conviction ×
/// position-sign biases toward aligned trades. Adam refines from
/// the prior via projection backward in c51_grad_kernel.
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.
@@ -6908,23 +6907,11 @@ pub struct GpuDqnTrainer {
/// 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.
/// Written by `c51_grad_kernel`'s projection backward (dW[a] =
/// Σ_b state_121[b] * dz_shift[b, a] where dz_shift comes from
/// dloss/dz_n_effective); 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
@@ -20848,34 +20835,24 @@ impl GpuDqnTrainer {
.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)
};
// SP22 H6 Phase 3 α scalar-bias kernel loads — DELETED 2026-05-13.
// Per the architectural revision (audit doc "Phase 3c — α
// architectural finding"), the scalar-bias α design is replaced
// by atom-position shift threading through compute_expected_q +
// c51_loss_kernel + c51_grad_kernel + mag_concat_qdir. No
// separate α forward / backward kernel is launched — the
// gradient flows via c51_grad_kernel's existing projection
// backward. The kernel handles (aux_to_q_dir_bias_kernel,
// aux_to_q_dir_bias_backward_dw_kernel,
// aux_to_q_dir_bias_backward_dstate_kernel) and the CUBIN
// statics (SP22_AUX_TO_Q_DIR_BIAS_CUBIN,
// SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN) are deleted in this same
// commit. The W + Adam moments + dW buffer survive — they
// still hold the Adam-trained W vector used by atom-shift
// threading. Phase 3-final implementation writes the
// structural prior init `[-0.5, 0.0, +0.5, 0.0]` into
// `w_aux_to_q_dir` after `alloc_zeros` above (via a small
// fill-kernel launch or repurposed fill_f32 sequence).
// SP14 EGF kernel loads (B.3–B.6).
// Handles are held in struct fields but not launched until B.9/B.11;
@@ -25226,17 +25203,18 @@ 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.
// SP22 H6 Phase 3 α (2026-05-13 revised — atom-position shift).
// 4-element Adam-trained W + Adam moments + grad accumulator.
// No separate forward/backward kernel handles — the
// scalar-bias design was replaced by atom-shift threading
// through compute_expected_q + c51_loss_kernel + c51_grad_kernel
// + mag_concat_qdir (Phase 3-final implementation). W is
// initialized to the structural prior `[-0.5, 0.0, +0.5, 0.0]`
// (Phase 3-final step in `new()`); Adam refines from there.
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,

20
docs/dqn-wire-up-audit.md
View File

@@ -16978,4 +16978,22 @@ Shifts atom positions state-dependently. C51 Bellman projection re-projects targ
**β + SP11 controller + A2 eval-side + telemetry**: unchanged from previous spec sections — those parts of Phase 3 were architecturally sound. Only the α design needed correction.
Next session: execute the revised α design per the updated spec. The runbook (`docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md`) needs corresponding revision before execution — the runbook's α tasks (B9 Steps 4-9, C1) describe the original scalar-bias implementation and need to be rewritten for atom-shift threading.
Next session: execute the revised α design per the updated spec.
#### Cleanup commit (2026-05-13) — runbook revised, dead scalar-bias kernel registrations removed
Same-session cleanup of Phase 3b scalar-bias residue (commit `cb80b74ce`'s additions that became architecturally redundant under the atom-shift revision):
- `crates/ml/build.rs`: removed `"aux_to_q_dir_bias_kernel.cu"` and `"aux_to_q_dir_bias_backward_kernel.cu"` from `kernels_with_common` registration. Replaced with a comment block documenting the C51 softmax-invariance reason and pointing to the spec/audit references.
- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`: removed `pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN` and `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` static declarations. Removed 3 struct fields (`aux_to_q_dir_bias_kernel`, `aux_to_q_dir_bias_backward_dw_kernel`, `aux_to_q_dir_bias_backward_dstate_kernel`) and their `new()` loading blocks and struct construction entries. **Kept**: `w_aux_to_q_dir` + `adam_m_w_aux` + `adam_v_w_aux` + `dw_aux_buf` — the Adam-trained W buffer + moments + grad accumulator are still needed for the atom-shift design.
- `crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu` + `aux_to_q_dir_bias_backward_kernel.cu`: source files stay on disk as committed dead code (commit `464bc5f7a` history preserved); future cleanup may delete them. nvcc no longer compiles them (build.rs unregistered). No `include_bytes!` references remain.
Runbook revision (`docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md`):
- Tasks A3, A4, A5 marked SUPERSEDED; A5 retains cleanup instructions for removing the build.rs registrations.
- Task B9 Steps 4-10 rewritten as Steps 4 (cleanup), 5 (W structural prior init), 6-10 (atom-shift threading through compute_expected_q, c51_loss_kernel, c51_grad_kernel, mag_concat_qdir, quantile_q_select/iqn_dual_head investigation), 11 (Adam wireup), 12 (verify + capture-graph integrity).
- Task C1 rewritten for atom-shift threading through the collector's rollout-time `compute_expected_q` launcher (NO new kernel — just pass 4 more args through the existing launcher).
Verification: cargo check 0 errors, 21 pre-existing warnings (Phase 3b baseline parity).
Phase 3-final implementation is now atom-shift design end-to-end. The runbook is consistent with the spec.

104
docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md
View File

@@ -80,30 +80,26 @@ Total: **36 tasks**. (Phase C is folded inside Phase B Task B9 — alpha plumbin
- [ ] **Step 3: Verify compile.** If `training_loop.rs:9008` (the name-string match-expression on registry entries) needs handling for the new names, surface and fix.
### Task A3: Create alpha forward kernel
### Task A3 (SUPERSEDED 2026-05-13): create alpha forward kernel — ARCHITECTURALLY REDUNDANT under atom-shift design
**Files:**
- Create: `crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu`
**Status**: Phase A landed this kernel at commit `464bc5f7a` under the original scalar-bias spec. The 2026-05-13 architectural revision (see audit doc "Phase 3c — α architectural finding") replaced scalar-bias α with per-(b, a) atom-position shift. The atom-shift gradient computation integrates into `c51_grad_kernel`'s existing projection backward — no separate α forward kernel is launched. This kernel file (`aux_to_q_dir_bias_kernel.cu`) stays committed as dead code; the cleanup commit (Phase A→atom-shift transition) deletes it from `build.rs` registration and removes the loading code from `gpu_dqn_trainer.rs`. **No action in this task for Phase 3-final execution.**
- [ ] **Step 1: Write the kernel source.** Signature: `aux_to_q_dir_bias(float* q_dir [INOUT, B*b0_size], const float* w_aux [b0_size], const float* batch_states [B*state_dim], int aux_dir_prob_index, int state_dim, int batch_size, int b0_size)`. Kernel body: `int tid = blockIdx.x * blockDim.x + threadIdx.x; if (tid >= B*b0_size) return; int b = tid / b0_size; int a = tid - b * b0_size; float state_121 = batch_states[(long long)b * state_dim + aux_dir_prob_index]; q_dir[(long long)b * b0_size + a] += w_aux[a] * state_121;`. Document the architectural role (reads encoder INPUT, bypassing cold encoder), discipline (no atomicAdd, no host branches, capture-safe), and launch shape (grid = ceil(B*b0_size / 256), block = 256). Header preamble cites `feedback_no_atomicadd`, `pearl_no_host_branches_in_captured_graph`, `feedback_no_cpu_compute_strict`.
### Task A4 (SUPERSEDED 2026-05-13): create alpha backward kernel — ARCHITECTURALLY REDUNDANT under atom-shift design
- [ ] **Step 2: Verify file created** via `ls -la crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_kernel.cu`.
**Status**: same as A3. The backward dW + dstate gradients in the atom-shift design come from `c51_grad_kernel`'s projection-derivative arithmetic (per-(b, a) `dz_shift = Σ_n dloss/dz_n_effective` × state_121 for dW; × W for dstate_121). The standalone backward kernel (`aux_to_q_dir_bias_backward_kernel.cu`) stays committed as dead code; the cleanup commit deletes it. **No action in this task for Phase 3-final execution.**
### Task A4: Create alpha backward kernel
**Files:**
- Create: `crates/ml/src/cuda_pipeline/aux_to_q_dir_bias_backward_kernel.cu`
- [ ] **Step 1: Write the kernel source with TWO kernels in one .cu file.** Kernel `aux_to_q_dir_bias_backward_dw`: per-action block tree-reduce over batch; signature `(const float* dq_dir [B*b0_size], const float* batch_states [B*state_dim], float* dw_aux [b0_size, OUT], int aux_dir_prob_index, int state_dim, int batch_size, int b0_size)`. Launch one block per action, 256 threads, standard log2 tree reduce in shared mem; thread 0 writes `dw_aux[a]`. Kernel `aux_to_q_dir_bias_backward_dstate`: per-sample sum; signature `(const float* dq_dir, const float* w_aux, float* dbatch_states [B*state_dim, INOUT], int aux_dir_prob_index, int state_dim, int batch_size, int b0_size)`. Launch ceil(B/256), 256 threads; each thread computes `sum = Sum_a w_aux[a] * dq_dir[b*b0_size+a]` and ACCUMULATES (`+=`) into `dbatch_states[b*state_dim + aux_dir_prob_index]`. Both kernels: no atomicAdd, no host branches, capture-safe. Document option-(i) gradient routing (both alpha and encoder paths get gradient).
- [ ] **Step 2: Verify file created.**
### Task A5: Register both alpha kernels in build.rs
### Task A5 (SUPERSEDED 2026-05-13): register both alpha kernels in build.rs — REMOVE registrations
**Files:**
- Modify: `crates/ml/build.rs`
- [ ] **Step 1: Find the existing `aux_softmax_to_per_env_kernel.cu` entry** via `grep -nE "aux_softmax_to_per_env_kernel.cu" crates/ml/build.rs`.
**Cleanup action (Phase 3-final precondition)**:
- [ ] **Step 1: Find existing registrations** via `grep -nE "aux_to_q_dir_bias_kernel.cu|aux_to_q_dir_bias_backward_kernel.cu" crates/ml/build.rs`.
- [ ] **Step 2: Delete both lines** (and the SP22 H6 Phase 3 α comment block above them). The kernel `.cu` files stay on disk as committed dead code; future cleanup may delete the source files too.
- [ ] **Step 3: Verify compile** via `SQLX_OFFLINE=true cargo check -p ml --features cuda`. Expected: 0 errors, 21 pre-existing warnings (parity). The two cubin `include_bytes!` statics in `gpu_dqn_trainer.rs` (Phase 3b commit `cb80b74ce` Steps 4) reference the cubin paths — those statics also need removal in Phase B9 cleanup (see revised Task B9).
- [ ] **Step 2: Add the two new kernel filenames** right after, in the `kernels_with_common` array: `"aux_to_q_dir_bias_kernel.cu",` and `"aux_to_q_dir_bias_backward_kernel.cu",`. Each gets a brief comment citing SP22 H6 Phase 3 alpha and the kernel role.
@@ -224,19 +220,67 @@ This is the largest single change. It covers both the 7-component contract migra
- [ ] **Step 3: Update HEALTH_DIAG print for `reward_split`** to include `aux_align` (same as B8 Step 3 at the trainer's own print site).
- [ ] **Step 4: Add `pub(crate) static SP22_AUX_TO_Q_DIR_BIAS_CUBIN` and `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` statics** at the top of the file alongside the other SP22_* CUBIN statics, using `include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_q_dir_bias_kernel.cubin"))` pattern (and similar for the backward cubin).
**REVISED 2026-05-13 — atom-shift design (replaces scalar-bias Steps 4-9)**
- [ ] **Step 5: Add alpha-related struct fields**: `w_aux_to_q_dir: CudaSlice<f32>` [b0_size=4], `adam_m_w_aux: CudaSlice<f32>` [4], `adam_v_w_aux: CudaSlice<f32>` [4], `dw_aux_buf: CudaSlice<f32>` [4], `aux_to_q_dir_bias_kernel: CudaFunction`, `aux_to_q_dir_bias_backward_dw_kernel: CudaFunction`, `aux_to_q_dir_bias_backward_dstate_kernel: CudaFunction`. Document each with SP22 H6 Phase 3 alpha pointer.
The original Steps 4-9 prescribed a separate α forward + backward kernel that biases scalar Q_dir post-encoder. That design is **mathematically ineffective in C51 distributional Q-learning** because scalar Q bias is softmax-shift-invariant → W gradient = 0 (see audit doc Phase 3c entry). The revised α design shifts atom POSITIONS per-(b, a), threading `aux_atom_shift[b, a] = W[a] * state_121[b]` through every kernel that uses atom positions for the direction branch.
- [ ] **Step 6: Allocate buffers + load kernels in `new()`.** Allocate the 4 alpha-related `CudaSlice<f32>` buffers with `alloc_zeros::<f32>(b0_size)`. Load the alpha forward kernel from `SP22_AUX_TO_Q_DIR_BIAS_CUBIN`; load both backward kernels (`aux_to_q_dir_bias_backward_dw` + `aux_to_q_dir_bias_backward_dstate`) from `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN`. Add all 7 new fields to the struct construction at the end of `new()`.
**Cleanup of Phase 3b residue (precondition)** — Phase 3b commit `cb80b74ce` landed code that's now redundant under atom-shift; clean it up before threading atom-shift through C51 kernels:
- [ ] **Step 7: Wire alpha forward into the training-time forward captured graph.** Find the Q-head forward site (grep for `q_dir`). After Q_dir is computed but before Q_dir is consumed by downstream branches (mag concat / action selection), add a launch of `aux_to_q_dir_bias_kernel` with args (q_dir_buf, w_aux_to_q_dir, batch_states, AUX_DIR_PROB_INDEX, STATE_DIM, batch_size, b0_size). Launch shape: grid = ceil(B*b0_size / 256), block = 256.
- [ ] **Step 4 (cleanup): Remove dead CUBIN statics + kernel handle struct fields.**
- Delete `SP22_AUX_TO_Q_DIR_BIAS_CUBIN` + `SP22_AUX_TO_Q_DIR_BIAS_BWD_CUBIN` statics (gpu_dqn_trainer.rs).
- Delete struct fields `aux_to_q_dir_bias_kernel`, `aux_to_q_dir_bias_backward_dw_kernel`, `aux_to_q_dir_bias_backward_dstate_kernel`.
- Delete the `let *_kernel = { ... }` loading block in `new()` for those 3 kernel handles.
- Delete the corresponding entries from the struct construction list.
- **KEEP** `w_aux_to_q_dir`, `adam_m_w_aux`, `adam_v_w_aux`, `dw_aux_buf` (still needed for atom-shift Adam-trained W).
- Verify via `cargo check`.
- [ ] **Step 8: Wire alpha backward into the training-time backward captured graph.** Find the site where `dq_dir` is computed (after C51 / IQN heads emit gradients). Launch `aux_to_q_dir_bias_backward_dw_kernel` (grid = b0_size blocks, 256 threads each) and `aux_to_q_dir_bias_backward_dstate_kernel` (grid = ceil(B/256), 256 threads). Both consume `dq_dir`; the first writes `dw_aux_buf`, the second accumulates into `dbatch_states_buf[, AUX_DIR_PROB_INDEX]`.
- [ ] **Step 5: Update W initialization to structural prior.**
In `gpu_dqn_trainer.rs::new()`, replace `stream.alloc_zeros::<f32>(b0_size_usize)` for `w_aux_to_q_dir` with allocation + immediate write of the structural prior `[-0.5, 0.0, +0.5, 0.0]` via a `fill_f32`-style GPU init (NOT HtoD — use a small new init kernel, or repurpose the existing `fill_f32` with per-slot writes, or write a `w_aux_init_kernel.cu` that takes a 4-element constant array and fills the buffer). The Adam moments + dW buffer stay zero-init via `alloc_zeros`.
- [ ] **Step 9: Wire alpha Adam-step update.** Find the existing Adam-step site (where other parameter buffers get their per-step updates). Add an Adam call for `w_aux_to_q_dir` using the existing generic per-param Adam kernel. Launch shape: grid = (1, 1, 1), block = (4, 1, 1) since only 4 parameters.
Rationale: structural prior `[-0.5, 0, +0.5, 0]` reflects domain belief that aux conviction × position-sign biases toward aligned trades (see audit doc Phase 3c α revision). Adam refines W from this prior.
- [ ] **Step 10: Verify compile.** Expected: 0 errors. If `q_dir_buf` / `dq_dir_buf` / `dbatch_states_buf` / `adam_step_kernel` / `lr_current` symbol names differ in the trainer source, grep for the actual names and adapt.
- [ ] **Step 6: Thread `aux_atom_shift` into `compute_expected_q` (`experience_kernels.cu:4932`).**
Add 4 new kernel args at the end of the signature: `const float* w_aux`, `const float* batch_states`, `int aux_dir_prob_index`, `int state_dim`. Inside the per-action loop where `z_val` is computed for the direction branch (`d == 0`), add:
```c
float aux_atom_shift = 0.0f;
if (d == 0 && w_aux != NULL && batch_states != NULL) {
float state_121 = batch_states[(long long)i * state_dim + aux_dir_prob_index];
aux_atom_shift = w_aux[a] * state_121;
}
float z_val_eff = z_val + aux_atom_shift;
// Use z_val_eff in place of z_val for the expected-Q accumulation.
```
NULL-safe: when callers don't supply W (test scaffolds), the shift collapses to zero and `compute_expected_q` behaves bit-identically to pre-Phase-3.
Update the launcher in `gpu_dqn_trainer.rs` to pass the 4 new args (W ptr, batch_states ptr, AUX_DIR_PROB_INDEX, STATE_DIM). For test-side callers that don't have access, pass `NULL` ptrs.
- [ ] **Step 7: Thread `aux_atom_shift` into `c51_loss_kernel` (`c51_loss_kernel.cu`).**
Identify the kernel's Bellman projection code (`block_bellman_project_f` and callers — see line ~1122 of c51_loss_kernel.cu). When computing `t_z = clamp(reward + gamma * z_atom, v_min, v_max)` for the target distribution, the kernel needs to know the shifted atom positions for both:
- The PREDICTED distribution's atom positions (the predicted Q^(s, a) atoms, shifted by `aux_atom_shift[s_curr, a_curr]` for direction branch)
- The TARGET distribution's atom positions (Bellman target at next-state, shifted by `aux_atom_shift[s_next, a_next_argmax]` for direction branch)
Add same 4 args + per-call invocation logic. For non-direction branches, atom shift is 0 (existing behavior). Per `pearl_no_atomicadd`: gradient computation in the same kernel pass stays block-reduce.
- [ ] **Step 8: Thread `aux_atom_shift` into `c51_grad_kernel` (`c51_grad_kernel.cu`).**
The C51 gradient kernel computes `d_predicted_logit` from the projection. Under atom shift, the projection of target probability mass onto predicted atoms depends on the shifted positions. Compute `dz_shift[b, a] = Σ_n dloss/dz_n_effective` per (b, a) for direction branch. Then:
- `dW[a] += Σ_b state_121[b] * dz_shift[b, a]` (per-action block tree-reduce, no atomicAdd)
- `dstate_121[b] += Σ_a W[a] * dz_shift[b, a]` (per-sample sum, accumulate into `dbatch_states[b * state_dim + 121]`)
Both write to existing trainer buffers (`dw_aux_buf` and `dbatch_states_buf` slot 121).
- [ ] **Step 9: Thread `aux_atom_shift` into `mag_concat_qdir` (`experience_kernels.cu:5505`).**
This kernel computes internal Q_dir for magnitude-branch conditioning. To keep magnitude-branch conditioning consistent with action selection, apply the same atom shift inline. Add same 4 args; shift `z_val = v_min + (float)z * dz + aux_atom_shift` per atom inside the per-action loop (lines 5560-5594 area).
- [ ] **Step 10: Investigate `quantile_q_select` / `iqn_dual_head_kernel` for atom-position usage.**
Grep for `z_val` / atom-position arithmetic in these kernels. If they read atom positions for the direction branch in any consumer path (e.g., IQN's quantile-based action selection), thread the same atom shift. If they don't (e.g., IQN uses different parameterization), document and skip.
- [ ] **Step 11: Wire α Adam-step update.**
Same as original Step 9 in the prior runbook revision. Add an Adam call for `w_aux_to_q_dir` (4 params) at the existing Adam-step site. Generic per-param Adam kernel (β1=0.9, β2=0.999, ε=1e-8, lr from schedule). The `dw_aux_buf` populated by Step 8 feeds this Adam call.
- [ ] **Step 12: Verify compile + capture-graph integrity.**
`cargo check`, `gpu_backtest_validation`, `compute-sanitizer`. Then smoke + observe captured-graph integrity (no host branches added in any kernel; capture-safe). Confirm `W_aux_to_Q_dir` moves off initial prior values after epoch 1 (Adam updates accumulated → W gradient non-zero).
### Task B10: reward_component_monitor.rs — 7-component EMA reader
@@ -264,20 +308,24 @@ This is the largest single change. It covers both the 7-component contract migra
## Phase C — alpha plumbing (collector / rollout-side)
### Task C1: gpu_experience_collector.rs — wire alpha forward into rollout-time captured graph
### Task C1 (REVISED 2026-05-13): gpu_experience_collector.rs — atom-shift in rollout-time `compute_expected_q` launcher
**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
- [ ] **Step 1: Locate the rollout-time Q-head forward site** via `grep -nE "q_dir|on_dir_logits|forward_online" crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`.
Under the revised atom-shift design, α at rollout-time means: when the collector calls `compute_expected_q` (or whatever the rollout-time per-step expected-Q computer is), it must pass the W + batch_states + AUX_DIR_PROB_INDEX + STATE_DIM args so the atom shift fires during rollout action selection. No new kernel launches — just plumbing the new args through the existing call.
- [ ] **Step 2: Load the alpha forward kernel handle in `new()`** via `super::gpu_dqn_trainer::SP22_AUX_TO_Q_DIR_BIAS_CUBIN.to_vec()` → `load_function("aux_to_q_dir_bias")`. Store in `exp_aux_to_q_dir_bias_kernel: CudaFunction`. Add `w_aux_to_q_dir_dev_ptr: u64` field initialized to 0 (set via setter post-construction by the trainer).
The W tensor is owned by the trainer (Adam-trained). The collector reads it via a shared device pointer (same source-of-truth pattern as the trainer's GEMM weights). The trainer post-constructs the collector and calls a setter to wire the W ptr.
- [ ] **Step 1: Locate the rollout-time `compute_expected_q` launcher** via `grep -nE "compute_expected_q|expected_q_kernel" crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`.
- [ ] **Step 2: Add `w_aux_to_q_dir_dev_ptr: u64` struct field** initialized to 0 (set via setter post-construction by trainer). NO kernel handle field needed — `compute_expected_q` is an existing kernel; only the launcher args change.
- [ ] **Step 3: Add `set_w_aux_ptr(&mut self, ptr: u64)` setter method.**
- [ ] **Step 4: Wire alpha forward into the rollout-time forward graph.** Inside the rollout-step forward loop after Q_dir is computed and before action selection, add the same launch as B9 Step 7 with collector's local variables (n_episodes instead of B; w_aux_ptr from the setter).
- [ ] **Step 4: Update the rollout `compute_expected_q` launcher** to pass the 4 new args (W ptr, batch_states ptr from existing `self.batch_states.raw_ptr()`, `AUX_DIR_PROB_INDEX as i32`, `STATE_DIM as i32`). When `w_aux_to_q_dir_dev_ptr == 0` (setter not called yet, e.g., during test-only construction), pass NULL → atom shift collapses to zero → behavior identical to pre-Phase-3.
- [ ] **Step 5: Trigger the setter from the trainer**: in `gpu_dqn_trainer.rs`, after constructing the collector, call `collector.set_w_aux_ptr(self.w_aux_to_q_dir.raw_ptr())`.
- [ ] **Step 5: Trigger the setter from the trainer**: in `gpu_dqn_trainer.rs`, after constructing the collector, call `collector.set_w_aux_ptr(self.w_aux_to_q_dir.raw_ptr())`. Ensures rollout-time atom-shift sees the same Adam-trained W as the training-time path.
- [ ] **Step 6: Verify compile.**