feat(sp22): H6 Phase 3 α — atom-shift wired through compute_expected_q (B9 Step 5+6)
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user