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>
This commit is contained in:
jgrusewski
2026-05-14 11:20:56 +02:00
parent d2331f2f62
commit 93aa4cd6a2
7 changed files with 335 additions and 165 deletions

View File

@@ -18151,3 +18151,78 @@ The K=3 head's input is now [h_s2_aux (256) || plan_params (6)] = 262-dim, match
- `cargo test -p ml --lib`**1016/0 green**.
**Phase D next**: 12-weight W atom-shift (4 actions × 3 outcomes). Extends the existing SP22 Phase 3 atom-shift mechanism from 4 weights (per-action) to 12 weights (per-action × per-outcome). The K=3 head's softmax probs (now plan-conditioned at training time) modulate Q-target atom positions per action.
#### Phase D — 12-weight W atom-shift (4 actions × 3 outcomes) (2026-05-14)
**The K=3 head's predictions now directly modulate Q-target atom positions.** 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**: the atom-shift for action `a` at sample `b` changes from:
```
shift[a, b] = W[a] × state[121, b]
```
to:
```
shift[a, b] = Σ_k W[a*K + k] × state[121+k, b]
= 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-batch path): the per-atom inner loop's `aux_atom_shift` accumulator now sums Σ_k W[a*K+k] × state[121+k] for k ∈ [0, 3). Replaces single-multiplication scalar with #pragma-unrolled 3-term sum.
2. **`experience_kernels.cu::mag_concat_qdir`** (rollout-step path, ~line 5820): hoisted `state_121` scalar replaced by `state_outcome[3]` array (3-element local). Per-action eq adjustment becomes 3-term unrolled sum.
3. **`experience_kernels.cu::quantile_q_select`** (~line 6488): per-action `q_blended` adjustment becomes 3-term unrolled sum.
4. **`c51_loss_kernel.cu`** loss numerator + Bellman target sites (~lines 836, 1094, 1247): hoisted `state_121` + `next_state_121` scalars replaced by `state_outcome[3]` + `next_state_outcome[3]` arrays. Both `aux_shift_online` (current state) and `aux_shift_target` (next state) computed as 3-term unrolled sums.
**5 supporting kernel updates**:
- **`aux_w_prior_init_kernel.cu`**: writes 12 K=3 structural prior values instead of 4. Spec's prior matrix:
```
Profit Stop Timeout
Short +0.5 -0.5 0.0
Hold 0.0 +0.5 0.0
Long +0.5 -0.5 0.0
Flat 0.0 +0.5 0.0
```
Block dim bumped 4 → 12.
- **`c51_aux_dw_kernel.cu`**: grid bumped from `(b0_size=4)` to `(b0_size × K = 12)`. Each block decodes `blockIdx.x → (a, k)`, reads state slot `aux_dir_prob_index + k`, writes `dw_aux[a*K + k]`. New kernel arg `aux_outcome_k` (= 3 in production).
- **`adam_w_aux_kernel.cu`**: `W_AUX_DIM` bumped 4 → 12. Block dim 12 threads. No other change — Adam math is element-wise so doubling-plus elements just adds threads.
**Trainer-side buffer resizes** (gpu_dqn_trainer.rs):
- `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)``(b0×K=12,1,1)`, new arg `aux_kto_i32 = 3`
- `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**: at rollout step 0 the state slots [121..124) are 0.0 (no K=3 prediction yet from Phase C-1 producer). So `Σ_k W[a*K+k] × 0 = 0` → atom-shift is zero across all actions → no Q-target modulation. After step 1+ when the K=3 producer fires, state slots get real softmax probs and the prior W (and Adam-refined W) modulate Q targets as designed.
**End-to-end K=3 → Q-target chain now active**:
```
K=3 forward (B3 rollout / B4 replay) → softmax probs
→ C-1 producer → prev_aux_outcome_probs [N, 3]
→ C-2 state gather → state[121..124)
→ Phase D atom-shift (5 kernel sites) → Q-target z_n + Σ_k W[a*K+k] × prob[k]
→ Bellman projection / argmax / action select biased by aux's outcome prediction
```
The K=3 head now influences policy behavior via TWO paths:
1. **State input** (Phase C-2): policy reads softmax probs as features in next-step state input.
2. **Q-target modulation** (Phase D): each action's effective Q-target atom positions shift by a learned W×outcome combination, biasing both the Bellman target and action selection.
**Verification**:
- `cargo check -p ml` clean (21 warnings, none new).
- `cargo test -p ml --lib`**1016/0 green**.
**Remaining vNext work**:
- **Phase E**: dW backward gradient validation — the Phase D dW kernel emits 12 partials; Phase E would add unit tests confirming the gradient direction is correct + Adam updates the prior in the right direction.
- **Phase F**: Validation smoke at structural prior — first end-to-end smoke run with the full K=3 + plan-conditioning + atom-shift stack. Decisive test of the spec's hypothesis.
- **B5b-2** (deferred): Collector trade plan launch (resolves train/inference asymmetry from B5b).