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>
5d163e0e1d