gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp22-vnext): Phase D — 12-weight W atom-shift (4 actions × 3 outcomes)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu
View File

@@ -1,6 +1,7 @@
// 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].
// SP22 H6 Phase 3 α Step 11 → SP22 H6 vNext Phase D Step 11 — Adam
// optimizer update for w_aux_to_q_dir [b0_size=4 × K=3] = 12 weights.
//
// Standard Adam (Kingma & Ba 2015) with bias correction:
// m_t = β1 × m_{t-1} + (1 - β1) × g_t
@@ -11,9 +12,10 @@
//
// Discipline
// ──────────
// - 4 threads, single block — trivial cost (<1µs). Launched once per
// - 12 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).
// the existing main-params Adam launch). Phase D bumped from 4 to 12
// threads to cover the K=3 per-outcome weights per direction action.
// - β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).
@@ -27,7 +29,10 @@
#include <cuda_runtime.h>
#define W_AUX_DIM 4
/* SP22 H6 vNext Phase D (2026-05-14): W_AUX_DIM bumped 4 → 12. Layout:
* 4 direction actions × 3 outcome classes = 12 weights, indexed
* as `w[a * K + k]` where K=3 = AUX_OUTCOME_K. */
#define W_AUX_DIM 12
/* Graph-capture-safe signature: `step` and `lr` are POINTERS to
* device-mapped pinned scratch the host updates each training step
@@ -36,10 +41,10 @@
* 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 */
float* __restrict__ w, /* [W_AUX_DIM=12] in-place update */
const float* __restrict__ dw, /* [W_AUX_DIM=12] gradient from c51_aux_dw_kernel */
float* __restrict__ m, /* [W_AUX_DIM=12] first moment in-place */
float* __restrict__ v, /* [W_AUX_DIM=12] second moment in-place */
const float* __restrict__ lr_ptr, /* [1] device-mapped pinned f32 */
float beta1,
float beta2,

95
crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu
View File

@@ -1,14 +1,36 @@
// crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu
//
// SP22 H6 Phase 3 α (2026-05-13): structural-prior initialization for
// the Adam-trained `w_aux_to_q_dir [b0_size=4]` weight vector.
// SP22 H6 Phase 3 α (2026-05-13) — DORMANT, superseded by Phase D below.
// SP22 H6 vNext Phase D (2026-05-14): structural-prior initialization for
// the Adam-trained `w_aux_to_q_dir [b0_size=4, K=3] = 12` weight matrix.
//
// Writes the per-action prior values:
// W[0] = Short → -0.5 (aux conviction × position-sign favors Long;
// negate to discourage Short when aux is positive)
// W[1] = Hold → 0.0 (no atom shift for Hold action)
// W[2] = Long → +0.5 (encourage Long when aux predicts up)
// W[3] = Flat → 0.0 (no atom shift for Flat action)
// K=3 outcome-conditioned atom-shift structural prior (12 weights):
// For each direction action a ∈ {Short, Hold, Long, Flat} and each
// outcome k ∈ {Profit, Stop, Timeout}, W[a*K + k] biases the Q-target
// atom position for action a when aux predicts outcome k with high
// probability.
//
// W[Profit] W[Stop] W[Timeout]
// Short (a=0) +0.5 -0.5 0.0
// Hold (a=1) 0.0 +0.5 0.0
// Long (a=2) +0.5 -0.5 0.0
// Flat (a=3) 0.0 +0.5 0.0
//
// Semantic rationale:
// - Trade-directional actions (Short, Long) get atom-shift +0.5 when
// aux predicts Profit (the trade WILL succeed) — encourages taking
// the trade. Get -0.5 when aux predicts Stop (the trade WILL fail)
// — discourages taking the trade. Timeout = 0 (uncertain).
// - Non-directional actions (Hold, Flat) get atom-shift +0.5 when aux
// predicts Stop — encourages staying flat/holding rather than
// entering a doomed trade. Get 0 on Profit (no need to bias toward
// no-action when a profitable trade is available). Timeout = 0.
//
// Replaces the Phase 3 K=2 prior {Short=-0.5, Hold=0, Long=+0.5, Flat=0}
// which was based on directional `p_up` — that hypothesis was FALSIFIED
// (smoke train-xrkb7 @ ebc7144434, WR=0.4346 = baseline). The K=3
// outcome-conditioned prior is the vNext replacement, aligned with the
// system's multi-bar trade decisions.
//
// Rationale (audit doc Phase 3c α revision, spec doc α section):
// state_121 ∈ [-1, +1] is the recentered aux p_up. When aux predicts
@@ -38,45 +60,40 @@
// apply; conventional kernel discipline (no atomicAdd, no host
// branches inside the kernel body) preserved anyway.
//
// Launch: grid=(1, 1, 1), block=(4, 1, 1). Each of the 4 threads
// writes one slot.
// Launch: grid=(1, 1, 1), block=(12, 1, 1). Each thread writes one slot.
#include <cuda_runtime.h>
extern "C" __global__ void aux_w_prior_init(
/* [b0_size=4] f32 output buffer; overwrites alloc_zeros'd content
* with the structural prior values. */
/* [b0_size=4, K=3] f32 output buffer (12 elements row-major); overwrites
* alloc_zeros'd content with the structural prior values. */
float* __restrict__ w_aux_to_q_dir
) {
int a = threadIdx.x;
if (a >= 4) return;
int idx = threadIdx.x;
if (idx >= 12) return;
/* SP22 H6 Phase 3 α — DORMANT (hypothesis falsified 2026-05-13):
/* SP22 H6 vNext Phase D (2026-05-14) — K=3 outcome-conditioned
* structural prior. Indexed as `w[a * K + k]` where a is the
* direction action (Short=0, Hold=1, Long=2, Flat=3) and k is the
* outcome class (Profit=0, Stop=1, Timeout=2).
*
* Decisive smoke `train-xrkb7` @ ebc7144434 with Phase 3 active +
* enlarged aux head (AUX_HIDDEN_DIM=128) produced WR=0.4346 —
* statistically identical to all baselines (0.4338-0.4346).
* The prior is DORMANT-by-design — values land in W at trainer
* construction but the atom-shift Q-target modulation is gated by
* the K=3 head's softmax outputs in state[121..124). At cold-start
* those slots are 0.0 (no Profit/Stop/Timeout prediction yet) so
* the atom-shift is effectively zero even with W populated. After
* the K=3 head trains and emits real probs, the prior's structural
* bias kicks in; Adam refines W from there.
*
* Root cause of falsification: aux head's "70% accuracy" in fold 0
* was mostly majority-class prediction on imbalanced labels (88% down).
* Fold 1 with more balanced labels showed accuracy drop to 52% —
* aux has minimal per-bar discriminative power. The H6 hypothesis
* (use aux's per-bar directional prediction to bias the policy)
* cannot improve WR at this data regime + horizon.
*
* **Architectural insight (next-session work)**: the aux head is
* predicting THE WRONG THING. The system makes MULTI-BAR TRADE
* decisions with target_bars / profit_target / stop_loss / conviction
* plan parameters, but the aux head predicts next-bar close direction.
* The decision horizons mismatch — see
* `docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md` for the
* proposed trade-outcome aux head that aligns with the actual
* decision the policy makes.
*
* Infrastructure (atom-shift kernels, NaN guards, Step 8/11 backward
* + Adam, Phase C1 collector wireup, enlarged aux trunk/head) remains
* intact. vNext re-uses all of it with the new aux signal source.
* Adam refines W from this prior via the per-(a, k) gradient
* computed in c51_aux_dw_kernel + applied in adam_w_aux_kernel
* (both extended to 12 weights in this phase).
*/
float prior_values[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
w_aux_to_q_dir[a] = prior_values[a];
float prior_values[12] = {
/* Short × {Profit, Stop, Timeout} */ +0.5f, -0.5f, 0.0f,
/* Hold × {Profit, Stop, Timeout} */ 0.0f, +0.5f, 0.0f,
/* Long × {Profit, Stop, Timeout} */ +0.5f, -0.5f, 0.0f,
/* Flat × {Profit, Stop, Timeout} */ 0.0f, +0.5f, 0.0f,
};
w_aux_to_q_dir[idx] = prior_values[idx];
}

92
crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
View File

@@ -1,12 +1,21 @@
// 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].
// SP22 H6 Phase 3 α Step 8 → SP22 H6 vNext Phase D Step 8 — backward dW
// kernel for w_aux_to_q_dir [b0_size=4, K=3 outcomes] = 12 weights.
//
// 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
// Computes gradient of C51 loss w.r.t. the per-(action, outcome)
// atom-shift weights:
// dW[a*K + k] =
// Σ_b indicator(a_d_b == a) × isw_b × inv_batch × (SP_b / Δz_b)
// × state[121 + k]_b
// + Σ_b indicator(a*_b == a) × isw_b × inv_batch × (-γ_eff_b × (1-done_b))
// × (SP_b / Δz_b) × next_state[121 + k]_b
//
// Phase D shift from K=2 / Phase 3:
// - Phase 3 used state[121] only (scalar; K=2 head's recentered p_up)
// - Phase D uses state[121+k] for k ∈ {0=Profit, 1=Stop, 2=Timeout}
// (K=3 head's softmax probs from state slots [121..124))
// - W is indexed (a, k) instead of just (a), so 12 weights total.
//
// where:
// - SP_b = projection log-diff sum (saved by c51_loss_kernel forward):
@@ -20,12 +29,14 @@
// - 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).
// Output: dw_aux_buf[b0_size × K = 12] 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 pearl_no_atomicadd: ONE BLOCK PER (action, outcome) PAIR
// (grid=(12,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
@@ -55,7 +66,7 @@ extern "C" __global__ void c51_aux_dw_kernel(
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 */
float* __restrict__ dw_aux, /* [b0_size × K = 12] gradient — written by tid==0 per block */
/* Config */
int batch_size,
@@ -63,13 +74,29 @@ extern "C" __global__ void c51_aux_dw_kernel(
int b1_size,
int b2_size,
int b3_size,
int aux_dir_prob_index, /* = 121 */
int state_dim /* = 128 */
int aux_dir_prob_index, /* = 121 (= AUX_OUTCOME_PROFIT_INDEX) */
int state_dim, /* = 128 */
/* SP22 H6 vNext Phase D (2026-05-14): K=3 outcome dimension (3 in
* production = AUX_OUTCOME_K). Each block computes one (a, k) pair;
* grid = (b0_size × K, 1, 1). Decoded as `a = blockIdx.x / K`,
* `k = blockIdx.x % K`. The state slot is `aux_dir_prob_index + k`
* (slots 121, 122, 123 for k = 0, 1, 2). */
int aux_outcome_k
) {
int a = blockIdx.x; /* which W index this block accumulates */
/* SP22 H6 vNext Phase D (2026-05-14): decode (a, k) from blockIdx.x.
* Layout: grid block index = a * K + k. */
int block_idx = blockIdx.x;
int a = block_idx / aux_outcome_k;
int k = block_idx - a * aux_outcome_k;
int tid = threadIdx.x;
if (a >= b0_size) return;
if (a >= MAX_DIR_ACTIONS) return;
if (k >= aux_outcome_k) return;
/* State slot for this outcome class (k=0 → slot 121 = p_Profit,
* k=1 → slot 122 = p_Stop, k=2 → slot 123 = p_Timeout). */
int state_slot = aux_dir_prob_index + k;
if (state_slot < 0 || state_slot >= state_dim) return;
float inv_batch = 1.0f / (float)batch_size;
@@ -93,12 +120,7 @@ extern "C" __global__ void c51_aux_dw_kernel(
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];
/* DEFENSIVE NaN guards (2026-05-13): the dz check `dz < 1e-7` does
* NOT catch NaN (NaN comparison is always false), so NaN dz would
* slip through and produce NaN dW. Bisect cut #1 proved this
* kernel + Adam are the NaN source. Guard ALL upstream-derived
* inputs: sp / is_weights / dz / gamma / done. Any NaN/Inf in
* any input → skip sample (no contribution to dW). */
/* DEFENSIVE NaN guards: any NaN/Inf in inputs → skip sample. */
if (!isfinite(sp) || !isfinite(isw) || !isfinite(dz) || dz < 1e-7f) continue;
float gamma_eff = gamma_buf[b];
float done = dones[b];
@@ -108,19 +130,16 @@ extern "C" __global__ void c51_aux_dw_kernel(
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];
/* Defensive NaN guard (2026-05-13): mirrors atom-shift kernels.
* State_121 NaN would make local_sum NaN, dW NaN, Adam NaN,
* w_aux NaN — corrupting all subsequent forward atom-shifts.
* Clamp NaN to 0 → no W gradient contribution from that sample
* (correct behavior; the gradient was undefined anyway). */
if (!isfinite(s121)) s121 = 0.0f;
local_sum += inv_batch * isw * dL_dDelta_online * s121;
/* Phase D: state slot = aux_dir_prob_index + k (= 121, 122,
* or 123 for k = 0, 1, 2 outcome classes). */
float s_k = batch_states[(long long)b * state_dim + state_slot];
if (!isfinite(s_k)) s_k = 0.0f;
local_sum += inv_batch * isw * dL_dDelta_online * s_k;
}
if (a == a_star) {
float ns121 = next_batch_states[(long long)b * state_dim + aux_dir_prob_index];
if (!isfinite(ns121)) ns121 = 0.0f;
local_sum += inv_batch * isw * dL_dDelta_target * ns121;
float ns_k = next_batch_states[(long long)b * state_dim + state_slot];
if (!isfinite(ns_k)) ns_k = 0.0f;
local_sum += inv_batch * isw * dL_dDelta_target * ns_k;
}
}
@@ -141,15 +160,10 @@ extern "C" __global__ void c51_aux_dw_kernel(
if (tid == 0) {
float total = 0.0f;
for (int w = 0; w < NUM_WARPS; w++) total += shmem_warp[w];
/* SP22 H6 Phase 3 α Step 8 DEFENSIVE NaN guard (2026-05-13):
* Bisect cut #1 proved the NaN source is in Step 8 (this kernel)
* or Step 11 (adam_w_aux). With forward atom-shift kernels
* disabled outputs are clean; with Step 8+11 enabled NaN appears.
* If `total` is NaN/Inf (e.g., from sp/dz overflow or NaN
* inputs through scratch buffers), set dW to 0 instead of
* propagating NaN into W via Adam. This prevents W → NaN →
* atom-shift = NaN × sanitized_state = NaN cascade. */
/* Defensive NaN guard — set dW to 0 instead of propagating
* NaN into W via Adam. Phase D: write to dw_aux[a * K + k]
* instead of dw_aux[a]. */
if (!isfinite(total)) total = 0.0f;
dw_aux[a] = total;
dw_aux[a * aux_outcome_k + k] = total;
}
}

66
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -828,23 +828,27 @@ extern "C" __global__ void c51_loss_batched(
* Read once per sample (uniform across threads in the block; could
* be hoisted to a single-thread broadcast but per-thread reads from
* mapped pinned / device memory are cached identically). */
float state_121 = 0.0f;
float next_state_121 = 0.0f;
/* SP22 H6 vNext Phase D (2026-05-14): hoist 3 outcome-class state
* slots per sample (was: single state_121 + next_state_121 scalars).
* Slots: state[121] = p_Profit, state[122] = p_Stop, state[123] =
* p_Timeout. Defensive NaN guards per-element. */
float state_outcome[3] = { 0.0f, 0.0f, 0.0f };
float next_state_outcome[3] = { 0.0f, 0.0f, 0.0f };
bool aux_shift_active = (w_aux != NULL
&& batch_states != NULL
&& next_batch_states != NULL);
if (aux_shift_active) {
state_121 = batch_states[(long long)sample_id * state_dim
+ aux_dir_prob_index];
next_state_121 = next_batch_states[(long long)sample_id * state_dim
+ aux_dir_prob_index];
/* Defensive NaN guard (2026-05-13): state_121 SHOULD be in
* [-1, +1] per the H6 spec, but IEEE-754 `0 * NaN = NaN`
* defeats W=0 safety if an upstream NaN propagates here. The
* real source (aux head or state assembly) is investigated
* separately. */
if (!isfinite(state_121)) state_121 = 0.0f;
if (!isfinite(next_state_121)) next_state_121 = 0.0f;
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float s_k = batch_states[(long long)sample_id * state_dim
+ aux_dir_prob_index + kk];
float ns_k = next_batch_states[(long long)sample_id * state_dim
+ aux_dir_prob_index + kk];
if (!isfinite(s_k)) s_k = 0.0f;
if (!isfinite(ns_k)) ns_k = 0.0f;
state_outcome[kk] = s_k;
next_state_outcome[kk] = ns_k;
}
}
for (int d = 0; d < NUM_BRANCHES; d++) {
@@ -1088,11 +1092,15 @@ extern "C" __global__ void c51_loss_batched(
* the same amount (linearity of conditional expectation).
*/
if (d == 0 && aux_shift_active) {
/* DEFENSIVE NaN guard (2026-05-13): W can carry NaN from
* Adam corruption; 0 * NaN = NaN propagates. */
float w_a = w_aux[a];
if (!isfinite(w_a)) w_a = 0.0f;
cvar += w_a * next_state_121;
/* SP22 H6 vNext Phase D (2026-05-14): K=3 atom-shift
* sum on the target side (next_state).
* cvar += Σ_k W[a*K + k] × next_state[121+k] */
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float w_ak = w_aux[a * 3 + kk];
if (!isfinite(w_ak)) w_ak = 0.0f;
cvar += w_ak * next_state_outcome[kk];
}
}
eq_per_action[a] = cvar;
}
@@ -1237,11 +1245,23 @@ extern "C" __global__ void c51_loss_batched(
* branches and when aux_shift_active == false. */
float effective_reward = reward;
if (d == 0 && aux_shift_active) {
/* DEFENSIVE NaN guard (2026-05-13): W can carry NaN. */
float w_ad = w_aux[a0]; if (!isfinite(w_ad)) w_ad = 0.0f;
float w_star = w_aux[best_next_a]; if (!isfinite(w_star)) w_star = 0.0f;
float aux_shift_online = w_ad * state_121;
float aux_shift_target = w_star * next_state_121;
/* SP22 H6 vNext Phase D (2026-05-14): K=3 atom-shift sums
* for both the online (taken action a0 → current state) and
* the target (sampled action best_next_a → next state):
* aux_shift_online = Σ_k W[a0*K + k] × state[121+k]
* aux_shift_target = Σ_k W[a*_K + k] × next_state[121+k]
*/
float aux_shift_online = 0.0f;
float aux_shift_target = 0.0f;
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float w_ad = w_aux[a0 * 3 + kk];
float w_star = w_aux[best_next_a * 3 + kk];
if (!isfinite(w_ad)) w_ad = 0.0f;
if (!isfinite(w_star)) w_star = 0.0f;
aux_shift_online += w_ad * state_outcome[kk];
aux_shift_target += w_star * next_state_outcome[kk];
}
effective_reward = reward
+ gamma_eff * aux_shift_target * (1.0f - done)
- aux_shift_online;

92
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -5194,16 +5194,24 @@ extern "C" __global__ void compute_expected_q(
* at zero (only direction branch is α-influenced). */
float aux_atom_shift = 0.0f;
if (d == 0 && w_aux != NULL && batch_states != NULL) {
float state_121 = batch_states[(long long)i * state_dim
+ aux_dir_prob_index];
float w_a = w_aux[a];
/* SP22 H6 Phase 3 α DEFENSIVE GUARDS (2026-05-13):
* Both state_121 and W can carry upstream NaN. IEEE-754
* `0 * NaN = NaN` propagates through atom-shift regardless
* of which factor is zero. Sanitize both before multiply. */
if (!isfinite(state_121)) state_121 = 0.0f;
if (!isfinite(w_a)) w_a = 0.0f;
aux_atom_shift = w_a * state_121;
/* SP22 H6 vNext Phase D (2026-05-14): K=3 outcome-
* conditioned atom-shift. Sum over outcome classes
* (Profit/Stop/Timeout):
* atom_shift = Σ_k w_aux[a*K + k] × state[121 + k]
*
* The K=3 head's softmax probs land in state slots
* [121..124) via Phase C-2's state gather. W is sized
* [b0_size × K] = 12 floats. NaN guards per-element
* because state slots can carry upstream NaN. */
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float s_k = batch_states[(long long)i * state_dim
+ aux_dir_prob_index + kk];
float w_ak = w_aux[a * 3 + kk];
if (!isfinite(s_k)) s_k = 0.0f;
if (!isfinite(w_ak)) w_ak = 0.0f;
aux_atom_shift += w_ak * s_k;
}
}
/* ── Online softmax: 2-pass instead of 3-pass ────────────────
@@ -5723,15 +5731,22 @@ extern "C" __global__ void mag_concat_qdir(
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
/* SP22 H6 Phase 3 α atom-shift: hoist state_121 once per sample.
* NULL-safe per the same pattern as compute_expected_q + c51_loss.
* Defensive NaN guard (2026-05-13): see compute_expected_q for
* rationale — IEEE-754 `0 * NaN = NaN` defeats W=0 safety. */
float state_121 = 0.0f;
/* SP22 H6 vNext Phase D (2026-05-14): hoist 3 outcome-class state
* slots per sample (was: single state_121 scalar). Slots:
* state[121] = p_Profit, state[122] = p_Stop, state[123] = p_Timeout
* NULL-safe — when w_aux or batch_states is NULL, the atom-shift
* computation below short-circuits to zero. Defensive NaN guards
* per pattern in compute_expected_q. */
float state_outcome[3] = { 0.0f, 0.0f, 0.0f };
bool aux_shift_active = (w_aux != NULL && batch_states != NULL);
if (aux_shift_active) {
state_121 = batch_states[(long long)b * state_dim + aux_dir_prob_index];
if (!isfinite(state_121)) state_121 = 0.0f;
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float s_k = batch_states[(long long)b * state_dim
+ aux_dir_prob_index + kk];
if (!isfinite(s_k)) s_k = 0.0f;
state_outcome[kk] = s_k;
}
}
/* Plan 4 Task 2c.3c.6: read trunk-output RMS once into a register.
@@ -5809,17 +5824,16 @@ extern "C" __global__ void mag_concat_qdir(
float z_val = v_min + (float)z * dz;
eq += prob * z_val;
}
/* SP22 H6 Phase 3 α atom-shift: per-action shift = W[a]*state_121.
* Effective atom positions for this action: z_n + W[a]*state_121.
* Since Σ_n prob_n = 1, the expected Q shifts by exactly the
* scalar W[a]*state_121 (linearity of expectation).
* DEFENSIVE NaN guard (2026-05-13): same pattern as
* compute_expected_q — W can carry NaN from Adam corruption, and
* 0 * NaN = NaN propagates. Verify finite. */
/* SP22 H6 vNext Phase D (2026-05-14): K=3 atom-shift sum.
* eq += Σ_k W[a*K + k] × state[121+k]
* (linearity of expectation: Σ_n prob_n = 1 → shift is scalar). */
if (aux_shift_active) {
float w_a = w_aux[a];
if (!isfinite(w_a)) w_a = 0.0f;
eq += w_a * state_121;
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float w_ak = w_aux[a * 3 + kk];
if (!isfinite(w_ak)) w_ak = 0.0f;
eq += w_ak * state_outcome[kk];
}
}
eq_local[a] = eq;
}
@@ -6465,16 +6479,22 @@ extern "C" __global__ void quantile_q_select(
float quantile_blend = (1.0f - alpha) * q90 + alpha * q10;
float q_blended = util_factor * quantile_blend + (1.0f - util_factor) * q50;
/* SP22 H6 Phase 3 α atom-shift: dir branch (d==0) only.
* Per-action shift = W[a]*state_121 uniformly shifts q10/q50/q90
* (they're all positions on the shifted support), so the blend
* shifts by exactly W[a]*state_121.
* DEFENSIVE NaN guard (2026-05-13): same pattern as
* compute_expected_q + mag_concat_qdir. */
/* SP22 H6 vNext Phase D (2026-05-14): K=3 outcome-conditioned
* atom-shift sum.
* q_blended += Σ_k W[a*K + k] × state[121+k]
* The shift is scalar (linearity of expectation on the
* shifted support), so it applies uniformly to q10/q50/q90
* and the blend shifts by exactly Σ_k W[a*K+k] × state[121+k]. */
if (d == 0 && aux_shift_active) {
float w_a = w_aux[a];
if (!isfinite(w_a)) w_a = 0.0f;
q_blended += w_a * state_121;
#pragma unroll 3
for (int kk = 0; kk < 3; kk++) {
float s_k = batch_states[(long long)i * state_dim
+ aux_dir_prob_index + kk];
float w_ak = w_aux[a * 3 + kk];
if (!isfinite(s_k)) s_k = 0.0f;
if (!isfinite(w_ak)) w_ak = 0.0f;
q_blended += w_ak * s_k;
}
}
q_select[(long long)i * total_actions + action_idx] = q_blended;

59
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -21225,31 +21225,39 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("sp13 aux_pred_to_isv_tanh_kernel load: {e}")))?
};
// SP22 H6 Phase 3 α (2026-05-13): post-encoder Q_dir bias.
// Allocate the 4-element weight + 4-element Adam moments +
// 4-element gradient accumulator. Load forward + backward
// kernels. All buffers zero-init per `alloc_zeros`; Adam
// moves W from cold-start zero based on the backward signal.
// SP22 H6 Phase 3 α (2026-05-13) → SP22 H6 vNext Phase D
// (2026-05-14): post-encoder Q_dir bias W matrix.
//
// Phase D extends the Phase 3 W shape from [b0_size=4] (per-action
// scalar weight) to [b0_size=4, K=3=AUX_OUTCOME_K] (per-action ×
// per-outcome). Same Adam machinery (m, v moments + dW
// accumulator) — just 3× the elements. The W shape change is
// coordinated with the atom-shift kernel reads (compute_expected_q,
// c51_loss, mag_concat_qdir, quantile_q_select), the dW kernel
// (emits 12 partials), the Adam kernel (updates 12 weights), and
// the prior init kernel (writes 12 structural priors).
let b0_size_usize: usize = 4; // factored space b0_size = 4 (direction)
let aux_kto_init = AUX_OUTCOME_K; // = 3
let w_aux_count: usize = b0_size_usize * aux_kto_init; // 12
let w_aux_to_q_dir = stream
.alloc_zeros::<f32>(b0_size_usize)
.alloc_zeros::<f32>(w_aux_count)
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: alloc w_aux_to_q_dir: {e}"
"sp22-vnext D: alloc w_aux_to_q_dir [4, 3]=12: {e}"
)))?;
let adam_m_w_aux = stream
.alloc_zeros::<f32>(b0_size_usize)
.alloc_zeros::<f32>(w_aux_count)
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: alloc adam_m_w_aux: {e}"
"sp22-vnext D: alloc adam_m_w_aux [12]: {e}"
)))?;
let adam_v_w_aux = stream
.alloc_zeros::<f32>(b0_size_usize)
.alloc_zeros::<f32>(w_aux_count)
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: alloc adam_v_w_aux: {e}"
"sp22-vnext D: alloc adam_v_w_aux [12]: {e}"
)))?;
let dw_aux_buf = stream
.alloc_zeros::<f32>(b0_size_usize)
.alloc_zeros::<f32>(w_aux_count)
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: alloc dw_aux_buf: {e}"
"sp22-vnext D: alloc dw_aux_buf [12]: {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
@@ -21313,11 +21321,14 @@ impl GpuDqnTrainer {
.arg(&w_ptr)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (4, 1, 1),
// SP22 H6 vNext Phase D (2026-05-14): bumped 4 → 12
// threads. Kernel writes 12 values (4 actions × 3
// outcomes) per the K=3 structural prior table.
block_dim: (12, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α: aux_w_prior_init launch: {e}"
"sp22-vnext D: aux_w_prior_init launch: {e}"
)))?;
}
}
@@ -30657,6 +30668,11 @@ impl GpuDqnTrainer {
let next_states_ptr = self.ptrs.next_states_buf;
let dw_aux_ptr = self.dw_aux_buf.raw_ptr();
// SP22 H6 vNext Phase D (2026-05-14): bumped grid from 4 to
// b0_size × K = 12. The kernel decodes blockIdx.x → (a, k);
// each block tree-reduces dW[a*K + k] over the batch dim.
let aux_kto_i32 = AUX_OUTCOME_K as i32;
let dw_grid_blocks: u32 = (b0 as u32) * (aux_kto_i32 as u32);
unsafe {
self.stream
.launch_builder(&self.c51_aux_dw_kernel)
@@ -30677,13 +30693,14 @@ impl GpuDqnTrainer {
.arg(&b3)
.arg(&aux_dir_prob_index)
.arg(&state_dim_i32)
.arg(&aux_kto_i32) // Phase D: K=3 outcome dimension
.launch(LaunchConfig {
grid_dim: (4, 1, 1), // one block per W index (b0_size=4)
block_dim: (256, 1, 1), // tree-reduce across batch
grid_dim: (dw_grid_blocks, 1, 1), // Phase D: b0_size × K = 12 blocks
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}"
"sp22-vnext D Step 8: c51_aux_dw_kernel launch: {e}"
)))?;
}
Ok(())
@@ -30725,11 +30742,13 @@ impl GpuDqnTrainer {
.arg(&t_ptr)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (4, 1, 1), // one thread per W slot
// SP22 H6 vNext Phase D (2026-05-14): bumped 4 → 12
// threads. 12 weights = b0_size × K = 4 × 3.
block_dim: (12, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp22-h6-phase3 α Step 11: adam_w_aux_kernel launch: {e}"
"sp22-vnext D Step 11: adam_w_aux_kernel launch: {e}"
)))?;
}
Ok(())