From 7fc9799343a27ae0262ff6905113187d7797caf5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 12 May 2026 21:31:32 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22):=20H6=20Phase=201=20=E2=80=94=20aux?= =?UTF-8?q?=E2=86=92policy=20state=20bridge=20(atomic)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the aux head's per-env directional probability into policy STATE slot 121 (AUX_DIR_PROB_INDEX = PADDING_START + 0), preserving the trunk- separation invariant from `pearl_separate_aux_trunk_when_shared_starves`. H1 (label horizon) confirmed aux learns 78% dir-acc within-fold at H=200 but the policy was walled off; this commit conducts that signal through the state input with a one-step lag. Mechanism (rollout-time, collector-only) ──────────────────────────────────────── - State[env, t] reads `prev_aux_dir_prob[env]` (= p_up from step t-1). - After aux forward at step t, the new copy kernel writes `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` for step t+1. - Cold-start + FoldReset seed the buffer to 0.5 (neutral; p_up = 50%) via the pure-GPU `fill_f32` kernel — no HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md`. - Launch sits in the same `isv_signals && trainer_params != 0` gate as the aux forward, so when aux is skipped the cache keeps its previous (sentinel or last-good) value instead of copying stale `alloc_zeros`. Three state-gather kernels updated atomically (per `feedback_no_partial_refactor.md`): - `experience_state_gather` — training, reads `aux_dir_prob_per_env[i]` - `backtest_state_gather` — eval (single-step), NULL → 0.5 (A3 fallback) - `backtest_state_gather_chunk` — eval (chunked), NULL → 0.5 (A3 fallback) `assemble_state` gained a 7th param `float aux_dir_prob` written to `out[SL_PADDING_START + 0]`; the remaining 6 padding slots stay zero for 8-alignment. Phase 1 scope = training-side + eval A3 NULL fallback. A2 (aux trunk forward in eval) is deferred per the runbook — gates on whether the smoke moves WR off the 50.1–50.2% plateau. New files ───────── - crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu Modified ──────── - crates/ml-core/src/state_layout.rs (+AUX_DIR_PROB_INDEX) - crates/ml/src/cuda_pipeline/state_layout.cuh (+assemble_state param) - crates/ml/src/cuda_pipeline/experience_kernels.cu (3 state-gather kernels + NULL-defensive sentinel) - crates/ml/src/cuda_pipeline/gpu_experience_collector.rs (per-env buffer + 2 kernel handles + cold-start fill + copy launch + FoldReset re-fill + state-gather arg) - crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs (NULL aux_dir_prob_per_env at all 3 launchers for A3) - crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN static) - crates/ml/src/cuda_pipeline/gpu_action_selector.rs (EPSILON_GREEDY_CUBIN → pub(crate) so collector reuses fill_f32) - crates/ml/build.rs (register new kernel) - docs/dqn-wire-up-audit.md (## 2026-05-12 — SP22 H6 implementation: Phase 1 entry) Verification (all three gates clean, no smoke yet) ────────────────────────────────────────────────── - cargo check -p ml --features cuda: 0 errors - gpu_backtest_validation: 4/4 expected-passing tests still pass (the 2 PnL-assertion failures are pre-existing per the runbook) - compute-sanitizer --tool=memcheck: 0 CUDA errors Refs ──── - docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md - docs/plans/2026-05-12-sp22-h6-next-session-prompt.md - pearl_separate_aux_trunk_when_shared_starves - pearl_first_observation_bootstrap (sentinel = 0.5 cold-start) - feedback_no_htod_htoh_only_mapped_pinned (fill_f32 not HtoD) - feedback_no_partial_refactor (3 state-gather kernels atomic) Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-core/src/state_layout.rs | 7 + crates/ml/build.rs | 13 ++ .../aux_softmax_to_per_env_kernel.cu | 79 +++++++++ .../src/cuda_pipeline/experience_kernels.cu | 45 +++++- .../src/cuda_pipeline/gpu_action_selector.rs | 6 +- .../cuda_pipeline/gpu_backtest_evaluator.rs | 21 ++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 12 ++ .../cuda_pipeline/gpu_experience_collector.rs | 152 ++++++++++++++++++ crates/ml/src/cuda_pipeline/state_layout.cuh | 11 +- docs/dqn-wire-up-audit.md | 120 ++++++++++++++ 10 files changed, 453 insertions(+), 13 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu diff --git a/crates/ml-core/src/state_layout.rs b/crates/ml-core/src/state_layout.rs index e0d1bbe98..2af6ef5ee 100644 --- a/crates/ml-core/src/state_layout.rs +++ b/crates/ml-core/src/state_layout.rs @@ -55,6 +55,13 @@ pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 106 pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 114 pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 121 +// SP22 H6 (2026-05-12) — aux directional probability from previous step. +// Lives in the first padding slot. PADDING_DIM stays 7 (no STATE_DIM change); +// `assemble_state` writes aux_dir_prob to [AUX_DIR_PROB_INDEX] and zeros the +// remaining 6 padding slots [PADDING_START+1..STATE_DIM). +// Sentinel = 0.5 (neutral; p_up = 50%) at cold-start and FoldReset. +pub const AUX_DIR_PROB_INDEX: usize = PADDING_START; // 121 + const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM"); const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores"); const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS"); diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 67d0c990f..4eee7ac42 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -792,6 +792,19 @@ fn main() { // producer chain has fresh labels every rollout step. Same per-thread // O(1) map; same lookahead window semantics; same skip sentinel. "aux_sign_label_per_step_kernel.cu", + // SP22 H6 (2026-05-12): per-env p_up extractor that copies + // `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` after each aux + // forward. The cache buffer is read by `experience_state_gather` + // and `backtest_state_gather*` at the next step's state assembly, + // wiring the aux head's directional probability into policy state + // (slot `SL_PADDING_START + 0 = 121`). One-step lag is intentional: + // state is assembled BEFORE aux forward at step t, so step t+1 + // reads step t's aux output. Pure per-env gather; no reduction, no + // atomicAdd per `feedback_no_atomicadd.md`; pure GPU compute per + // `feedback_no_cpu_compute_strict.md`. Launched on the collector's + // stream right after `aux_next_bar_forward` in + // `gpu_experience_collector.rs`. + "aux_softmax_to_per_env_kernel.cu", // SP14 Layer C Phase C.4b (2026-05-08): adaptive aux prediction // horizon producer. Single-thread kernel writing // `ISV[AUX_PRED_HORIZON_BARS_INDEX=450]` from diff --git a/crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu b/crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu new file mode 100644 index 000000000..9be0527ae --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu @@ -0,0 +1,79 @@ +// crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu +// +// SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up extractor. +// +// Role +// ──── +// At rollout step t, after the aux head's K=2 softmax forward produces the +// per-env tile `aux_softmax [n_envs, K=2]`, this kernel copies the second +// column (p_up = softmax[1]) into the per-env cache buffer `prev_aux_dir_prob +// [n_envs]`. The cache is read by `experience_state_gather` at step t+1 to +// populate `state[AUX_DIR_PROB_INDEX = SL_PADDING_START + 0 = 121]`. +// +// Why a separate copy kernel (not direct softmax → state) +// ─────────────────────────────────────────────────────── +// Step t's state was already assembled BEFORE step t's aux forward ran (state +// is the network input, aux output is a forward result). So the state at step +// t+1 reads the aux output from step t — a one-step lag. The cache buffer is +// the lag carrier. H=200 label horizon makes the aux signal slow-moving, so +// one-bar lag is fine (per `pearl_first_observation_bootstrap`'s tolerance for +// slow signals on shared scalars). +// +// Why one slot per env (not batch mean) +// ───────────────────────────────────── +// Unlike `aux_pred_to_isv_tanh_kernel` (which reduces the batch to a single +// shared scalar in ISV[375]), this kernel preserves per-env granularity +// because each env's state row needs its own aux signal. The aux softmax is +// already per-env (`[n_envs, K]`), so this is just a strided gather into +// `[n_envs]`. +// +// Sentinel / cold-start / FoldReset +// ───────────────────────────────── +// The buffer is filled with 0.5 (neutral; p_up = 50%) at construction and at +// every FoldReset by the GPU `fill_f32` kernel (epsilon_greedy_kernel.cu). +// This kernel ALWAYS overwrites — no sentinel branch needed here. The first +// rollout step's state assembly reads 0.5 (no aux signal yet), then this +// kernel writes the real value for step 2's state, and so on. +// +// Discipline +// ────────── +// - No atomicAdd per `feedback_no_atomicadd.md` (no reduction, just gather). +// - No HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md` (GPU↔GPU only). +// - Stream-ordered with the producer that wrote the softmax tile; same-stream +// barrier suffices (cf. `aux_pred_to_isv_tanh_kernel`'s same-stream contract +// with `aux_next_bar_loss_reduce`). +// +// Launch +// ────── +// grid = ((n_envs + 255) / 256), block = 256. K is passed runtime so the +// kernel ABI stays stable if the aux head's K flips (currently 2). +// +// Cost +// ──── +// One coalesced 4-byte load + one 4-byte store per env. n_envs is O(1k); ~1µs. + +#include + +extern "C" __global__ void aux_softmax_to_per_env_kernel( + /* Per-env K-way softmax tile [n_envs, K]. Written by + * `aux_next_bar_forward` during the captured rollout-step graph; valid + * after same-stream sync. Layout row-major; this kernel reads + * `aux_softmax[env * K + 1]` per env. */ + const float* __restrict__ aux_softmax, + /* Per-env cache buffer [n_envs]. Written here, read by + * `experience_state_gather` and `backtest_state_gather*` at the next + * step's state assembly. */ + float* __restrict__ prev_aux_dir_prob, + /* Number of envs (rollout-time batch width = `alloc_episodes`). */ + int n_envs, + /* Softmax K (= AUX_NEXT_BAR_K = 2 in B1.1a). Kernel reads slot K=1 + * (p_up). Passed runtime so the kernel signature stays stable across + * head-dim flips — the only index read is `env * K + 1`. */ + int K +) { + int env = blockIdx.x * blockDim.x + threadIdx.x; + if (env >= n_envs) return; + + /* Write p_up = softmax[1]. Used by next step's state assembly. */ + prev_aux_dir_prob[env] = aux_softmax[(size_t)env * (size_t)K + 1]; +} diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 14a4cc0c4..c8545be38 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -651,7 +651,13 @@ extern "C" __global__ void experience_state_gather( const float* __restrict__ plan_params_ptr, /* [N, 6] trade plan params. NULL = no plan. */ const float* __restrict__ isv_signals_ptr, /* [12] pinned device-mapped ISV signals. NULL = static. */ const float* __restrict__ ofi_features, /* [total_bars, ofi_dim] OFI features. NULL = no OFI. */ - int ofi_dim /* OFI feature dimension (20 with deltas+book+dur). 0 = no OFI. */ + int ofi_dim, /* OFI feature dimension (20 with deltas+book+dur). 0 = no OFI. */ + /* SP22 H6 (2026-05-12): per-env aux directional probability from previous step. + * Length [N]; populated by `aux_softmax_to_per_env_kernel` after each aux forward. + * NULL = no aux signal yet (cold-start or producer not wired) → kernel uses + * 0.5 sentinel (neutral). See `state_layout.cuh::assemble_state` slot + * `SL_PADDING_START+0`. */ + const float* __restrict__ aux_dir_prob_per_env ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; @@ -968,8 +974,16 @@ extern "C" __global__ void experience_state_gather( ofi[k] = ofi_row[k]; } + /* ── SP22 H6 aux→state bridge: previous-step p_up for this env ── + * NULL = no aux producer wired yet (cold-start, eval A3 fallback) → 0.5 + * sentinel (neutral, no signal). Lag of 1 bar is intentional — H=200 + * label horizon makes the signal slow-moving, so one-step lag is fine. */ + float aux_dir_prob = (aux_dir_prob_per_env != NULL) + ? aux_dir_prob_per_env[i] + : 0.5f; + /* ── Final assembly: canonical layout via state_layout.cuh ── */ - assemble_state(out, market, ofi, mtf, portfolio, plan_isv); + assemble_state(out, market, ofi, mtf, portfolio, plan_isv, aux_dir_prob); } /* ================================================================== */ @@ -8633,7 +8647,11 @@ extern "C" __global__ void backtest_state_gather( int max_len, int feat_dim, /* = market_dim + ofi_dim = 62 */ int current_step, - int padded_sd + int padded_sd, + /* SP22 H6 (2026-05-12): per-window aux directional probability from + * previous step. Length [n_windows]. NULL = eval-time A3 fallback + * (Phase 1 — eval has no aux infrastructure yet) → 0.5 sentinel. */ + const float* __restrict__ aux_dir_prob_per_env ) { int w = blockIdx.x * blockDim.x + threadIdx.x; if (w >= n_windows) return; @@ -8709,8 +8727,13 @@ extern "C" __global__ void backtest_state_gather( for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++) pisv[k] = plan_isv[w * SL_PORTFOLIO_PLAN_DIM + k]; + /* SP22 H6 — per-window aux p_up (NULL → 0.5 sentinel for Phase 1 A3 eval). */ + float aux_dir_prob = (aux_dir_prob_per_env != NULL) + ? aux_dir_prob_per_env[w] + : 0.5f; + /* Single shared assembly call — produces identical layout to training */ - assemble_state(out, market, ofi, mtf, pf, pisv); + assemble_state(out, market, ofi, mtf, pf, pisv, aux_dir_prob); /* Zero-pad [SL_STATE_DIM .. padded_sd) for cuBLAS K-tile alignment */ for (int k = SL_STATE_DIM; k < padded_sd; k++) @@ -8752,7 +8775,12 @@ extern "C" __global__ void backtest_state_gather_chunk( int feat_dim, /* = market_dim + ofi_dim */ int start_step, int chunk_len, - int padded_sd + int padded_sd, + /* SP22 H6 (2026-05-12): per-window aux directional probability. + * Constant within a chunk (matches portfolio/plan_isv constancy + * semantics — env updates aux only at chunk boundary, same as + * portfolio). Length [n_windows]. NULL = A3 fallback → 0.5. */ + const float* __restrict__ aux_dir_prob_per_env ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_threads = n_windows * chunk_len; @@ -8841,8 +8869,13 @@ extern "C" __global__ void backtest_state_gather_chunk( for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++) pisv[k] = plan_isv[w * SL_PORTFOLIO_PLAN_DIM + k]; + /* SP22 H6 — per-window aux p_up (NULL → 0.5 for Phase 1 A3 eval). */ + float aux_dir_prob = (aux_dir_prob_per_env != NULL) + ? aux_dir_prob_per_env[w] + : 0.5f; + /* Single shared assembly call — produces identical layout to training */ - assemble_state(out, market, ofi, mtf, pf, pisv); + assemble_state(out, market, ofi, mtf, pf, pisv, aux_dir_prob); /* Zero-pad [SL_STATE_DIM .. padded_sd) for cuBLAS K-tile alignment */ for (int k = SL_STATE_DIM; k < padded_sd; k++) diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 30f8743fe..7d1730274 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -22,7 +22,11 @@ use crate::MLError; use super::mapped_pinned::{MappedF32Buffer, MappedU32Buffer}; /// Precompiled epsilon_greedy_kernel cubin, embedded at compile time by build.rs. -static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin")); +// SP22 H6 (2026-05-12): visibility lifted to `pub(crate)` so the collector +// can load the `fill_f32` kernel for `prev_aux_dir_prob` 0.5 sentinel init +// (cold-start + FoldReset). Same cubin, multiple in-crate consumers per the +// existing pattern shared by SP13/SP14 CUBIN statics in `gpu_dqn_trainer.rs`. +pub(crate) static EPSILON_GREEDY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_greedy_kernel.cubin")); /// GPU-fused epsilon-greedy action selector -- pure cudarc API. pub struct GpuActionSelector { diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index fa51d0650..d982c1987 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -1570,7 +1570,13 @@ impl GpuBacktestEvaluator { let padded_sd_i32 = ml_core::state_layout::STATE_DIM_PADDED as i32; // Safety: argument order matches `backtest_state_gather` signature exactly: - // features, portfolio, plan_isv, states_out, n_windows, max_len, feat_dim, current_step, padded_sd + // features, portfolio, plan_isv, states_out, n_windows, max_len, feat_dim, + // current_step, padded_sd, aux_dir_prob_per_env (SP22 H6). + // SP22 H6 (2026-05-12) Phase 1 A3: eval has no aux producer + // infrastructure yet, so pass NULL → kernel uses 0.5 sentinel for + // state[AUX_DIR_PROB_INDEX]. A2 follow-up (aux trunk forward in + // eval) would wire a real per-window p_up buffer here. + let aux_dir_prob_null: u64 = 0; unsafe { self.stream .launch_builder(&self.gather_kernel) @@ -1583,6 +1589,7 @@ impl GpuBacktestEvaluator { .arg(&feat_dim_i32) .arg(&step_i32) .arg(&padded_sd_i32) + .arg(&aux_dir_prob_null) // SP22 H6 A3: NULL → 0.5 sentinel .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!("backtest_state_gather launch step {step}: {e}")))?; } @@ -3164,7 +3171,10 @@ impl GpuBacktestEvaluator { let padded_sd_i32 = ml_core::state_layout::STATE_DIM_PADDED as i32; // Safety: argument order matches `backtest_state_gather` signature exactly: - // features, portfolio, plan_isv, states_out, n_windows, max_len, feat_dim, current_step, padded_sd + // features, portfolio, plan_isv, states_out, n_windows, max_len, feat_dim, + // current_step, padded_sd, aux_dir_prob_per_env (SP22 H6). + // SP22 H6 Phase 1 A3: NULL → kernel uses 0.5 sentinel. + let aux_dir_prob_null: u64 = 0; unsafe { self.stream .launch_builder(&self.gather_kernel) @@ -3177,6 +3187,7 @@ impl GpuBacktestEvaluator { .arg(&feat_dim_i32) .arg(&step_i32) .arg(&padded_sd_i32) + .arg(&aux_dir_prob_null) // SP22 H6 A3: NULL → 0.5 sentinel .launch(launch_cfg) .map_err(|e| { MLError::ModelError(format!("backtest_state_gather launch step {step}: {e}")) @@ -3226,7 +3237,10 @@ impl GpuBacktestEvaluator { // Safety: argument order matches `backtest_state_gather_chunk` signature exactly: // features, portfolio, plan_isv, chunked_states_out (raw u64 ptr), - // n_windows, max_len, feat_dim, start_step, chunk_len, padded_sd + // n_windows, max_len, feat_dim, start_step, chunk_len, padded_sd, + // aux_dir_prob_per_env (SP22 H6). + // SP22 H6 Phase 1 A3: NULL → kernel uses 0.5 sentinel. + let aux_dir_prob_null: u64 = 0; unsafe { self.stream .launch_builder(&self.gather_chunk_kernel) @@ -3240,6 +3254,7 @@ impl GpuBacktestEvaluator { .arg(&start_step_i32) .arg(&chunk_len_i32) .arg(&padded_sd_i32) + .arg(&aux_dir_prob_null) // SP22 H6 A3: NULL → 0.5 sentinel .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!( "backtest_state_gather_chunk launch start_step={start_step} chunk_len={chunk_len}: {e}" diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 23553e9b4..499247d6a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -672,6 +672,18 @@ pub(crate) static SP13_AUX_PRED_TO_ISV_TANH_CUBIN: &[u8] = pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_disagreement_update_kernel.cubin")); +/// SP22 H6 (2026-05-12): per-env p_up extractor. Copies +/// `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` after each aux forward +/// during rollout. The cache buffer is read at the next step's state +/// assembly to populate `state[AUX_DIR_PROB_INDEX = SL_PADDING_START + 0]`, +/// wiring the aux head's directional probability into policy STATE. +/// Loaded from `aux_softmax_to_per_env_kernel.cubin`. Consumed by the +/// collector (`GpuExperienceCollector::collect_experiences_gpu`); no +/// trainer-side consumer (this is a rollout-time bridge, not a training- +/// step kernel). +pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_softmax_to_per_env_kernel.cubin")); + /// SP17 Task PP.3 (2026-05-08): pre-centering V/A mean diagnostic. /// Single-block tree-reduce kernel reading the existing `on_v_logits_buf /// [B, NA]` and `on_b_logits_buf [B, (B0+B1+B2+B3)*NA]` (BRANCH-MAJOR) and diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index ec1279ba8..53ef4d86c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1189,6 +1189,27 @@ pub struct GpuExperienceCollector { /// softmax diff into ISV[375]), `exp_sp14_q_disagreement_update_kernel` /// (vs Q-argmax, into ISV[383/384/389]). exp_aux_nb_softmax_buf: cudarc::driver::CudaSlice, + /// SP22 H6 (2026-05-12) — per-env aux directional probability cache. + /// Sized `[alloc_episodes]` f32 (one slot per env). Written end-of-step + /// by `exp_aux_softmax_to_per_env_kernel` (gathers `aux_softmax[env, 1]` + /// = p_up). Read start-of-next-step by `experience_state_gather` for + /// `state[AUX_DIR_PROB_INDEX = SL_PADDING_START + 0 = 121]`. Cold-start + /// and FoldReset both initialize to 0.5 (neutral; p_up = 50%) via the + /// `fill_f32` GPU kernel — no HtoD per + /// `feedback_no_htod_htoh_only_mapped_pinned.md`. + prev_aux_dir_prob: cudarc::driver::CudaSlice, + /// SP22 H6 (2026-05-12) — copy kernel handle for + /// `aux_softmax_to_per_env_kernel`. Loaded from + /// `SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN`. Launched on the collector's + /// stream right after `aux_next_bar_forward` populates + /// `exp_aux_nb_softmax_buf`. + exp_aux_softmax_to_per_env_kernel: CudaFunction, + /// SP22 H6 (2026-05-12) — `fill_f32` kernel handle (from + /// `EPSILON_GREEDY_CUBIN`). Used to initialize `prev_aux_dir_prob` to + /// 0.5 at construction and at every FoldReset. Pure GPU compute — no + /// HtoD memcpy of an `init_vec` per + /// `feedback_no_htod_htoh_only_mapped_pinned.md`. + exp_fill_f32_kernel: CudaFunction, /// SP14 β-migration step 2 (2026-05-07): rollout-sized aux next-bar /// label tile. Sized `[alloc_episodes]` i32. Filled per-step by a /// thin per-step variant of `aux_sign_label_kernel` that derives @@ -2551,6 +2572,66 @@ impl GpuExperienceCollector { "sp14-β: alloc exp_aux_nb_softmax_buf ({} f32, collector): {e}", alloc_episodes * AUX_NEXT_BAR_K )))?; + // SP22 H6 (2026-05-12): aux→policy state bridge. Load the per-env + // softmax extractor (`aux_softmax_to_per_env_kernel`) + the shared + // `fill_f32` from the epsilon_greedy cubin (for 0.5 sentinel init at + // cold-start and FoldReset). Allocate `prev_aux_dir_prob [alloc_episodes]` + // and prime it to 0.5 via fill_f32 — no HtoD per + // `feedback_no_htod_htoh_only_mapped_pinned.md`. + let exp_aux_softmax_to_per_env_kernel = { + use super::gpu_dqn_trainer::SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN; + let m = stream.context() + .load_cubin(SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: aux_softmax_to_per_env cubin (collector): {e}" + )))?; + m.load_function("aux_softmax_to_per_env_kernel") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: aux_softmax_to_per_env_kernel load (collector): {e}" + )))? + }; + let exp_fill_f32_kernel = { + use super::gpu_action_selector::EPSILON_GREEDY_CUBIN; + let m = stream.context() + .load_cubin(EPSILON_GREEDY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: epsilon_greedy cubin (collector, for fill_f32): {e}" + )))?; + m.load_function("fill_f32") + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: fill_f32 load (collector): {e}" + )))? + }; + let prev_aux_dir_prob = stream + .alloc_zeros::(alloc_episodes) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: alloc prev_aux_dir_prob ({} f32, collector): {e}", + alloc_episodes + )))?; + { + // Cold-start fill: write 0.5 sentinel to every slot. Pure GPU + // compute via `fill_f32` (no HtoD). Matches the launch shape used + // by `GpuActionSelector::run_branching` for the epsilon buffer. + let n_envs_i32 = alloc_episodes as i32; + let sentinel: f32 = 0.5; + let blocks = ((alloc_episodes as u32).div_ceil(256)).max(1); + let prev_aux_dir_prob_ptr = prev_aux_dir_prob.raw_ptr(); + unsafe { + stream + .launch_builder(&exp_fill_f32_kernel) + .arg(&prev_aux_dir_prob_ptr) + .arg(&sentinel) + .arg(&n_envs_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: fill_f32(prev_aux_dir_prob=0.5) cold-start: {e}" + )))?; + } + } let exp_aux_nb_label_buf = stream .alloc_zeros::(alloc_episodes) .map_err(|e| MLError::ModelError(format!( @@ -3066,6 +3147,11 @@ impl GpuExperienceCollector { exp_aux_nb_hidden_buf, exp_aux_nb_logits_buf, exp_aux_nb_softmax_buf, + // SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up + // cache + extractor kernel + fill_f32 handle for 0.5 sentinel. + prev_aux_dir_prob, + exp_aux_softmax_to_per_env_kernel, + exp_fill_f32_kernel, exp_aux_nb_label_buf, exp_aux_dir_acc_buf, // SP14 Layer C Phase C.5a (2026-05-08): rollout-sized aux trunk @@ -5206,6 +5292,10 @@ impl GpuExperienceCollector { let isv_ptr = self.isv_signals_dev_ptr; let ofi_ptr = self.ofi_gpu.dev_ptr; let ofi_dim_i32 = self.ofi_dim as i32; + // SP22 H6 (2026-05-12): per-env aux p_up cache (filled by + // `exp_aux_softmax_to_per_env_kernel` at the END of the prior + // step; sentinel 0.5 at cold-start / FoldReset via fill_f32). + let prev_aux_dir_prob_ptr = self.prev_aux_dir_prob.raw_ptr(); self.stream .launch_builder(&self.state_gather_kernel) @@ -5227,6 +5317,7 @@ impl GpuExperienceCollector { .arg(&isv_ptr) // isv_signals_ptr (NULL=static) .arg(&ofi_ptr) // ofi_features (NULL=no OFI) .arg(&ofi_dim_i32) // ofi_dim (0=no OFI) + .arg(&prev_aux_dir_prob_ptr) // SP22 H6: per-env aux p_up [N] .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!( "experience_state_gather t={t}: {e}" @@ -5730,6 +5821,39 @@ impl GpuExperienceCollector { } } + // SP22 H6 (2026-05-12): aux→policy state bridge. Copy + // per-env p_up = `aux_softmax[env, 1]` into the + // `prev_aux_dir_prob` cache for the NEXT step's state + // assembly. Lives in the same gate as Producers 1a–1d + // because if `trainer_params_ptr == 0` then aux forward + // was skipped and the softmax tile is stale `alloc_zeros` + // data — copying it would write zero everywhere, which + // collapses to a worse-than-sentinel signal (p_up=0 ⇒ + // "down with 100% conviction"). When the gate is false, + // `prev_aux_dir_prob` keeps its previous value (0.5 + // sentinel at cold-start, or last good p_up otherwise). + // Pure GPU per `feedback_no_atomicadd` + no-HtoD discipline. + { + let prev_aux_dir_prob_dev = self.prev_aux_dir_prob.raw_ptr(); + let blocks = ((n as u32).div_ceil(256)).max(1); + unsafe { + self.stream + .launch_builder(&self.exp_aux_softmax_to_per_env_kernel) + .arg(&aux_softmax_dev) + .arg(&prev_aux_dir_prob_dev) + .arg(&n_i32_dir) + .arg(&k_i32_dir) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: aux_softmax_to_per_env t={t}: {e}" + )))?; + } + } + // Producer 2: SP14 q_disagreement_update. Reads aux // softmax (K=2) + q_values (K=4 over the first 4 columns // per row, q_stride=13). Writes ISV[383/384/389] @@ -6984,6 +7108,34 @@ impl GpuExperienceCollector { upload_host_to_cuda_f32( &self.stream, &portfolio_init, &mut self.portfolio_states, "portfolio_states reset")?; + // SP22 H6 (2026-05-12): FoldReset of the per-env aux directional + // probability cache. Re-seed every slot to the 0.5 neutral sentinel + // so the first state-gather call of the new fold reads "no aux + // signal yet" instead of the previous fold's last p_up. Pure GPU + // compute via fill_f32 — no HtoD per + // `feedback_no_htod_htoh_only_mapped_pinned.md`. + { + let n_envs_i32 = self.alloc_episodes as i32; + let sentinel: f32 = 0.5; + let blocks = ((self.alloc_episodes as u32).div_ceil(256)).max(1); + let prev_aux_dir_prob_ptr = self.prev_aux_dir_prob.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.exp_fill_f32_kernel) + .arg(&prev_aux_dir_prob_ptr) + .arg(&sentinel) + .arg(&n_envs_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "sp22-h6: fill_f32(prev_aux_dir_prob=0.5) FoldReset: {e}" + )))?; + } + } + debug!( initial_capital, "Episode state reset complete" diff --git a/crates/ml/src/cuda_pipeline/state_layout.cuh b/crates/ml/src/cuda_pipeline/state_layout.cuh index d7645725b..2e6ed4656 100644 --- a/crates/ml/src/cuda_pipeline/state_layout.cuh +++ b/crates/ml/src/cuda_pipeline/state_layout.cuh @@ -628,7 +628,8 @@ __device__ __forceinline__ void assemble_state( const float* __restrict__ ofi, const float mtf[SL_MTF_DIM], const float portfolio[SL_PORTFOLIO_BASE_DIM], - const float plan_isv[SL_PORTFOLIO_PLAN_DIM] + const float plan_isv[SL_PORTFOLIO_PLAN_DIM], + const float aux_dir_prob ) { // Market features [0..42) for (int k = 0; k < SL_MARKET_DIM; k++) @@ -654,7 +655,11 @@ __device__ __forceinline__ void assemble_state( for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++) out[SL_PLAN_ISV_START + k] = plan_isv[k]; - // Padding [121..128) — 7 zeros for 8-alignment (D.8 Plan 2 Task 6C) - for (int k = 0; k < SL_PADDING_DIM; k++) + // Padding [121..128) — SP22 H6 (2026-05-12): slot 0 = aux directional + // probability (p_up from previous step's aux softmax); slots 1..7 zero + // for 8-alignment. Caller passes 0.5 sentinel at cold-start / FoldReset + // and for eval-time NULL fallback (A3). + out[SL_PADDING_START + 0] = aux_dir_prob; + for (int k = 1; k < SL_PADDING_DIM; k++) out[SL_PADDING_START + k] = 0.0f; } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index a51778a57..e3d9f83cb 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -16276,3 +16276,123 @@ optimizer can park V at Q(Flat) → Q(Flat) attractor. Plan doc: `docs/plans/2026-05-12-sp22-wr-plateau-investigation.md` section H6. + +## 2026-05-12 — SP22 H6 implementation: Phase 1 (training-side + eval A3 fallback) + +### Runbook reference + +- Design + 10-step runbook: `docs/plans/2026-05-12-sp22-h6-aux-policy-state-bridge.md` +- Next-session prompt: `docs/plans/2026-05-12-sp22-h6-next-session-prompt.md` +- Implementation predecessor: `## 2026-05-12 — SP22 H6 design: aux→policy state bridge` (above) + +### What was implemented (Steps 1-10, Phase 1 only) + +State slot 121 (`SL_PADDING_START + 0`) is now the per-env aux directional +probability cache, sourced from the previous rollout step's aux softmax tile. +A2 (eval-side aux integration — adding an aux trunk forward to the closure- +path evaluator) is deferred; eval uses an A3 NULL fallback that resolves to +the 0.5 neutral sentinel. + +Data flow at rollout step `t`: + +1. `experience_state_gather` reads `prev_aux_dir_prob[env]` (written at step + `t-1` end-of-step, or 0.5 sentinel at cold-start / FoldReset) and writes + it into `state[121]`. +2. Forward pass: policy Q-head + (in parallel, separate trunk) aux head → + `aux_softmax [N, K=2]`. +3. New copy kernel `aux_softmax_to_per_env_kernel` gathers + `aux_softmax[env, 1]` (= p_up) into `prev_aux_dir_prob[env]` for step + `t+1`'s state. + +The trunk-separation invariant from `pearl_separate_aux_trunk_when_shared_starves` +is preserved: there is no autograd graph crossing between policy Q-head and +aux trunk. The new wire is a raw `CudaSlice` write/read; stop-grad is +implicit (we're in raw CUDA, not Candle). + +### Critical discipline points + +**No HtoD for the buffer init**. Plan Step 6's snippet used +`stream.memcpy_htod(&init_vec=vec![0.5; N], ...)` which violates +`feedback_no_htod_htoh_only_mapped_pinned.md` (caught mid-flight). Replaced +with project-canonical pattern: pure-GPU `fill_f32` kernel from +`epsilon_greedy_kernel.cu` (same kernel `GpuActionSelector::run_branching` +uses for the per-sample epsilon buffer init). Cold-start fill in +`GpuExperienceCollector::new()`; FoldReset re-fill in `reset_episodes()`. + +**Three state-gather kernels, not one**. Plan named only +`backtest_state_gather`. Actual code has THREE consumers of `assemble_state`: +`experience_state_gather` (training), `backtest_state_gather` (eval per-step), +`backtest_state_gather_chunk` (eval chunked). The signature change to +`assemble_state` (added 7th param `float aux_dir_prob`) required atomic +updates to all three (per `feedback_no_partial_refactor.md`). Two +launchers for `backtest_state_gather` + one for `backtest_state_gather_chunk` +all gained the NULL-arg fallback for Phase 1 A3. + +**Gate-coupling of the copy kernel launch**. The launch sits inside +`if isv_signals_dev_ptr != 0 && trainer_params_ptr != 0` — the SAME +predicate that determines whether aux forward ran this step. When the gate +is false, the softmax tile is stale `alloc_zeros` data; copying it would +write `p_up = 0` ("100% conviction down") into every env's cache, which is +strictly worse than the 0.5 sentinel. Gating preserves the +"sentinel-or-last-good" invariant. + +### Files touched + +New: + +- `crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu` + +Modified: + +- `crates/ml-core/src/state_layout.rs` — `AUX_DIR_PROB_INDEX = PADDING_START` + (just an alias; no STATE_DIM change) +- `crates/ml/src/cuda_pipeline/state_layout.cuh` — `assemble_state` gained + `float aux_dir_prob`; written to `out[SL_PADDING_START + 0]`, remaining + 6 padding slots stay zero for 8-alignment +- `crates/ml/src/cuda_pipeline/experience_kernels.cu` — all 3 state-gather + kernels gained `aux_dir_prob_per_env` arg with NULL-defensive `0.5f` + fallback +- `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` — + `prev_aux_dir_prob` buffer, 2 kernel handles + (`exp_aux_softmax_to_per_env_kernel`, `exp_fill_f32_kernel`), cold-start + fill in `new()`, FoldReset re-fill in `reset_episodes()`, copy-kernel + launch right after aux forward (inside the same isv/trainer-params gate + as Producers 1a–1d), state-gather arg threading +- `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` — NULL + `aux_dir_prob_per_env` arg at all 3 launchers (Phase 1 A3 fallback) +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — + `SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN` static (matches the existing + per-kernel cubin static pattern; placed alongside + `SP13_AUX_PRED_TO_ISV_TANH_CUBIN` and `SP14_Q_DISAGREEMENT_CUBIN`) +- `crates/ml/src/cuda_pipeline/gpu_action_selector.rs` — + `EPSILON_GREEDY_CUBIN` visibility lifted to `pub(crate)` so the collector + can load `fill_f32` from the same cubin +- `crates/ml/build.rs` — registered `aux_softmax_to_per_env_kernel.cu` in + `kernels_with_common` (gets `common_device_functions.cuh` prepended; + trivial since the kernel uses only `cuda_runtime.h` primitives) + +### Verification (pre-smoke) + +| Gate | Result | +|---|---| +| `SQLX_OFFLINE=true cargo check -p ml --features cuda` | 0 errors (21 pre-existing warnings) | +| `cargo test -p ml --test gpu_backtest_validation --features cuda --release` | 4/4 expected-passing tests pass; 2 pre-existing PnL-assertion failures unchanged | +| `/usr/local/cuda/bin/compute-sanitizer --tool=memcheck ` | `ERROR SUMMARY: 0 errors` | + +All three gates were required (per runbook) before smoke dispatch. + +### Next decision points + +After smoke (3-epoch baseline on `ci-training-l40s`): + +- **WR > 50.5%** → H6 confirmed; aux signal is policy-usable when + bridged. Justify A2 follow-up (aux trunk forward in eval) for + production parity. Likely also worth investigating cross-fold aux + generalization (78% within-fold → 49% across-fold per H1 HD[2]) — + Phase 1 deliberately tests with the overfit signal; generalization is + a separate concern. +- **WR pinned at 50.1–50.2%** → H6 falsified. Aux signal is conducted + into state but policy can't use it (or doesn't trust it). Pivot to + H3 (reward density mismatch per + `pearl_event_driven_reward_density_alignment`) or audit the V/A + unidentifiability fix (`project_dueling_va_unidentifiable`).