From 195c051e7ad43e9d65fb558cd91af615e02f9a2f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 13 May 2026 01:41:26 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22):=20H6=20Phase=203=20=CE=B1=20?= =?UTF-8?q?=E2=80=94=20atom-shift=20wired=20through=20compute=5Fexpected?= =?UTF-8?q?=5Fq=20(B9=20Step=205+6)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit W structural prior init kernel (Step 5): - aux_w_prior_init_kernel.cu: 4-thread one-time init writing W[a] = [-0.5, 0.0, +0.5, 0.0] (Short / Hold / Long / Flat). Hardcoded values in kernel — no HtoD per feedback_no_htod_htoh_only_mapped_pinned.md. - build.rs registration + gpu_dqn_trainer.rs cubin static + handle field + new() launch after alloc_zeros for w_aux_to_q_dir. compute_expected_q atom-shift (Step 6): - experience_kernels.cu signature grows 4 args (w_aux, batch_states, aux_dir_prob_index, state_dim). NULL-safe — collapses to 0 shift when either pointer is NULL, bit-identical to pre-Phase-3-α. - Per-action inner loop computes aux_atom_shift = w_aux[a] * state_121 once per (b, a) for d==0 only. Inner z-loop applies z_val += aux_atom_shift before all S/TZ/TZ²/TLM accumulators consume it. Launcher updates (3 trainer sites + 1 collector site): - populate_q_out: passes W + current states (online path). - replay_forward_for_q_values: passes W + current states (online replay). - compute_denoise_target_q: passes W + next_states (target on s'; W shared across online/target). - gpu_experience_collector.rs: passes NULL W + NULL states until Phase C1 wires the trainer's W ptr through a setter. Architectural notes: - Action selection (every compute_expected_q call) now uses shifted atom positions for direction branch (d==0). Other branches stay bit-identical (shift = 0). Step 7 (c51_loss_kernel) + Step 8 (c51_grad_kernel) will close the loop on training loss + W gradient in the same atomic commit to avoid the gradient mismatch trap per feedback_no_partial_refactor.md. - Adam wireup for w_aux_to_q_dir lands at Step 11; until then W stays at structural prior values (no Adam step modifies it). Verification: - cargo check -p ml --lib: 0 errors, 21 pre-existing warnings (Phase 3b baseline parity). - Audit doc updated with B9 Step 5+6 checkpoint entry. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 9 +++ .../cuda_pipeline/aux_w_prior_init_kernel.cu | 62 ++++++++++++++++ .../src/cuda_pipeline/experience_kernels.cu | 37 +++++++++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 73 +++++++++++++++++++ .../cuda_pipeline/gpu_experience_collector.rs | 14 ++++ docs/dqn-wire-up-audit.md | 49 +++++++++++++ 6 files changed, 243 insertions(+), 1 deletion(-) create mode 100644 crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index c3be05695..105421c55 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -818,6 +818,15 @@ fn main() { // 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. + // + // SP22 H6 Phase 3 α (atom-shift design 2026-05-13): structural-prior + // init for w_aux_to_q_dir [4] — writes per-action prior + // [-0.5, 0.0, +0.5, 0.0] for [Short, Hold, Long, Flat] at trainer + // construction. One-time GPU init (NOT in captured graph). The + // Adam-trained W refines from this prior via c51_grad_kernel's + // projection backward (Phase 3-final B9 Step 8). Pure 4-thread + // single-block; trivial cost. + "aux_w_prior_init_kernel.cu", // 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 diff --git a/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu b/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu new file mode 100644 index 000000000..c8936358a --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu @@ -0,0 +1,62 @@ +// crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu +// +// SP22 H6 Phase 3 α (2026-05-13): structural-prior initialization for +// the Adam-trained `w_aux_to_q_dir [b0_size=4]` weight vector. +// +// Writes the per-action prior values: +// W[0] = Short → -0.5 (aux conviction × position-sign favors Long; +// negate to discourage Short when aux is positive) +// W[1] = Hold → 0.0 (no atom shift for Hold action) +// W[2] = Long → +0.5 (encourage Long when aux predicts up) +// W[3] = Flat → 0.0 (no atom shift for Flat action) +// +// Rationale (audit doc Phase 3c α revision, spec doc α section): +// state_121 ∈ [-1, +1] is the recentered aux p_up. When aux predicts +// up (state_121 = +1): +// - W[Short] * state_121 = -0.5 → Q_short atoms shift DOWN by 0.5 +// → action selection avoids Short. +// - W[Long] * state_121 = +0.5 → Q_long atoms shift UP by 0.5 +// → action selection biases toward Long. +// - Hold / Flat: unchanged. +// +// When aux predicts down (state_121 = -1): symmetrically biases +// toward Short and away from Long. +// +// Adam refines W from this prior via projection backward in +// c51_grad_kernel.cu (B9 Step 8 — the dz_shift gradient flows to +// dW via per-action block tree-reduce). +// +// Discipline +// ────────── +// - One-time init call at trainer construction; NOT in the captured +// training loop. +// - No HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md` (pure +// GPU compute; hardcoded values in the kernel). +// - 4 threads, single block — trivial cost (<1 µs). +// - Capture-safety: the kernel is not invoked from inside any +// captured graph (per-step or per-fold), so capture rules don't +// apply; conventional kernel discipline (no atomicAdd, no host +// branches inside the kernel body) preserved anyway. +// +// Launch: grid=(1, 1, 1), block=(4, 1, 1). Each of the 4 threads +// writes one slot. + +#include + +extern "C" __global__ void aux_w_prior_init( + /* [b0_size=4] f32 output buffer; overwrites alloc_zeros'd content + * with the structural prior values. */ + float* __restrict__ w_aux_to_q_dir +) { + int a = threadIdx.x; + if (a >= 4) return; + + /* Per-action structural prior: + * a == 0 (Short) → -0.5 + * a == 1 (Hold) → 0.0 + * a == 2 (Long) → +0.5 + * a == 3 (Flat) → 0.0 + */ + float prior_values[4] = { -0.5f, 0.0f, +0.5f, 0.0f }; + w_aux_to_q_dir[a] = prior_values[a]; +} diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index ca1c448e8..3895a7618 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -4943,7 +4943,20 @@ extern "C" __global__ void compute_expected_q( float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2]: per-block partial sums (phase 1). * NULL to skip atom-stat collection entirely. */ float* __restrict__ q_variance, /* [N, total_actions] Var[Q] per action. NULL to skip. */ - const float* __restrict__ atom_positions /* [4, num_atoms] adaptive positions. NULL = use linear. */ + const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = use linear. */ + /* SP22 H6 Phase 3 α atom-shift (2026-05-13): direction-branch (d==0) + * per-(b, a) atom position shift. Computed inline per action a as + * aux_atom_shift = w_aux[a] * batch_states[b * state_dim + aux_dir_prob_index] + * and added to z_val inside the per-atom expected-Q accumulation. + * NULL-safe: when callers don't supply W (test scaffolds or non-α + * launches), shift collapses to 0 and the kernel behaves + * bit-identically to pre-Phase-3-α. + * Shift applies ONLY to direction branch (d == 0); other branches + * use the unshifted per_sample_support atom positions. */ + const float* __restrict__ w_aux, /* [b0_size=4] f32, or NULL */ + const float* __restrict__ batch_states, /* [N, state_dim] f32, or NULL */ + int aux_dir_prob_index, /* = SL_PADDING_START + 0 = 121 */ + int state_dim /* = STATE_DIM = 128 */ ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; @@ -5020,6 +5033,20 @@ extern "C" __global__ void compute_expected_q( for (int a = 0; a < n_d; a++) { const float* adv_a = branch_base + (long long)a * num_atoms; + /* SP22 H6 Phase 3 α atom-shift (2026-05-13): direction-branch + * (d == 0) per-(b, a) atom shift hoisted outside the per-z + * loops below. Computed once per (b, a), applied uniformly + * to every atom of action a's distribution. NULL-safe: 0 + * when W or batch_states absent → bit-identical to + * pre-Phase-3-α for d == 0 too. For d != 0 the shift stays + * at zero (only direction branch is α-influenced). */ + 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; + } + /* ── Online softmax: 2-pass instead of 3-pass ──────────────── * Pass 1: streaming (running_max, running_sum_exp_shifted) plus * running accumulators for E[Z], E[Z²], E[Z*(logit - max)] — @@ -5056,6 +5083,14 @@ extern "C" __global__ void compute_expected_q( float z_val = (atom_pos_d != NULL) ? atom_pos_d[z] : v_min + (float)z * dz; + /* SP22 H6 Phase 3 α atom-shift: shift atom position by + * per-(b, a) state-dependent bias (direction branch only). + * Mathematically identical to E[Q] += aux_atom_shift for + * action selection, but here it flows through the C51 + * distributional support so c51_loss_kernel + c51_grad_kernel + * see the shifted positions when computing the Bellman + * projection (gradient flows via dproj/dz). */ + z_val += aux_atom_shift; if (logit > M) { /* Rescale running accumulators down to the new max. */ float scale = expf(M - logit); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 954c5038b..d811bc832 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -684,6 +684,15 @@ 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 α structural-prior init (atom-shift design 2026-05-13). +/// One-time GPU init kernel that writes the per-action prior +/// `[-0.5, 0.0, +0.5, 0.0]` for `[Short, Hold, Long, Flat]` into the +/// Adam-trained `w_aux_to_q_dir [4]` buffer at trainer construction. +/// NOT invoked inside any captured graph — pure init-time launch. +/// Loaded from `aux_w_prior_init_kernel.cubin`. +pub(crate) static SP22_AUX_W_PRIOR_INIT_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_w_prior_init_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- @@ -20835,6 +20844,40 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!( "sp22-h6-phase3 α: alloc dw_aux_buf: {e}" )))?; + + // SP22 H6 Phase 3 α (atom-shift design 2026-05-13): structural-prior + // init for w_aux_to_q_dir. Writes per-action prior + // `[-0.5, 0.0, +0.5, 0.0]` for `[Short, Hold, Long, Flat]`. The + // Adam-trained W refines from this prior via c51_grad_kernel's + // projection backward (Phase 3-final B9 Step 8). One-time + // GPU init — NOT inside any captured graph. + { + let init_module = stream.context() + .load_cubin(SP22_AUX_W_PRIOR_INIT_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_w_prior_init cubin: {e}" + )))?; + let init_kernel = init_module + .load_function("aux_w_prior_init") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_w_prior_init load: {e}" + )))?; + let w_ptr = w_aux_to_q_dir.raw_ptr(); + unsafe { + stream + .launch_builder(&init_kernel) + .arg(&w_ptr) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (4, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α: aux_w_prior_init launch: {e}" + )))?; + } + } + // 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 @@ -28180,6 +28223,12 @@ impl GpuDqnTrainer { let atom_stats_block_sums_ptr = self.atom_stats_block_sums_buf.raw_ptr(); let q_var_ptr = self.q_var_buf_trainer.raw_ptr(); let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); + // SP22 H6 Phase 3 α atom-shift: pass W + current states + AUX_DIR_PROB_INDEX + // so the kernel shifts direction-branch atom positions by W[a] * state_121. + // Online forward → use current states_buf. + let w_aux_ptr = self.w_aux_to_q_dir.raw_ptr(); + let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; + let state_dim_i32 = ml_core::state_layout::STATE_DIM as i32; // Phase 1: per-block sums (deterministic, no atomic) → block_sums unsafe { self.stream @@ -28197,6 +28246,10 @@ impl GpuDqnTrainer { .arg(&atom_stats_block_sums_ptr) .arg(&q_var_ptr) .arg(&atom_positions_buf_ptr) + .arg(&w_aux_ptr) + .arg(&self.ptrs.states_buf) + .arg(&aux_dir_prob_index) + .arg(&state_dim_i32) .launch(LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), @@ -28326,6 +28379,10 @@ impl GpuDqnTrainer { let null_atom_stats = 0u64; let null_q_var = 0u64; let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); + // SP22 H6 Phase 3 α atom-shift: online replay → current states. + let w_aux_ptr = self.w_aux_to_q_dir.raw_ptr(); + let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; + let state_dim_i32 = ml_core::state_layout::STATE_DIM as i32; unsafe { self.stream .launch_builder(&self.expected_q_kernel) @@ -28342,6 +28399,10 @@ impl GpuDqnTrainer { .arg(&null_atom_stats) .arg(&null_q_var) .arg(&atom_positions_buf_ptr) + .arg(&w_aux_ptr) + .arg(&self.ptrs.states_buf) + .arg(&aux_dir_prob_index) + .arg(&state_dim_i32) .launch(LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), @@ -28768,6 +28829,14 @@ impl GpuDqnTrainer { let null_atom_stats = 0u64; let null_q_var = 0u64; let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); + // SP22 H6 Phase 3 α atom-shift: target Q evaluated on s' → next_states_buf. + // W is shared between online and target (no separate target W: aux head + // is online-only — the atom-shift on target Q uses the SAME W and reads + // next-state aux p_up that was written into next_states_buf[121] by the + // state-assembly pipeline). + let w_aux_ptr = self.w_aux_to_q_dir.raw_ptr(); + let aux_dir_prob_index = ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; + let state_dim_i32 = ml_core::state_layout::STATE_DIM as i32; unsafe { self.stream .launch_builder(&self.expected_q_kernel) @@ -28784,6 +28853,10 @@ impl GpuDqnTrainer { .arg(&null_atom_stats) .arg(&null_q_var) .arg(&atom_positions_buf_ptr) + .arg(&w_aux_ptr) + .arg(&self.ptrs.next_states_buf) + .arg(&aux_dir_prob_index) + .arg(&state_dim_i32) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index e0a726b04..3d165b063 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -5646,6 +5646,16 @@ impl GpuExperienceCollector { unsafe { let null_atom_stats = 0u64; let null_atom_positions = 0u64; // NULL = use linear atom spacing + // SP22 H6 Phase 3 α atom-shift collector-side: Phase C1 will + // wire trainer's W ptr in via a setter (collector doesn't own + // the Adam-trained weight). For now NULL → kernel collapses + // shift to 0, bit-identical to pre-Phase-3-α rollout behavior. + let null_w_aux = 0u64; + let null_batch_states = 0u64; + let aux_dir_prob_index = + ml_core::state_layout::AUX_DIR_PROB_INDEX as i32; + let state_dim_i32 = + ml_core::state_layout::STATE_DIM as i32; self.stream .launch_builder(&self.expected_q_kernel) .arg(&self.exp_v_logits) @@ -5661,6 +5671,10 @@ impl GpuExperienceCollector { .arg(&null_atom_stats) .arg(&mut self.q_var_buf) .arg(&null_atom_positions) // atom_positions (NULL = linear) + .arg(&null_w_aux) + .arg(&null_batch_states) + .arg(&aux_dir_prob_index) + .arg(&state_dim_i32) .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!( "compute_expected_q t={t}: {e}" diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index acbbee867..ec06f7b8f 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -16997,3 +16997,52 @@ Runbook revision (`docs/plans/2026-05-13-sp22-h6-phase3-alpha-beta-runbook.md`): 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. + +#### Atom-shift wiring checkpoint — Step 5 (W prior init) + Step 6 (compute_expected_q threading) — 2026-05-13 + +**Step 5 — W structural prior init kernel** (committed alongside this checkpoint): + +- New file: `crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu`. 4-thread kernel `aux_w_prior_init(float* w_aux_to_q_dir)` writes per-action structural prior `W[0..4] = [-0.5, 0.0, +0.5, 0.0]` (Short / Hold / Long / Flat). +- `build.rs`: registered new kernel. +- `gpu_dqn_trainer.rs`: added `SP22_AUX_W_PRIOR_INIT_CUBIN` static + kernel handle field `aux_w_prior_init_kernel`; loaded in `new()` and launched once after `alloc_zeros(b0_size=4)` for `w_aux_to_q_dir`. One-time init, NOT inside any captured graph; cost <1µs. +- Per `feedback_no_htod_htoh_only_mapped_pinned.md`: pure GPU compute, prior values hardcoded in kernel — no HtoD memcpy. + +**Step 6 — `compute_expected_q` atom-shift threading**: + +`crates/ml/src/cuda_pipeline/experience_kernels.cu::compute_expected_q` signature grew 4 args (after `atom_positions`): +- `const float* __restrict__ w_aux` (`[b0_size=4]` f32, NULL-safe) +- `const float* __restrict__ batch_states` (`[N, state_dim]` f32, NULL-safe) +- `int aux_dir_prob_index` (= `SL_PADDING_START` = 121) +- `int state_dim` (= `STATE_DIM` = 128) + +Per-action inner loop (`for a = 0 .. n_d`) computes once-per-action: +``` +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; +} +``` +Per-z body adds `z_val += aux_atom_shift;` after `z_val = atom_pos_d[z]` (or linear). All running accumulators (S/TZ/TZ²/TLM) consume the shifted z. Shift applies ONLY to direction branch (d==0); other branches stay bit-identical. + +NULL-safety: when either pointer is NULL the shift collapses to 0 and the kernel produces bit-identical output to pre-Phase-3-α — required so test scaffolds (and the collector during the C1 interim window) can call the same kernel without W wired. + +**Launcher updates (3 trainer call sites + 1 collector site)**: + +- `gpu_dqn_trainer.rs::populate_q_out` — online forward, passes `w_aux_to_q_dir.raw_ptr()` + `ptrs.states_buf` + `AUX_DIR_PROB_INDEX` + `STATE_DIM`. +- `gpu_dqn_trainer.rs::replay_forward_for_q_values` — same args, online replay path. +- `gpu_dqn_trainer.rs::compute_denoise_target_q` — target Q on s' → passes `ptrs.next_states_buf` (W is shared across online/target; aux p_up for s' is written into `next_states_buf[121]` by the state-assembly pipeline). +- `gpu_experience_collector.rs` (rollout-time launcher) — passes NULL for W + states; Phase C1 will add a setter so the trainer's W ptr flows here. State_dim/aux_dir_prob_index args still passed (the kernel reads them but ignores when W is NULL). + +`cargo check -p ml --lib`: 0 errors, 21 pre-existing warnings (Phase 3b baseline parity). + +**What this commit does NOT yet do** (remaining B9 work): +- Step 7: `c51_loss_kernel` projection-side atom-shift threading (large — Bellman projection arithmetic). +- Step 8: `c51_grad_kernel` backward (dW + dstate from projection backward). +- Step 9: `mag_concat_qdir` atom-shift inline. +- Step 10: `quantile_q_select` / `iqn_dual_head` investigation. +- Step 11: Adam wireup for `w_aux_to_q_dir`. +- Step 12: capture-graph integrity verification. +- Phase C1: collector W ptr setter. + +Action selection now uses shifted atom positions everywhere `compute_expected_q` runs. C51 loss + grad still see unshifted positions until Step 7-8. This is an architecturally consistent checkpoint: Q-eval shift is end-to-end; loss-side shift will land atomically with grad-side in Steps 7-8 to avoid the gradient mismatch trap (`feedback_no_partial_refactor.md`).