feat(sp22-vnext): Phase D — 12-weight W atom-shift (4 actions × 3 outcomes)
THE K=3 HEAD NOW DIRECTLY MODULATES Q-TARGETS. Phase D extends the
Phase 3 atom-shift mechanism from per-action W[4] reading single state
slot 121 to per-(action, outcome) W[4, 3] = 12 weights reading 3 state
slots [121..124).
Mathematical change: shift[a, b] now sums Σ_k W[a*K+k] × state[121+k]
= W[a, 0]*p_Profit + W[a, 1]*p_Stop + W[a, 2]*p_Timeout.
5 atom-shift kernel sites updated (all coordinated):
1. experience_kernels.cu::compute_expected_q (replay path)
2. experience_kernels.cu::mag_concat_qdir (rollout path)
3. experience_kernels.cu::quantile_q_select
4. c51_loss_kernel.cu loss numerator (next_state CVaR side)
5. c51_loss_kernel.cu Bellman target (online + target combined)
Each site replaces `W[a] × state_121` with `Σ_k W[a*K+k] × state[121+k]`
unrolled 3 times. State hoist points lift 1 scalar → 3-element array.
5 supporting kernel updates:
- aux_w_prior_init_kernel.cu: writes 12 K=3 structural priors instead
of 4. Spec prior matrix:
Short × {Profit=+0.5, Stop=-0.5, Timeout=0}
Hold × {Profit= 0, Stop=+0.5, Timeout=0}
Long × {Profit=+0.5, Stop=-0.5, Timeout=0}
Flat × {Profit= 0, Stop=+0.5, Timeout=0}
Block dim bumped 4 → 12.
- c51_aux_dw_kernel.cu: grid bumped (4,1,1) → (b0_size×K=12,1,1).
blockIdx.x decoded as (a, k); reads state slot 121+k, writes
dw_aux[a*K + k]. New kernel arg aux_outcome_k=3.
- adam_w_aux_kernel.cu: W_AUX_DIM 4 → 12. Block dim 12 threads.
Trainer-side buffer resizes:
w_aux_to_q_dir [4] → [12]
adam_m_w_aux [4] → [12]
adam_v_w_aux [4] → [12]
dw_aux_buf [4] → [12]
Rust launcher updates:
- c51_aux_dw_kernel launch: grid (4,1,1) → (12,1,1) + new aux_kto arg
- adam_w_aux_kernel launch: block (4,1,1) → (12,1,1)
- aux_w_prior_init launch: block (4,1,1) → (12,1,1)
Cold-start gracefulness preserved: state[121..124] = 0.0 at step 0
(no K=3 prediction yet from C-1 producer). Σ_k W[a*K+k] × 0 = 0 →
zero atom-shift across all actions. After step 1+ when C-1 producer
fires, real softmax probs activate the prior W's structural bias and
Adam refines from there.
End-to-end K=3 → Q-target chain now active:
K=3 fwd (B3/B4) → softmax → C-1 producer → prev_aux_outcome_probs
→ C-2 state gather → state[121..124) → Phase D atom-shift
→ Q-target z_n + Σ_k W[a*K+k] × prob[k]
→ Bellman target + argmax + action_select all see aux's outcome
prediction.
The K=3 head now influences policy via TWO paths: state input (Phase
C-2) AND Q-target modulation (Phase D).
Verification:
- cargo check -p ml clean.
- cargo test -p ml --lib → 1016/0 green.
Remaining vNext work:
- Phase E: dW backward gradient validation tests
- Phase F: Validation smoke at structural prior — decisive spec test
- B5b-2 (deferred): collector trade plan launch
Audit: docs/dqn-wire-up-audit.md Phase D section.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
93aa4cd6a2