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

feat(sp22-vnext): Phase B3 — collector-side rollout buffers + forward chain wireup

Browse Source 
Adds collector-side trade-outcome head: 5 struct fields + allocations
+ per-step forward + per-step label producer launches in the rollout
loop. Mirrors the K=2 next-bar head's collector wireup at K=3.

Collector struct additions:
- exp_aux_to_fwd: AuxTradeOutcomeForwardOps (3 kernel handles)
- exp_aux_to_hidden_buf   [alloc_episodes × H=128] saved post-ELU
- exp_aux_to_logits_buf   [alloc_episodes × K=3]   saved logits
- exp_aux_to_softmax_buf  [alloc_episodes × K=3]   softmax tile
- exp_aux_to_label_buf    [alloc_episodes] i32     sparse {-1, 0, 1, 2}

Per-step launches in collect_experiences_gpu rollout loop:

1. aux_trade_outcome_forward — launched immediately after the K=2
   sibling's forward_next_bar, parallel on the same stream. Reads
   exp_h_s2_aux + weights at flat-buffer indices [163..167) (Phase
   B1 additions). Writes hidden/logits/softmax tiles. No consumer
   yet — Phase C wires state assembly; Phase B4 wires trainer
   scatter.

2. trade_outcome_label_kernel — launched immediately after
   experience_env_step on the same stream, reading the save-for-
   backward buffers (pnl_vs_target_at_close_per_env, pnl_vs_stop_at_
   close_per_env) that env_step just wrote at segment_complete.
   Stream-implicit producer→consumer ordering. Emits per-env
   {-1, 0, 1, 2} labels — sparse, ~95-99% bars produce -1 (mask).

Dead-code discipline per feedback_wire_everything_up: every kernel arg
+ producer site is real wiring (not NULL placeholder) — only the
absence of consumers reading the produced tiles is "dead". The smoke
run produces softmax tiles + labels every step bit-identical to
pre-vNext baseline (no consumer = no effect on training behavior).

Phase B4 next: trainer-side replay-batch chain (forward + loss_reduce
+ backward + Adam SAXPY for the 4 new weight tensors).

Audit: docs/dqn-wire-up-audit.md Phase B3 section.

Verification:
- cargo check -p ml clean (21 warnings, none new on aux_to_*).
- cargo test -p ml --lib → 1016 passing / 0 failing (unchanged from
  post-fix-sweep baseline at ebc1b1502).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-14 01:16:45 +02:00
parent ebc1b15023
commit b28b349ac3
2 changed files with 193 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1215,6 +1215,56 @@ 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 vNext Phase B3 (2026-05-14): trade-outcome aux head ──
// collector-side rollout buffers + ops handle. Mirrors the K=2
// `exp_aux_nb_*` pattern at K=3. Forward runs per-rollout-step on
// the collector's stream; the label kernel runs after each
// `experience_env_step` to extract the K=3 outcome label from the
// save-for-backward buffers populated at segment_complete.
//
/// Forward orchestrator (Phase B0 scaffold). Holds the 3 forward-
/// side kernel handles: aux_trade_outcome_forward,
/// aux_trade_outcome_loss_reduce, trade_outcome_label_kernel.
/// Loss reduce isn't invoked in the collector path (backward runs
/// in the trainer's replay-batch path; loss numerator only matters
/// for training); kept on the handle to share the ops struct with
/// the trainer-side wiring (Phase B4) without two parallel loads.
exp_aux_to_fwd: super::gpu_aux_heads::AuxTradeOutcomeForwardOps,
/// `[alloc_episodes × AUX_HIDDEN_DIM=128]` saved post-ELU
/// activation. Mirrors `exp_aux_nb_hidden_buf` for the K=3 head.
/// Backward never runs in the collector path; allocated so the
/// forward kernel's hidden_out arg is non-NULL per
/// `feedback_wire_everything_up`.
exp_aux_to_hidden_buf: cudarc::driver::CudaSlice<f32>,
/// `[alloc_episodes × AUX_OUTCOME_K=3]` saved pre-softmax logits.
/// Parity with `exp_aux_nb_logits_buf`; backward reads the softmax
/// tile not the logits, but the buffer is allocated for contract
/// symmetry.
exp_aux_to_logits_buf: cudarc::driver::CudaSlice<f32>,
/// `[alloc_episodes × AUX_OUTCOME_K=3]` saved softmax tile —
/// THREE eventual consumers:
/// - Phase C 3-slot state assembly (state[121..124] = p_Profit,
/// p_Stop, p_Timeout) — replaces Phase 2's single-slot
/// `prev_aux_dir_prob` for the K=2 head
/// - Trainer-side backward (Phase B4): reads via replay-buffer
/// scatter into trainer's `aux_to_softmax_buf`
/// - Future diagnostic / HEALTH_DIAG slots (Phase F)
/// Phase B3 (this commit): produced per-step but no consumer yet.
exp_aux_to_softmax_buf: cudarc::driver::CudaSlice<f32>,
/// `[alloc_episodes] i32` per-env trade-outcome label in
/// `{-1=mask, 0=Profit, 1=Stop, 2=Timeout}`. Written per-step by
/// `trade_outcome_label_kernel` reading
/// `trade_close_per_sample[env*L+t]` and the save-for-backward
/// `pnl_vs_target_at_close_per_env` / `pnl_vs_stop_at_close_per_env`
/// buffers populated at segment_complete (Phase A3 wireup).
/// Sparse: most bars (~95-99%) emit -1. Consumed (eventual): trainer-
/// side via replay-buffer label scatter (Phase B4 wireup, mirrors
/// `aux_sign_labels` per-(i, t) scatter pattern).
/// Phase B3 (this commit): produced per-step but no replay-buffer
/// scatter yet.
exp_aux_to_label_buf: cudarc::driver::CudaSlice<i32>,
/// SP22 H6 Phase 2 (2026-05-12) — per-env aux directional probability
/// cache, RECENTERED encoding. Sized `[alloc_episodes]` f32 (one slot
/// per env). Written end-of-step by `exp_aux_softmax_to_per_env_kernel`
@@ -2636,6 +2686,40 @@ impl GpuExperienceCollector {
"sp14-β: alloc exp_aux_nb_softmax_buf ({} f32, collector): {e}",
alloc_episodes * AUX_NEXT_BAR_K
)))?;
// ── SP22 H6 vNext Phase B3 (2026-05-14): trade-outcome aux head ──
// collector-side ops + 4 rollout buffers. Mirrors aux_nb_* allocs
// at K=3 (AUX_OUTCOME_K=3) instead of K=2.
//
// Forward orchestrator loads 3 kernel handles from 3 separate
// cubins (forward, loss_reduce, label_producer) — each in its
// own .cu file for diagnostic isolation per Phase A2-A4 design.
let exp_aux_to_fwd = super::gpu_aux_heads::AuxTradeOutcomeForwardOps::new(&stream)?;
use super::gpu_aux_heads::AUX_OUTCOME_K;
let exp_aux_to_hidden_buf = stream
.alloc_zeros::<f32>(alloc_episodes * AUX_HIDDEN_DIM)
.map_err(|e| MLError::ModelError(format!(
"sp22-vnext B3: alloc exp_aux_to_hidden_buf ({} f32): {e}",
alloc_episodes * AUX_HIDDEN_DIM
)))?;
let exp_aux_to_logits_buf = stream
.alloc_zeros::<f32>(alloc_episodes * AUX_OUTCOME_K)
.map_err(|e| MLError::ModelError(format!(
"sp22-vnext B3: alloc exp_aux_to_logits_buf ({} f32): {e}",
alloc_episodes * AUX_OUTCOME_K
)))?;
let exp_aux_to_softmax_buf = stream
.alloc_zeros::<f32>(alloc_episodes * AUX_OUTCOME_K)
.map_err(|e| MLError::ModelError(format!(
"sp22-vnext B3: alloc exp_aux_to_softmax_buf ({} f32): {e}",
alloc_episodes * AUX_OUTCOME_K
)))?;
let exp_aux_to_label_buf = stream
.alloc_zeros::<i32>(alloc_episodes)
.map_err(|e| MLError::ModelError(format!(
"sp22-vnext B3: alloc exp_aux_to_label_buf ({} i32): {e}",
alloc_episodes
)))?;
// 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
@@ -3217,6 +3301,14 @@ impl GpuExperienceCollector {
exp_aux_nb_hidden_buf,
exp_aux_nb_logits_buf,
exp_aux_nb_softmax_buf,
// SP22 H6 vNext Phase B3 (2026-05-14): trade-outcome aux head
// collector-side ops + rollout buffers (Phase B0 scaffold +
// K=3 mirror of aux_nb_* pattern).
exp_aux_to_fwd,
exp_aux_to_hidden_buf,
exp_aux_to_logits_buf,
exp_aux_to_softmax_buf,
exp_aux_to_label_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,
@@ -5557,6 +5649,9 @@ impl GpuExperienceCollector {
if self.trainer_params_ptr != 0 {
use crate::cuda_pipeline::batched_forward::f32_weight_ptrs_from_base;
use super::gpu_aux_heads::AUX_NEXT_BAR_K;
// SP22 H6 vNext Phase B3 (2026-05-14): trade-outcome head
// forward needs K=3 in scope alongside the K=2 sibling's K.
use super::gpu_aux_heads::AUX_OUTCOME_K;
// Param tensor pointers via the same offset arithmetic the
// trainer's `aux_heads_forward` uses (gpu_dqn_trainer.rs:17166).
@@ -5703,6 +5798,38 @@ impl GpuExperienceCollector {
self.exp_aux_nb_logits_buf.raw_ptr(),
self.exp_aux_nb_softmax_buf.raw_ptr(),
)?;
// ── SP22 H6 vNext Phase B3 (2026-05-14): trade-outcome ──
// head forward. Same input (`exp_h_s2_aux`) as the K=2
// sibling above; reads weights at flat-buffer indices
// [163..167) (aux_to_w1, aux_to_b1, aux_to_w2, aux_to_b2)
// added by Phase B1. Writes hidden + logits + softmax
// tiles (rollout-scratch, sized `[n_episodes × K=3]`).
//
// The softmax tile output has no consumer in this commit
// — Phase C will use it for 3-slot state assembly
// (state[121..124] = p_Profit, p_Stop, p_Timeout). Phase B4
// wires the trainer-side scatter into the replay buffer
// for backward. Per Phase B0's "dead-code scaffolding"
// discipline, the launch runs every step but produces
// tiles that nothing reads yet — this lets us validate
// the forward path doesn't crash / NaN before consumers
// land.
let to_w1 = aux_w_ptrs[163];
let to_b1 = aux_w_ptrs[164];
let to_w2 = aux_w_ptrs[165];
let to_b2 = aux_w_ptrs[166];
self.exp_aux_to_fwd.forward(
&self.stream,
self.exp_h_s2_aux.raw_ptr(),
to_w1, to_b1, to_w2, to_b2,
n, // batch dim = n_episodes
sh2,
AUX_OUTCOME_K,
self.exp_aux_to_hidden_buf.raw_ptr(),
self.exp_aux_to_logits_buf.raw_ptr(),
self.exp_aux_to_softmax_buf.raw_ptr(),
)?;
}
// ── 3. Convert logits → expected Q-values ───────────────────
@@ -6746,6 +6873,45 @@ impl GpuExperienceCollector {
)))?;
}
// ── SP22 H6 vNext Phase B3 (2026-05-14): trade-outcome label ──
// Producer launched IMMEDIATELY after experience_env_step on
// the same stream so the save-for-backward buffers
// (`pnl_vs_target_at_close_per_env`, `pnl_vs_stop_at_close
// _per_env`) and the per-(env, t) `trade_close_per_sample`
// flag are valid for reading (stream-implicit producer →
// consumer ordering).
//
// Per-env semantics: kernel emits `out_labels[env]` ∈
// `{-1, 0, 1, 2}` where -1 = no trade close at this bar
// (the common case ~95-99%). Labels persist across rollout
// steps until the next trade-close overwrites them; the
// sparse-CE consumer reads on `label != -1` so stale values
// don't reach the loss numerator.
//
// Phase B3 (this commit): producer wired and runs every
// step. The labels populate `exp_aux_to_label_buf` but
// have no consumer yet — Phase B4 will wire the trainer-
// side replay-buffer scatter alongside the trainer's
// backward chain. Dead-code semantics: the launch is real
// (validates kernel-arg threading + producer ordering)
// but the label buffer is overwritten every step without
// being consumed.
{
let t_i32 = t as i32;
let l_i32 = timesteps as i32;
self.exp_aux_to_fwd.compute_label(
&self.stream,
self.trade_close_per_sample.raw_ptr(),
self.pnl_vs_target_at_close_per_env.raw_ptr(),
self.pnl_vs_stop_at_close_per_env.raw_ptr(),
t_i32,
n, // n_envs in this rollout slice
timesteps,
self.exp_aux_to_label_buf.raw_ptr(),
)?;
let _ = l_i32; // documents the stride contract — consumer reads `env*L + t`
}
// ── 5b. SP15 Wave 2 (2026-05-06): per-step ISV-driven α
// producer + fused post-SP11 reward-axis composer.
//