diff --git a/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu index f6a98914c..e0d10d4b9 100644 --- a/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu +++ b/crates/ml/src/cuda_pipeline/adam_w_aux_kernel.cu @@ -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 -#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, diff --git a/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu b/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu index 621ba7a77..fc1dced47 100644 --- a/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu +++ b/crates/ml/src/cuda_pipeline/aux_w_prior_init_kernel.cu @@ -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 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]; } diff --git a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu index f052faa14..5fcdfa0fe 100644 --- a/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu @@ -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; } } diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 6e9ee0452..5946883c6 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 476da9faa..17ac3e371 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 964483c0c..ca64943af 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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::(b0_size_usize) + .alloc_zeros::(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::(b0_size_usize) + .alloc_zeros::(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::(b0_size_usize) + .alloc_zeros::(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::(b0_size_usize) + .alloc_zeros::(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(()) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index fac9aef62..2e6318bf7 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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).