feat(sp22): H6 Phase 3 α — adaptive W via dW backward + Adam (B9 Steps 8+11)
Step 8 — c51_aux_dw_kernel (new):
- Per-action block tree-reduce: grid=(4,1,1), block=(256,1,1). One block
per W index a, tree-reduces dW[a] across batch via warp shuffle + shmem.
Zero atomicAdd per pearl_no_atomicadd.
- Per-sample contributions:
a == a_d: dW[a] += inv_batch × isw × (SP_b/dz) × state_121
a == a*: dW[a] += inv_batch × isw × (-γ(1-done)) × (SP_b/dz) × next_state_121
c51_loss_kernel forward — new scratch outputs:
- aux_target_a_dir_buf[B] (i32): saves best_next_a for d==0 after Step c
sampling.
- aux_proj_logdiff_dir_buf[B] (f32): saves SP_b = Σ_n p_target_n ×
(current_lp[upper_n] - current_lp[lower_n]) after Step d's projection
via re-derivation of lower_n/upper_n (matching Huber compression +
clamp arithmetic of block_bellman_project_f).
Step 11 — adam_w_aux_kernel (new):
- Standard Adam with bias correction, grid=(1,1,1), block=(4,1,1).
- Graph-capture-safe: lr via self.lr_dev_ptr pointer arg; step via
self.ptrs.t_buf pointer arg (matches main Adam pattern). beta/eps
as value args from sp5_isv_slots constants.
- Bias-correction denominator floored at 1e-30 to avoid /0.
Trainer wiring (submit_adam_ops):
- launch_c51_aux_dw + launch_adam_w_aux added right after
launch_adam_update. Both inside the captured adam_child graph.
- New trainer fields: aux_target_a_dir_buf, aux_proj_logdiff_dir_buf,
c51_aux_dw_kernel, adam_w_aux_kernel. Cubin statics SP22_C51_AUX_DW_CUBIN
+ SP22_ADAM_W_AUX_CUBIN added.
NULL-safety:
- aux_shift_active=false in c51_loss_kernel forward → both scratch
buffers stay at alloc_zeros 0 → dW reads 0 → Adam W is a no-op.
- aux_target_a_dir_out / aux_proj_logdiff_dir_out are NULL-tolerant.
Deferred (deliberate scope):
- dL/dstate_121 backward (c51 → aux head): refinement, not correctness;
aux head trains via own supervised CE loss.
- Phase C1 collector W ptr setter.
- Phase D (eval-side aux infrastructure).
Verification:
- cargo build -p ml --lib: 0 errors, 21 pre-existing warnings.
- nvcc full recompile clean (1m05s for sm_89 target).
- All forward atom-shift consumers + adaptive W backward + Adam now wired.
End-state: adaptive W trains from structural prior [-0.5, 0, +0.5, 0]
via projection log-diff gradient. Aux head trains independently via
supervised CE. Together they form learned cross-coupling from aux
direction predictions to dir-branch Q distribution shifts. Smoke can
now measure adaptive W's effect on WR.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
72
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
Normal file
72
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
Normal file
@@ -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 <cuda_runtime.h>
|
||||
|
||||
#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);
|
||||
}
|
||||
132
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
Normal file
132
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
Normal file
@@ -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 <cuda_runtime.h>
|
||||
|
||||
#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;
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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<f32>,
|
||||
/// 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<i32>,
|
||||
/// 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<f32>,
|
||||
/// 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::<i32>(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::<f32>(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),
|
||||
|
||||
Reference in New Issue
Block a user