diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs index d3f1bd9e3..dfd9efffd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -83,6 +83,18 @@ pub(crate) const AUX_HIDDEN_DIM: usize = 128; /// runtime tanh squash per `pearl_bounded_modifier_outputs_require_structural_activation`. pub(crate) const AUX_NEXT_BAR_K: usize = 2; +/// SP22 H6 vNext Phase A3+ (2026-05-13): output cardinality of the +/// trade-outcome aux head. K=3 softmax over `{Profit=0, Stop=1, +/// Timeout=2}`. Replaces the per-bar next-bar direction prediction +/// (K=2) with a per-trade-close 3-way outcome prediction that ALIGNS +/// with the system's actual decision horizon (multi-bar trade +/// lifecycle). The architectural mismatch — per-bar prediction vs +/// multi-bar decisions — was the root cause of the H6 Phase 3 +/// falsification (smoke train-xrkb7: WR=0.4346 with K=2 head fully +/// active vs baseline 0.4338-0.4346). See +/// `docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md`. +pub(crate) const AUX_OUTCOME_K: usize = 3; + /// Output cardinality of the regime classification head (5-way softmax). pub(crate) const AUX_REGIME_K: usize = 5; @@ -629,3 +641,385 @@ impl AuxHeadsBackwardOps { Ok(()) } } + +// ──────────────────────────────────────────────────────────────────────────── +// SP22 H6 vNext: Trade-outcome aux head ops (Phase A2-A5 GPU side wired here) +// ──────────────────────────────────────────────────────────────────────────── +// +// The vNext aux head replaces per-bar direction prediction (K=2) with +// per-trade-close outcome prediction (K=3: {Profit, Stop, Timeout}) — a +// signal that aligns with the system's actual multi-bar decision horizon. +// See `docs/plans/2026-05-14-sp22-h6-vNext-trade-outcome-aux.md` for the +// full architecture spec. +// +// Phase B0 (this commit): orchestrator scaffolding only — additive, +// no production callers yet. Mirrors the AuxHeadsForwardOps / +// AuxHeadsBackwardOps shape so subsequent phases (trainer struct fields, +// `collect_experiences_gpu` wireup) plug in with minimal contract drift. +// +// Reuses existing K-generic `aux_param_grad_reduce` from +// `AuxHeadsBackwardOps` — no separate reducer needed. + +use super::gpu_dqn_trainer::{ + AUX_TRADE_OUTCOME_BACKWARD_CUBIN, AUX_TRADE_OUTCOME_FORWARD_CUBIN, + AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN, TRADE_OUTCOME_LABEL_CUBIN, +}; + +/// Forward + loss-reduce + label-producer orchestrator for the SP22 H6 +/// vNext trade-outcome aux head. Held separate from `AuxHeadsForwardOps` +/// because the trade-outcome head is structurally distinct — different +/// label semantics (per-trade-close vs per-bar), different K (3 vs 2), +/// different input (Phase B switches to `[h_s2_aux || plan_params]` = +/// 262-dim concat vs the K=2 head's 256-dim `h_s2_aux`-only input). +/// +/// Owns three kernel handles: +/// - `aux_trade_outcome_forward` (Phase A3 cubin) +/// - `aux_trade_outcome_loss_reduce` (Phase A4 cubin) +/// - `trade_outcome_label_kernel` (Phase A2 cubin) +/// +/// `pub(crate)` so the collector / trainer wireup commits can construct +/// it via `new()` and call the per-launch methods. +#[allow(missing_debug_implementations)] +pub(crate) struct AuxTradeOutcomeForwardOps { + forward_kernel: CudaFunction, + loss_reduce_kernel: CudaFunction, + label_kernel: CudaFunction, +} + +impl AuxTradeOutcomeForwardOps { + /// Load all three forward-side kernel handles from their precompiled + /// cubins. Mirrors `AuxHeadsForwardOps::new`'s loader pattern but + /// reads three SEPARATE cubins (one per kernel) because the vNext + /// kernels are in distinct `.cu` files for diagnostic isolation. + pub(crate) fn new(stream: &Arc) -> Result { + let context = stream.context(); + + let fwd_module = context + .load_cubin(AUX_TRADE_OUTCOME_FORWARD_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_forward cubin load: {e}" + )))?; + let forward_kernel = fwd_module + .load_function("aux_trade_outcome_forward") + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_forward load: {e}" + )))?; + + let loss_module = context + .load_cubin(AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_loss_reduce cubin load: {e}" + )))?; + let loss_reduce_kernel = loss_module + .load_function("aux_trade_outcome_loss_reduce") + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_loss_reduce load: {e}" + )))?; + + let label_module = context + .load_cubin(TRADE_OUTCOME_LABEL_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "trade_outcome_label_kernel cubin load: {e}" + )))?; + let label_kernel = label_module + .load_function("trade_outcome_label_kernel") + .map_err(|e| MLError::ModelError(format!( + "trade_outcome_label_kernel load: {e}" + )))?; + + Ok(Self { + forward_kernel, + loss_reduce_kernel, + label_kernel, + }) + } + + /// Launch `aux_trade_outcome_forward`: + /// `Linear(h_s2_aux) → ELU → Linear → softmax` over K=3 trade outcomes. + /// + /// Caller-owned buffers (raw `u64` device pointers for graph-capture + /// safety; mirrors `forward_next_bar`): + /// * `h_s2_aux_ptr` — `[B, SH2]` row-major. Phase A3-B0: SH2 = + /// 256 (aux trunk output only). Phase B + /// switches to SH2 = 262 (`h_s2_aux || + /// plan_params`) — kernel arg unchanged, + /// only the caller's W1 shape and the buffer + /// contents change. + /// * `w1_ptr` / `b1_ptr` — `[H, SH2]` / `[H]`. Trainer-owned; + /// Phase B1 lands the allocation. + /// * `w2_ptr` / `b2_ptr` — `[K=3, H]` / `[K=3]`. + /// * `hidden_out_ptr` — `[B, H]` SAVED post-ELU buffer + /// (consumed by backward). + /// * `logits_out_ptr` — `[B, K=3]` SAVED logits (consumed by + /// backward only — Phase A5 backward reads + /// this; current backward reads the softmax + /// tile, but the buffer is allocated for + /// contract symmetry with the K=2 head). + /// * `softmax_out_ptr` — `[B, K=3]` SAVED softmax tile (consumed + /// by loss reduce + backward + Phase C + /// 3-slot state assembly producer). + /// + /// Block: `AUX_BLOCK` threads. Grid: `B` blocks. + /// Shared mem: `(H + K) * sizeof(f32)` bytes. + #[allow(clippy::too_many_arguments)] + pub(crate) fn forward( + &self, + stream: &Arc, + h_s2_aux_ptr: u64, + w1_ptr: u64, + b1_ptr: u64, + w2_ptr: u64, + b2_ptr: u64, + b: usize, + sh2: usize, + k: usize, + hidden_out_ptr: u64, + logits_out_ptr: u64, + softmax_out_ptr: u64, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let sh2_i32 = sh2 as i32; + let k_i32 = k as i32; + let smem_bytes = + (AUX_HIDDEN_DIM as u32 + k as u32) * std::mem::size_of::() as u32; + unsafe { + stream + .launch_builder(&self.forward_kernel) + .arg(&h_s2_aux_ptr) + .arg(&w1_ptr) + .arg(&b1_ptr) + .arg(&w2_ptr) + .arg(&b2_ptr) + .arg(&b_i32) + .arg(&sh2_i32) + .arg(&k_i32) + .arg(&hidden_out_ptr) + .arg(&logits_out_ptr) + .arg(&softmax_out_ptr) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (AUX_BLOCK, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_forward: {e}" + )))?; + } + Ok(()) + } + + /// Launch `aux_trade_outcome_loss_reduce`: mean CE over the valid + /// (label != -1) rows of the K=3 softmax tile. + /// + /// Sparse-label semantic: most bars (~95-99%) have label=-1; the + /// kernel divides by `fmaxf(B_valid, 1.0)` so an all-skip batch + /// produces loss=0 + valid_count=0 (the backward sees valid_count=0 + /// and emits zero gradients across the board, no NaN). + pub(crate) fn loss_reduce( + &self, + stream: &Arc, + softmax_ptr: u64, + labels_ptr: u64, + b: usize, + k: usize, + loss_out_ptr: u64, + valid_count_out_ptr: u64, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let k_i32 = k as i32; + // Two parallel partial-reduction strips (loss + valid-count). + let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::() as u32; + unsafe { + stream + .launch_builder(&self.loss_reduce_kernel) + .arg(&softmax_ptr) + .arg(&labels_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&loss_out_ptr) + .arg(&valid_count_out_ptr) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (AUX_BLOCK, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_loss_reduce: {e}" + )))?; + } + Ok(()) + } + + /// Launch `trade_outcome_label_kernel`: per-env K=3 outcome + /// classification at trade-close. + /// + /// Inputs: + /// * `trade_close_per_sample_ptr` — `[N*L]` i32 flag (read at + /// offset `env*L + t` for the current step). + /// Produced by `experience_env_step` at + /// line ~3093 (`is_close`). + /// * `pnl_vs_target_at_close_ptr` — `[n_envs]` f32 save-for-backward + /// from `experience_env_step` segment_complete + /// (Phase A3 wireup — `pnl_vs_target_at_close + /// _per_env`). + /// * `pnl_vs_stop_at_close_ptr` — `[n_envs]` f32 save-for-backward + /// from `experience_env_step` segment_complete + /// (Phase A3 wireup — `pnl_vs_stop_at_close + /// _per_env`). + /// * `t` — current rollout step (0..timesteps); used to compute + /// the offset `env*L + t` into the per-sample + /// trade-close flag. + /// * `n_envs` — number of envs in the rollout slice. + /// * `l_timesteps` — buffer stride (timesteps per env). + /// + /// Output: `out_labels_ptr [n_envs]` i32 in `{-1, 0, 1, 2}`. `-1` + /// = no trade close at this bar (the common case ~95-99% of bars). + /// + /// Block: 256 threads. Grid: `ceil(n_envs / 256)` blocks (per-env + /// parallelism, no reduction, no atomicAdd). + #[allow(clippy::too_many_arguments)] + pub(crate) fn compute_label( + &self, + stream: &Arc, + trade_close_per_sample_ptr: u64, + pnl_vs_target_at_close_ptr: u64, + pnl_vs_stop_at_close_ptr: u64, + t: i32, + n_envs: usize, + l_timesteps: usize, + out_labels_ptr: u64, + ) -> Result<(), MLError> { + let n_envs_i32 = n_envs as i32; + let l_i32 = l_timesteps as i32; + let block: u32 = 256; + let grid = (n_envs as u32).div_ceil(block); + unsafe { + stream + .launch_builder(&self.label_kernel) + .arg(&trade_close_per_sample_ptr) + .arg(&pnl_vs_target_at_close_ptr) + .arg(&pnl_vs_stop_at_close_ptr) + .arg(&t) + .arg(&n_envs_i32) + .arg(&l_i32) + .arg(&out_labels_ptr) + .launch(LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "trade_outcome_label_kernel: {e}" + )))?; + } + Ok(()) + } +} + +/// Backward orchestrator for the SP22 H6 vNext trade-outcome aux head. +/// Held separate from `AuxHeadsBackwardOps` for the same reasons the +/// forward ops are separate. Owns ONE kernel handle (the K=3 backward) — +/// the per-sample → final-grad batch-reduce step reuses the existing +/// K-generic `aux_param_grad_reduce` from `AuxHeadsBackwardOps`. +/// +/// Wireup commits will hold BOTH `AuxHeadsBackwardOps` (for the +/// K-generic `param_grad_reduce`) AND this struct (for the K=3 +/// backward) on the trainer. +#[allow(missing_debug_implementations)] +pub(crate) struct AuxTradeOutcomeBackwardOps { + backward_kernel: CudaFunction, +} + +impl AuxTradeOutcomeBackwardOps { + /// Load the K=3 backward kernel handle from the Phase A5 cubin. + pub(crate) fn new(stream: &Arc) -> Result { + let context = stream.context(); + let bwd_module = context + .load_cubin(AUX_TRADE_OUTCOME_BACKWARD_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_backward cubin load: {e}" + )))?; + let backward_kernel = bwd_module + .load_function("aux_trade_outcome_backward") + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_backward load: {e}" + )))?; + Ok(Self { backward_kernel }) + } + + /// Launch `aux_trade_outcome_backward` — emits per-sample param-grad + /// partials + per-sample `dh_s2_aux [B, SH2]`. Caller follows up + /// with `param_grad_reduce` (from `AuxHeadsBackwardOps`) for each + /// of {`dW1`, `db1`, `dW2`, `db2`} to collapse along the batch dim. + /// + /// Sparse-label gradient amplification: `valid_count` from the loss + /// kernel is typically small (~1-5% of nominal B) so `inv_B = + /// 1/B_valid` is much larger than the K=2 sibling's `inv_B = + /// 1/(~B)`. Per-trade-close gradients have proportionally higher + /// magnitude — correct credit assignment (rare signal speaks louder) + /// but Phase E's Adam may need class-weighted CE or per-group LR + /// tuning to keep the trade-outcome head's update magnitude + /// comparable to the K=2 sibling. + /// + /// `dh_s2_aux_out_ptr` SAXPYs into the aux trunk's + /// `dh_s2_aux_accum` at the call site — the aux trunk's backward + /// kernel reads this as its upstream gradient. Encoder stop-grad + /// is enforced STRUCTURALLY by `aux_trunk_backward` (no `dx_in` + /// output param), so Q's encoder is protected from aux + /// contamination regardless of gradient magnitude. + #[allow(clippy::too_many_arguments)] + pub(crate) fn backward( + &self, + stream: &Arc, + h_s2_aux_ptr: u64, + w1_ptr: u64, + w2_ptr: u64, + hidden_post_ptr: u64, + softmax_ptr: u64, + labels_ptr: u64, + valid_count_ptr: u64, + b: usize, + sh2: usize, + k: usize, + dw1_partial_ptr: u64, + db1_partial_ptr: u64, + dw2_partial_ptr: u64, + db2_partial_ptr: u64, + dh_s2_aux_out_ptr: u64, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let sh2_i32 = sh2 as i32; + let k_i32 = k as i32; + // Shared mem: 2*H (d_h_pre + h_post caches) + K (d_logits cache). + let smem_bytes = + ((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::() as u32; + unsafe { + stream + .launch_builder(&self.backward_kernel) + .arg(&h_s2_aux_ptr) + .arg(&w1_ptr) + .arg(&w2_ptr) + .arg(&hidden_post_ptr) + .arg(&softmax_ptr) + .arg(&labels_ptr) + .arg(&valid_count_ptr) + .arg(&b_i32) + .arg(&sh2_i32) + .arg(&k_i32) + .arg(&dw1_partial_ptr) + .arg(&db1_partial_ptr) + .arg(&dw2_partial_ptr) + .arg(&db2_partial_ptr) + .arg(&dh_s2_aux_out_ptr) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (AUX_BLOCK, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .map_err(|e| MLError::ModelError(format!( + "aux_trade_outcome_backward: {e}" + )))?; + } + Ok(()) + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 32de49d20..efc68ec9f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2210,6 +2210,44 @@ pub(crate) static AUX_HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR" #[allow(dead_code)] pub(crate) static AUX_HEADS_LOSS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads_loss_ema_kernel.cubin")); +/// SP22 H6 vNext Phase A2 (2026-05-13): trade-outcome label producer. +/// Per-env K=3 outcome classification at trade-close: reads +/// `trade_close_per_sample[i*L+t]` + `pnl_vs_target_at_close[env]` + +/// `pnl_vs_stop_at_close[env]` → writes `out_labels[env]` in +/// {-1 mask, 0=Profit, 1=Stop, 2=Timeout}. Loaded by +/// `gpu_aux_heads::AuxTradeOutcomeForwardOps`. Additive in this commit +/// — wire-up into the collector chain lands in Phase B (atomic). +#[allow(dead_code)] +pub(crate) static TRADE_OUTCOME_LABEL_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/trade_outcome_label_kernel.cubin")); + +/// SP22 H6 vNext Phase A3 (2026-05-13): trade-outcome aux head forward. +/// K=3 softmax over `{Profit, Stop, Timeout}` mirroring +/// `aux_next_bar_forward`'s architecture (Linear → ELU → Linear → softmax) +/// at K=3 instead of K=2. Loaded by +/// `gpu_aux_heads::AuxTradeOutcomeForwardOps`. +#[allow(dead_code)] +pub(crate) static AUX_TRADE_OUTCOME_FORWARD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_trade_outcome_forward_kernel.cubin")); + +/// SP22 H6 vNext Phase A4 (2026-05-13): trade-outcome sparse CE reduce. +/// Single-block shmem-tree reduce over the K=3 softmax tile with `-1` +/// mask handling. Sparse-label arithmetic (~1-5% valid bars). Loaded by +/// `gpu_aux_heads::AuxTradeOutcomeForwardOps`. +#[allow(dead_code)] +pub(crate) static AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_trade_outcome_loss_reduce_kernel.cubin")); + +/// SP22 H6 vNext Phase A5 (2026-05-14): trade-outcome aux head backward. +/// K=3 backward emitting per-sample partials (dW1, db1, dW2, db2, +/// dh_s2_aux). Mirrors `aux_next_bar_backward`'s structure at K=3. +/// Loaded by `gpu_aux_heads::AuxTradeOutcomeBackwardOps`. Reuses the +/// existing K-generic `aux_param_grad_reduce` for the per-sample → +/// final-grad batch-reduce step. +#[allow(dead_code)] +pub(crate) static AUX_TRADE_OUTCOME_BACKWARD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_trade_outcome_backward_kernel.cubin")); + /// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward kernel cubin. /// 3-layer Linear→ELU→Linear→ELU→Linear MLP forward — reads encoder /// output, writes `h_s2_aux` plus saved-for-backward hidden activations. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 2fbca3275..016a7eae1 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17764,3 +17764,19 @@ dW*_partial + db*_partial ──→ aux_param_grad_reduce (existing, K-generic) Phase B is the next major commit and the largest scope shift (Rust-side struct fields for aux head W1/W2/b1/b2 weights, Adam state, hidden_post / logits / softmax / dW*_partial buffers; gpu_aux_heads.rs gets a parallel `AuxTradeOutcomeForwardOps` + `AuxTradeOutcomeBackwardOps` struct; Phase B wires the chain into `gpu_experience_collector.rs::collect_experiences_gpu` immediately after the K=2 head's launches). Cargo check clean. Phase A complete. Ready for Phase B. + +#### Phase B0 — AuxTradeOutcome ops struct scaffolding (2026-05-14) + +First Rust-side commit of Phase B. Pure scaffolding: adds orchestrator structs in `gpu_aux_heads.rs` parallel to the existing `AuxHeadsForwardOps` / `AuxHeadsBackwardOps`, plus the 4 cubin embeds in `gpu_dqn_trainer.rs`. Zero production callers — additive per `feedback_wire_everything_up`'s commit-by-commit ordering convention (mirrors the original SP14 Layer C aux trunk pattern: scaffold → trainer fields → collector wireup). + +**Changes**: +- `gpu_dqn_trainer.rs`: 4 new `pub(crate) static` cubin embeds — `TRADE_OUTCOME_LABEL_CUBIN`, `AUX_TRADE_OUTCOME_FORWARD_CUBIN`, `AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN`, `AUX_TRADE_OUTCOME_BACKWARD_CUBIN`. All `#[allow(dead_code)]` until wireup. +- `gpu_aux_heads.rs`: New constant `AUX_OUTCOME_K = 3` (parallel to `AUX_NEXT_BAR_K = 2`). +- `gpu_aux_heads.rs`: New struct `AuxTradeOutcomeForwardOps` holding 3 kernel handles (forward + loss_reduce + label producer). Launch methods: `forward()`, `loss_reduce()`, `compute_label()`. Mirrors `AuxHeadsForwardOps`'s contract shape for forward + loss_reduce; the label producer is per-env-only with grid `ceil(n_envs/256)` (no reduction). +- `gpu_aux_heads.rs`: New struct `AuxTradeOutcomeBackwardOps` holding 1 kernel handle (backward). Launch method: `backward()`. **Reuses the existing K-generic `aux_param_grad_reduce`** from `AuxHeadsBackwardOps` — no separate reducer needed. + +**Phase B1 next**: trainer struct fields for W1/b1/W2/b2 + Adam state (m/v per param) + Xavier init. Phase B1 lands the weight tensors and reset registry; B2 lands saved-tensor + per-sample partial buffers; B3 wires the forward chain into `collect_experiences_gpu`; B4 wires backward. After all that, the spec's actual Phase B (input concat 256→262 with `plan_params`) becomes a small change touching only the buffer fill + W1 shape. + +**Discipline**: `feedback_no_partial_refactor` honored — cubin embed + ops struct + load function names are atomic; the ops methods can't be called without the trainer struct fields, but having the methods land first lets the field-allocation commits compile-check the wireup endpoint by endpoint. + +Cargo check clean (21 warnings, none new — all pre-existing `DevicePtrMut` unused imports + sp14 unused isv_slots).