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

feat(sp22): H6 Phase 1 — aux→policy state bridge (atomic)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-12 21:31:32 +02:00
parent ff9ec76f18
commit 7fc9799343
10 changed files with 453 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml-core/src/state_layout.rs
View File

@@ -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");

13
crates/ml/build.rs
View File

@@ -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

79
crates/ml/src/cuda_pipeline/aux_softmax_to_per_env_kernel.cu Normal file
View File

@@ -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 <cuda_runtime.h>
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];
}

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

@@ -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++)

6
crates/ml/src/cuda_pipeline/gpu_action_selector.rs
View File

@@ -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 {

21
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -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}"

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

@@ -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

152
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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<f32>,
/// 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<f32>,
/// 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::<f32>(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::<i32>(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"

11
crates/ml/src/cuda_pipeline/state_layout.cuh
View File

@@ -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;
}

120
docs/dqn-wire-up-audit.md
View File

@@ -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<f32>` 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 <test_bin>` | `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`).