diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 105421c55..b72cf6c90 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -827,6 +827,17 @@ fn main() { // projection backward (Phase 3-final B9 Step 8). Pure 4-thread // single-block; trivial cost. "aux_w_prior_init_kernel.cu", + // SP22 H6 Phase 3 α Step 8 (2026-05-13): backward dW kernel for + // w_aux_to_q_dir [4]. Reads scratch from c51_loss_kernel forward + // (aux_target_a_dir + aux_proj_logdiff_dir) and computes per-action + // gradient via block tree-reduce (grid=(4,1,1), block=(256,1,1)) + // — no atomicAdd per pearl_no_atomicadd. Per-block writes dw_aux[a]. + "c51_aux_dw_kernel.cu", + // SP22 H6 Phase 3 α Step 11 (2026-05-13): Adam update kernel for + // w_aux_to_q_dir [4]. Standard Adam with bias correction; 4 + // threads single block; launched once per training step OUTSIDE + // the captured forward/backward graphs. + "adam_w_aux_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/adam_w_aux_kernel.cu b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu new file mode 100644 index 000000000..0aa4679ac --- /dev/null +++ b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu @@ -0,0 +1,72 @@ +// crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu +// +// SP22 H6 Phase 3 α Step 11 — Adam optimizer update for w_aux_to_q_dir [4]. +// +// Standard Adam (Kingma & Ba 2015) with bias correction: +// m_t = β1 × m_{t-1} + (1 - β1) × g_t +// v_t = β2 × v_{t-1} + (1 - β2) × g_t² +// m_hat = m_t / (1 - β1^t) +// v_hat = v_t / (1 - β2^t) +// W_t = W_{t-1} - lr × m_hat / (√v_hat + ε) +// +// Discipline +// ────────── +// - 4 threads, single block — trivial cost (<1µs). Launched once per +// training step OUTSIDE the captured forward/backward graph (mirrors +// the existing main-params Adam launch). +// - β1, β2, ε, lr passed by value from host — host-stable across captured +// launches per pearl_no_host_branches_in_captured_graph (this kernel is +// not captured anyway). +// - step is i32 (1-indexed) per the trainer's existing Adam step counter +// semantics. `__powf` for the bias-correction denominator (sufficient +// precision for the few-thousand-step horizon). +// - Per pearl_first_observation_bootstrap.md: m and v default to 0 at +// alloc_zeros init; the first step's bias-correction divides by +// (1 - β^1) = (1 - β), not (1 - β^0) = 0 → no NaN. +// - No HtoD inside kernel per feedback_no_htod_htoh_only_mapped_pinned.md. + +#include + +#define W_AUX_DIM 4 + +/* Graph-capture-safe signature: `step` and `lr` are POINTERS to + * device-mapped pinned scratch the host updates each training step + * (mirrors the main Adam launcher's `t_buf` + `lr_dev_ptr` pattern). + * The captured launch records the pointers; the kernel reads the + * current values at replay time. Beta/eps are constants — host-stable + * across all replays per pearl_no_host_branches_in_captured_graph. */ +extern "C" __global__ void adam_w_aux_update( + float* __restrict__ w, /* [W_AUX_DIM=4] in-place update */ + const float* __restrict__ dw, /* [W_AUX_DIM=4] gradient from c51_aux_dw_kernel */ + float* __restrict__ m, /* [W_AUX_DIM=4] first moment in-place */ + float* __restrict__ v, /* [W_AUX_DIM=4] second moment in-place */ + const float* __restrict__ lr_ptr, /* [1] device-mapped pinned f32 */ + float beta1, + float beta2, + float eps, + const int* __restrict__ step_ptr /* [1] device-mapped pinned i32, 1-indexed */ +) { + int a = threadIdx.x; + if (a >= W_AUX_DIM) return; + + float lr = lr_ptr[0]; + int step = step_ptr[0]; + + float g = dw[a]; + float m_new = beta1 * m[a] + (1.0f - beta1) * g; + float v_new = beta2 * v[a] + (1.0f - beta2) * g * g; + m[a] = m_new; + v[a] = v_new; + + /* Bias-corrected moments. step is at least 1 on the first invocation + * (trainer increments before submit_adam_ops), so (1 - β^step) > 0. */ + float bc1 = 1.0f - __powf(beta1, (float)step); + float bc2 = 1.0f - __powf(beta2, (float)step); + /* Floors at fp32 epsilon to avoid /0 in degenerate first-step + * configurations (β1 ≈ 0 or β2 ≈ 0). */ + bc1 = fmaxf(bc1, 1e-30f); + bc2 = fmaxf(bc2, 1e-30f); + float m_hat = m_new / bc1; + float v_hat = v_new / bc2; + w[a] = w[a] - lr * m_hat / (sqrtf(v_hat) + eps); +} diff --git a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu new file mode 100644 index 000000000..a24c839de --- /dev/null +++ b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu @@ -0,0 +1,132 @@ +// crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu +// +// SP22 H6 Phase 3 α Step 8 — backward dW kernel for w_aux_to_q_dir [4]. +// +// Computes gradient of C51 loss w.r.t. the per-action atom-shift weights: +// dW[a] = Σ_b indicator(a_d_b == a) × isw_b × inv_batch × (SP_b / Δz_b) +// × state_121_b +// + Σ_b indicator(a*_b == a) × isw_b × inv_batch × (-γ_eff_b × (1-done_b)) +// × (SP_b / Δz_b) × next_state_121_b +// +// where: +// - SP_b = projection log-diff sum (saved by c51_loss_kernel forward): +// SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n]) +// - a_d_b = taken direction action (= a0, computed from actions[b]) +// - a*_b = sampled target action (saved by c51_loss_kernel as +// aux_target_a_dir[b]) +// - Δz_b = per_sample_support[b * 12 + 0 * 3 + 2] (dir branch dz) +// - γ_eff_b = effective gamma for dir branch from gamma_buf[b] (γ^n_steps, +// ISV gamma_dir read inside) +// - state_121_b = batch_states[b * state_dim + AUX_DIR_PROB_INDEX] +// - next_state_121_b = next_batch_states[b * state_dim + AUX_DIR_PROB_INDEX] +// +// Output: dw_aux_buf[4] f32 (zeroed before launch — overwrite semantics). +// +// Discipline +// ────────── +// - Per pearl_no_atomicadd: ONE BLOCK PER ACTION (grid=(4,1,1), block=(256,1,1)); +// each block tree-reduces across batch samples via shmem. No atomicAdd. +// - Per feedback_no_htod_htoh_only_mapped_pinned.md: all reads from device / +// mapped pinned buffers; no HtoD inside kernel. +// - Per pearl_no_host_branches_in_captured_graph: launch happens OUTSIDE the +// captured training graph (mirrors the existing post-c51_grad launches). +// - NULL-safety: when forward didn't populate scratch (aux_shift_active=false), +// both aux_target_a_dir and aux_proj_logdiff_dir stay at alloc_zeros +// sentinels; gradient stays at 0 → Adam step on W is a no-op. + +#include + +#define BLOCK_THREADS 256 +#define NUM_WARPS (BLOCK_THREADS / 32) +#define MAX_DIR_ACTIONS 4 + +extern "C" __global__ void c51_aux_dw_kernel( + /* Inputs from c51_loss_kernel forward scratch */ + const int* __restrict__ aux_target_a_dir, /* [B] sampled a* per sample */ + const float* __restrict__ aux_proj_logdiff_dir,/* [B] SP_b per sample */ + + /* Inputs that the gradient formula needs */ + const int* __restrict__ actions, /* [B] factored action codes */ + const float* __restrict__ is_weights, /* [B] PER importance weights (clamped to 10) */ + const float* __restrict__ dones, /* [B] 0 or 1 */ + const float* __restrict__ gamma_buf, /* [B] effective γ^n_steps */ + const float* __restrict__ per_sample_support, /* [B, 4, 3] for Δz */ + const float* __restrict__ batch_states, /* [B, state_dim] for state_121 */ + const float* __restrict__ next_batch_states, /* [B, state_dim] for next_state_121 */ + + /* Output */ + float* __restrict__ dw_aux, /* [b0_size=4] gradient — written by tid==0 of each block */ + + /* Config */ + int batch_size, + int b0_size, + int b1_size, + int b2_size, + int b3_size, + int aux_dir_prob_index, /* = 121 */ + int state_dim /* = 128 */ +) { + int a = blockIdx.x; /* which W index this block accumulates */ + int tid = threadIdx.x; + if (a >= b0_size) return; + if (a >= MAX_DIR_ACTIONS) return; + + float inv_batch = 1.0f / (float)batch_size; + + /* Decode a_d (taken dir action = a0) inverse to c51_loss_kernel's decode: + * a0 = factored / (b1 * b2 * b3) + * Note: b0_size=4 in production but read from kernel arg for ABI safety. */ + int divisor_a0 = b1_size * b2_size * b3_size; + if (divisor_a0 <= 0) divisor_a0 = 1; + + /* Per-thread partial sum across samples assigned to this thread. */ + float local_sum = 0.0f; + for (int b = tid; b < batch_size; b += BLOCK_THREADS) { + int factored = actions[b]; + if (factored < 0) factored = 0; + int a_d = factored / divisor_a0; + if (a_d < 0 || a_d >= b0_size) a_d = 0; + + int a_star = aux_target_a_dir[b]; + if (a_star < 0 || a_star >= b0_size) a_star = 0; + + float sp = aux_proj_logdiff_dir[b]; + float isw = fminf(is_weights[b], 10.0f); + float dz = per_sample_support[b * 12 + 0 * 3 + 2]; + if (dz < 1e-7f) continue; /* degenerate dir support: no gradient (matches forward skip) */ + + float dL_dDelta_online = sp / dz; + float gamma_eff = gamma_buf[b]; + float done = dones[b]; + float dL_dDelta_target = -gamma_eff * (1.0f - done) * dL_dDelta_online; + + if (a == a_d) { + float s121 = batch_states[(long long)b * state_dim + aux_dir_prob_index]; + local_sum += inv_batch * isw * dL_dDelta_online * s121; + } + if (a == a_star) { + float ns121 = next_batch_states[(long long)b * state_dim + aux_dir_prob_index]; + local_sum += inv_batch * isw * dL_dDelta_target * ns121; + } + } + + /* Block tree-reduce (per pearl_no_atomicadd): warp shuffle + shmem. */ + __shared__ float shmem_warp[NUM_WARPS]; + unsigned int mask = 0xFFFFFFFF; + int lane = tid & 31; + int warp = tid >> 5; + + /* Warp-level reduce. */ + for (int offset = 16; offset > 0; offset >>= 1) { + local_sum += __shfl_xor_sync(mask, local_sum, offset); + } + if (lane == 0) shmem_warp[warp] = local_sum; + __syncthreads(); + + /* Final reduce across warps (sequential in tid==0). */ + if (tid == 0) { + float total = 0.0f; + for (int w = 0; w < NUM_WARPS; w++) total += shmem_warp[w]; + dw_aux[a] = total; + } +} diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index d8b9d463c..b6b4fb9ed 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -648,7 +648,28 @@ extern "C" __global__ void c51_loss_batched( const float* __restrict__ batch_states, /* [B, state_dim] f32, or NULL */ const float* __restrict__ next_batch_states, /* [B, state_dim] f32, or NULL */ int aux_dir_prob_index, /* = SL_PADDING_START = 121 */ - int state_dim /* = STATE_DIM = 128 */ + int state_dim, /* = STATE_DIM = 128 */ + + /* ── SP22 H6 Phase 3 α Step 8 backward scratch outputs ────────────── + * Per-sample scratch produced by the FORWARD pass for the dir-branch + * (d == 0) only; consumed by `c51_aux_dW_kernel` to compute W's + * gradient without re-running the projection arithmetic. + * + * `aux_target_a_dir_out[b]` (i32, [B]): best_next_a sampled in Step c + * for d == 0. The dW[a*] contribution accumulates per this index. + * + * `aux_proj_logdiff_dir_out[b]` (f32, [B]): SP_b = Σ_n p_target_n × + * (shmem_current_lp[upper_n] - shmem_current_lp[lower_n]) computed + * during the d==0 Bellman projection. dW gradient = (SP_b / Δz) × + * state_121 for the online side, × (-γ*(1-done)*next_state_121) for + * the target side. + * + * NULL-tolerant: when either is NULL OR aux_shift_active is false, + * no write occurs. Both default to alloc_zeros at trainer init so + * test scaffolds without Step 8 wiring see 0 → bit-identical pre- + * Phase-3-α gradient flow downstream. */ + int* __restrict__ aux_target_a_dir_out, /* [B] i32, or NULL */ + float* __restrict__ aux_proj_logdiff_dir_out /* [B] f32, or NULL */ ) { extern __shared__ float shmem_f[]; @@ -1166,6 +1187,16 @@ extern "C" __global__ void c51_loss_batched( __syncthreads(); int best_next_a = sampled_action; + /* SP22 H6 Phase 3 α Step 8 backward scratch save: + * record sampled target action for d == 0 only. Consumed by + * c51_aux_dW_kernel to look up W[best_next_a] for the target + * side dW gradient. */ + if (d == 0 && tid == 0 + && aux_shift_active + && aux_target_a_dir_out != NULL) { + aux_target_a_dir_out[sample_id] = best_next_a; + } + /* Target distribution for sampled action */ for (int j = tid; j < num_atoms; j += BLOCK_THREADS) shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * num_atoms + j] - shmem_proj[j]; @@ -1207,6 +1238,57 @@ extern "C" __global__ void c51_loss_batched( d, a0, isv_signals); __syncthreads(); + /* SP22 H6 Phase 3 α Step 8 backward scratch save: + * SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n]) + * for d == 0 only. Inputs: + * - shmem_proj[n] = target probabilities (input to projection, + * unchanged by block_bellman_project_f) + * - shmem_current_lp[k] = online log-probs for taken action a_d, + * intact from Step a (lines 925-929; Steps b/c/d don't touch + * shmem_current_lp by name — they use shmem_lp / shmem_proj as + * scratch). + * - effective_reward already includes Δ_target + Δ_online so + * b_pos_n maps target atom n to the correct bin on the + * UNSHIFTED support — matching the projection's actual bin + * selection. + * + * Per-thread local accumulation; final block_reduce_sum_f + * produces sample-scalar SP_b written by tid==0. Re-derives + * lower_n / upper_n by re-running the projection's per-atom + * arithmetic (cheap: ~10 fmul/atom per thread). */ + if (d == 0 && aux_shift_active + && aux_proj_logdiff_dir_out != NULL) { + float local_sp = 0.0f; + float v_range_for_clip = v_max - v_min; + (void)v_range_for_clip; + float gamma_d = gamma_eff; + for (int n = tid; n < num_atoms; n += BLOCK_THREADS) { + float z_n = shmem_support[n]; + float t_z = effective_reward + gamma_d * z_n * (1.0f - done); + /* Match block_bellman_project_f's Huber compression + clamp. */ + if (t_z < 0.0f) { + t_z = -10.0f * (1.0f - expf(t_z / 10.0f)); + } + t_z = fminf(fmaxf(t_z, v_min), v_max); + float b_pos = (t_z - v_min) / delta_z; + b_pos = fminf(b_pos, (float)(num_atoms - 1) - 0.001f); + b_pos = fmaxf(b_pos, 0.001f); + int lower = (int)floorf(b_pos); + int upper = lower + 1; + lower = max(min(lower, num_atoms - 1), 0); + upper = max(min(upper, num_atoms - 1), 0); + float p_target_n = shmem_proj[n]; + float lp_upper = shmem_current_lp[upper]; + float lp_lower = shmem_current_lp[lower]; + local_sp += p_target_n * (lp_upper - lp_lower); + } + float SP_b = block_reduce_sum_f(local_sp, shmem_reduce, tid); + if (tid == 0) { + aux_proj_logdiff_dir_out[sample_id] = SP_b; + } + __syncthreads(); /* barrier before downstream label-smoothing reads */ + } + /* Label smoothing — health-coupled. * * eps_eff = LABEL_SMOOTHING_BASE × (1 − health) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 4682986a1..babe9509c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -693,6 +693,21 @@ pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = 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 α Step 8 (2026-05-13): backward dW kernel for +/// `w_aux_to_q_dir [4]`. Reads scratch buffers from c51_loss_kernel +/// forward (`aux_target_a_dir` + `aux_proj_logdiff_dir`) and computes +/// per-action gradient via block tree-reduce (grid=(4,1,1), +/// block=(256,1,1)). Launched OUTSIDE captured graphs. +pub(crate) static SP22_C51_AUX_DW_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/c51_aux_dw_kernel.cubin")); + +/// SP22 H6 Phase 3 α Step 11 (2026-05-13): Adam update kernel for +/// `w_aux_to_q_dir [4]`. Standard Adam with bias correction. 4 threads +/// single block. Launched once per training step OUTSIDE captured +/// graphs. +pub(crate) static SP22_ADAM_W_AUX_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/adam_w_aux_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- @@ -6921,6 +6936,26 @@ pub struct GpuDqnTrainer { /// dloss/dz_n_effective); read by the Adam-step update. /// [b0_size=4] f32. dw_aux_buf: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α Step 8 — per-sample sampled target action for d == 0. + /// Written by c51_loss_kernel forward at the Expected SARSA sampling + /// site (`best_next_a` for dir branch). Read by `c51_aux_dw_kernel` + /// to look up W[a*] for the target-side dW contribution. + /// Shape `[batch_size]` i32. alloc_zeros default = 0 (Short action; + /// harmless when aux_shift inactive since SP_b stays at 0 too). + aux_target_a_dir_buf: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α Step 8 — per-sample projection log-diff sum + /// for d == 0. Written by c51_loss_kernel forward right after the + /// Bellman projection completes: + /// `SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])` + /// Read by `c51_aux_dw_kernel` to compute + /// `dL/dΔ_online = SP_b / Δz`, `dL/dΔ_target = -γ*(1-done)*dL/dΔ_online` + /// Shape `[batch_size]` f32. alloc_zeros default = 0 → no gradient + /// when aux_shift inactive. + aux_proj_logdiff_dir_buf: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 3 α Step 8 — kernel handle for `c51_aux_dw_kernel`. + c51_aux_dw_kernel: cudarc::driver::CudaFunction, + /// SP22 H6 Phase 3 α Step 11 — kernel handle for `adam_w_aux_update`. + adam_w_aux_kernel: cudarc::driver::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 @@ -20874,6 +20909,43 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!( "sp22-h6-phase3 α: alloc dw_aux_buf: {e}" )))?; + // SP22 H6 Phase 3 α Step 8 (2026-05-13): per-sample scratch buffers + // for the W gradient backward. aux_target_a_dir saves best_next_a + // for d==0 from the forward; aux_proj_logdiff_dir saves the + // per-sample SP scalar. + let aux_target_a_dir_buf = stream + .alloc_zeros::(config.batch_size) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 8: alloc aux_target_a_dir_buf: {e}" + )))?; + let aux_proj_logdiff_dir_buf = stream + .alloc_zeros::(config.batch_size) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 8: alloc aux_proj_logdiff_dir_buf: {e}" + )))?; + // Load Step 8 + Step 11 kernels. + let c51_aux_dw_kernel = { + let module = stream.context() + .load_cubin(SP22_C51_AUX_DW_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 8: c51_aux_dw cubin: {e}" + )))?; + module.load_function("c51_aux_dw_kernel") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 8: c51_aux_dw_kernel load: {e}" + )))? + }; + let adam_w_aux_kernel = { + let module = stream.context() + .load_cubin(SP22_ADAM_W_AUX_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 11: adam_w_aux cubin: {e}" + )))?; + module.load_function("adam_w_aux_update") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 11: adam_w_aux_update load: {e}" + )))? + }; // SP22 H6 Phase 3 α (atom-shift design 2026-05-13): structural-prior // init for w_aux_to_q_dir. Writes per-action prior @@ -25288,6 +25360,10 @@ impl GpuDqnTrainer { adam_m_w_aux, adam_v_w_aux, dw_aux_buf, + aux_target_a_dir_buf, + aux_proj_logdiff_dir_buf, + c51_aux_dw_kernel, + adam_w_aux_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, @@ -30124,6 +30200,114 @@ impl GpuDqnTrainer { } + /// SP22 H6 Phase 3 α Step 8 (2026-05-13): dW kernel launcher. + /// Reads scratch buffers populated by c51_loss_kernel forward + /// (aux_target_a_dir + aux_proj_logdiff_dir for d == 0 only) and + /// writes dw_aux_buf[4] via per-action block tree-reduce. No + /// atomicAdd per pearl_no_atomicadd. + /// + /// Graph-capture-safe: all kernel args are device pointers (the + /// scratch buffers / params buffer / state buffers are stable across + /// captured replays; their CONTENTS update each step via the captured + /// forward and the externally-written batch upload). + fn launch_c51_aux_dw(&self) -> Result<(), MLError> { + let b = self.config.batch_size as i32; + let b0 = self.config.branch_0_size as i32; + let b1 = self.config.branch_1_size as i32; + let b2 = self.config.branch_2_size as i32; + let b3 = self.config.branch_3_size as i32; + 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; + + let aux_target_a_ptr = self.aux_target_a_dir_buf.raw_ptr(); + let aux_proj_logdiff_ptr = self.aux_proj_logdiff_dir_buf.raw_ptr(); + let actions_ptr = self.ptrs.actions_buf; + let is_weights_ptr = self.ptrs.is_weights_buf; + let dones_ptr = self.ptrs.dones_buf; + let gamma_buf_ptr = self.gamma_buf.raw_ptr(); + let per_sample_support_ptr = self.per_sample_support_ptr; + let states_ptr = self.ptrs.states_buf; + let next_states_ptr = self.ptrs.next_states_buf; + let dw_aux_ptr = self.dw_aux_buf.raw_ptr(); + + unsafe { + self.stream + .launch_builder(&self.c51_aux_dw_kernel) + .arg(&aux_target_a_ptr) + .arg(&aux_proj_logdiff_ptr) + .arg(&actions_ptr) + .arg(&is_weights_ptr) + .arg(&dones_ptr) + .arg(&gamma_buf_ptr) + .arg(&per_sample_support_ptr) + .arg(&states_ptr) + .arg(&next_states_ptr) + .arg(&dw_aux_ptr) + .arg(&b) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) + .arg(&aux_dir_prob_index) + .arg(&state_dim_i32) + .launch(LaunchConfig { + grid_dim: (4, 1, 1), // one block per W index (b0_size=4) + block_dim: (256, 1, 1), // tree-reduce across batch + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 8: c51_aux_dw_kernel launch: {e}" + )))?; + } + Ok(()) + } + + /// SP22 H6 Phase 3 α Step 11 (2026-05-13): Adam-update launcher for + /// `w_aux_to_q_dir [4]`. Reads `dw_aux_buf` (written by + /// `launch_c51_aux_dw`) and updates W in-place using Adam moments. + /// + /// Graph-capture-safe: lr and step counter are read via + /// device-mapped pinned pointers (matching the main Adam pattern). + /// beta1/beta2/eps are constants from sp5_isv_slots — host-stable. + fn launch_adam_w_aux(&self) -> Result<(), MLError> { + use crate::cuda_pipeline::sp5_isv_slots::{ + ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE, + }; + + let w_ptr = self.w_aux_to_q_dir.raw_ptr(); + let dw_ptr = self.dw_aux_buf.raw_ptr(); + let m_ptr = self.adam_m_w_aux.raw_ptr(); + let v_ptr = self.adam_v_w_aux.raw_ptr(); + let lr_ptr = self.lr_dev_ptr; + let t_ptr = self.ptrs.t_buf; + let beta1 = ADAM_BETA1_BASE as f32; + let beta2 = ADAM_BETA2_BASE as f32; + let eps = ADAM_EPS_BASE as f32; + + unsafe { + self.stream + .launch_builder(&self.adam_w_aux_kernel) + .arg(&w_ptr) + .arg(&dw_ptr) + .arg(&m_ptr) + .arg(&v_ptr) + .arg(&lr_ptr) + .arg(&beta1) + .arg(&beta2) + .arg(&eps) + .arg(&t_ptr) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (4, 1, 1), // one thread per W slot + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6-phase3 α Step 11: adam_w_aux_kernel launch: {e}" + )))?; + } + Ok(()) + } + /// Submit the optimizer phase ops to the stream (used by adam_child capture). /// /// Steps: Adam → unflatten. @@ -30138,6 +30322,15 @@ impl GpuDqnTrainer { // ── 6. Adam update (f32 master weights) ───────────────── self.launch_adam_update()?; + // ── 6a. SP22 H6 Phase 3 α Step 8 + Step 11 (2026-05-13) ── + // Compute dW and Adam-update w_aux_to_q_dir [4] alongside the + // main Adam step. Both kernels are graph-capture-safe (all + // varying inputs via device-mapped pointers; constants via + // value args). The dW kernel reads scratch buffers populated + // by the captured c51_loss_kernel forward earlier in this step. + self.launch_c51_aux_dw()?; + self.launch_adam_w_aux()?; + // ── 6.5. Snapshot grad_buf → prev_grad_buf for next step's vaccine comparison ── // Graph-safe: submit_adam_ops is captured in adam_update child graph. self.graph_safe_copy_f32( @@ -30955,6 +31148,14 @@ impl GpuDqnTrainer { .arg(&self.ptrs.next_states_buf) .arg(&(ml_core::state_layout::AUX_DIR_PROB_INDEX as i32)) .arg(&(ml_core::state_layout::STATE_DIM as i32)) + // ── SP22 H6 Phase 3 α Step 8 backward scratch (2026-05-13) ── + // Per-sample scratch outputs for the W gradient kernel: + // aux_target_a_dir[B] — best_next_a for d==0 per sample + // aux_proj_logdiff_dir[B] — projection log-diff sum per sample + // Both populated by c51_loss_kernel forward d==0 path; + // consumed by c51_aux_dw_kernel post-loss. + .arg(&self.aux_target_a_dir_buf.raw_ptr()) + .arg(&self.aux_proj_logdiff_dir_buf.raw_ptr()) .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (256, 1, 1), diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 449ce7c52..4710f051c 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17096,3 +17096,84 @@ W stays at the structural prior `[-0.5, 0.0, +0.5, 0.0]` (Step 5 init). All grad **Defensible intermediate checkpoint rationale**: Steps 7+9+10 land all FORWARD-path consumers of atom positions. Steps 8+11 add the dW + Adam path. Without Step 8/11, W stays at the structural prior — functionally equivalent to the "fixed-W version" the user previously characterized as "would activate α immediately for action selection but never learn state-dependent corrections." The user explicitly chose the adaptive version, but this checkpoint is structurally honest: all forward consumers see consistent shifts; W is just frozen. A smoke at this checkpoint can answer "does the structural prior alone move WR?" before investing the Step 8+11 effort. The math derivation in `7eae832f2` makes Steps 8+11 a transcription job for the next session. Verification: cargo check -p ml --lib: 0 errors, 21 pre-existing warnings (baseline parity). + +#### Atom-shift backward + Adam complete — Step 8 (c51_aux_dw_kernel) + Step 11 (adam_w_aux_kernel) — 2026-05-13 + +**Step 8 — dW backward kernel** (`c51_aux_dw_kernel.cu`, new): + +Grid (b0_size=4, 1, 1), block (256, 1, 1). One block per W index a; each block tree-reduces dW[a] across all batch samples via warp shuffle + shmem. No atomicAdd per `pearl_no_atomicadd.md`. + +Per-sample contributions accumulated by each block (only the matching action contributes): +- Online side: when `a == a_d_b` (taken direction action a0 from `actions[b]` decode), `dW[a] += inv_batch × isw_b × (SP_b / Δz_b) × state_121_b`. +- Target side: when `a == a*_b` (sampled target action from `aux_target_a_dir[b]`), `dW[a] += inv_batch × isw_b × (-γ_eff_b × (1-done_b)) × (SP_b / Δz_b) × next_state_121_b`. + +Where: +- `SP_b` = projection log-diff sum saved by c51_loss_kernel forward in `aux_proj_logdiff_dir[b]`. Computed as `Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])` for d == 0 only. +- `Δz_b` = `per_sample_support[b * 12 + 0*3 + 2]` (dir branch dz). +- `isw_b` = `min(is_weights[b], 10.0)` matching c51_grad_kernel's PER IS-weight clamp. +- `inv_batch = 1/B` matching c51_grad_kernel's per-sample scaling. + +Degenerate-support skip: when `Δz_b < 1e-7`, contributions for that sample are skipped (matches c51_loss_kernel forward's `branch_degenerate` skip). + +**c51_loss_kernel forward scratch writes** (added in this commit): + +- After Step c's Expected SARSA sampling sets `best_next_a` for d == 0: `aux_target_a_dir_out[sample_id] = best_next_a` (single tid==0 write). +- After Step d's `block_bellman_project_f` returns: re-derive `b_pos_n → lower_n / upper_n` for each target atom n (re-running the projection's per-atom arithmetic; cheap ~10 fmul/atom) and accumulate `local_sp += p_target_n × (current_lp[upper_n] - current_lp[lower_n])`. Block-reduce → `aux_proj_logdiff_dir_out[sample_id]` (tid==0 write). + +The re-derivation matches `block_bellman_project_f`'s exact arithmetic (Huber compression on `t_z < 0`, clamp to [v_min, v_max], floor/ceil index extraction) — guarantees the lower_n/upper_n used in SP computation are the SAME bins the projection wrote mass into. + +**Step 11 — Adam-update kernel** (`adam_w_aux_kernel.cu`, new): + +Grid (1, 1, 1), block (4, 1, 1). One thread per W slot. Standard Adam (Kingma & Ba 2015) with bias correction: +``` +m_t = β1 m_{t-1} + (1-β1) g_t +v_t = β2 v_{t-1} + (1-β2) g_t² +m_hat = m_t / (1 - β1^t) +v_hat = v_t / (1 - β2^t) +W_t = W_{t-1} - lr × m_hat / (√v_hat + ε) +``` + +Graph-capture-safe via device-mapped pointer args: +- `lr_ptr` = `self.lr_dev_ptr` (matches main Adam's lr pointer pattern). +- `step_ptr` = `self.ptrs.t_buf` (device-mapped i32 step counter; host writes via `t_pinned` each step). +- `beta1` / `beta2` / `eps` passed by value from `sp5_isv_slots::{ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE}` (host-stable constants per `pearl_no_host_branches_in_captured_graph`). + +Bias-correction denominator floored at fp32 epsilon (`max(bc, 1e-30)`) to avoid /0 if β1 ≈ 0. + +**Trainer wiring** (`gpu_dqn_trainer.rs::submit_adam_ops`): + +Both launches added inside `submit_adam_ops` right after `launch_adam_update`: +```rust +self.launch_adam_update()?; +self.launch_c51_aux_dw()?; +self.launch_adam_w_aux()?; +``` + +This places W gradient computation + W Adam update INSIDE the captured `adam_child` graph alongside the main Adam step. All inputs to both kernels are stable pointers (the scratch buffers, params, states are updated each step but the pointers don't change). + +**c51_loss_kernel launcher updates** (`launch_c51_loss`): 2 new args at the end (`aux_target_a_dir_buf.raw_ptr()`, `aux_proj_logdiff_dir_buf.raw_ptr()`). + +**New trainer fields**: +- `aux_target_a_dir_buf: CudaSlice` shape `[batch_size]`. alloc_zeros default = 0. +- `aux_proj_logdiff_dir_buf: CudaSlice` shape `[batch_size]`. alloc_zeros default = 0. +- `c51_aux_dw_kernel: CudaFunction`. +- `adam_w_aux_kernel: CudaFunction`. + +Both scratch buffers default to 0 → when `aux_shift_active` is false in c51_loss_kernel forward (NULL W or NULL states), nothing is written → next step's dW kernel sees 0 SP → dW stays at 0 → Adam W is a no-op. So Steps 8+11 are NULL-safe at every level. + +**Scope NOT covered in this commit (deliberate deferral)**: + +- `dL/dstate_121`: the gradient path c51_loss → state_121 → aux head weights. This would refine the aux head to maximize C51 task performance. Currently aux head trains via its OWN supervised CE loss (next-bar direction); the missing path is a REFINEMENT, not a correctness issue. Adding it requires writing to a `dbatch_states` buffer + integrating with the existing state gradient infrastructure — non-trivial. Deferred per scope discipline. +- Phase C1 (collector W ptr setter): rollout-time `compute_expected_q` + `quantile_q_select` still pass NULL W. Until C1, W's effect on training is internal to the trainer (no rollout-time action distribution shift). Phase D (eval-side) needs C1's setter pattern. +- Phase D (A2 eval-side aux infrastructure): 7 tasks. Eval-time action selection doesn't see atom-shift yet. + +**Verification**: +- `cargo build -p ml --lib`: 0 errors, 21 pre-existing warnings (baseline parity). +- nvcc compiles all new .cu files cleanly (no warnings/errors on c51_aux_dw_kernel.cu, adam_w_aux_kernel.cu, c51_loss_kernel.cu). +- Full kernel recompile took 1m 05s — within expected wall-time for sm_89 target. + +**End-state for trainer-side adaptive α**: + +W now trains via the c51 loss gradient. Starting from structural prior `[-0.5, 0.0, +0.5, 0.0]`, Adam refines W per-action based on observed `state_121 × (loss reduction from atom-shift)`. The aux head trains independently via its supervised next-bar CE; together they form a learned cross-coupling from aux direction predictions to dir-branch Q distribution shifts. + +Next session can dispatch a smoke at this checkpoint to measure whether the adaptive W (vs. fixed structural prior) lifts WR off 50%. If WR moves on the trainer side without rollout-side activation, Phase C1 is the next priority. If not, the hypothesis needs reconsideration.