From d2331f2f6206c82a5d04f35522701af32b3ea6b7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 10:53:23 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22-vnext):=20Phase=20B5b=20=E2=80=94=20f?= =?UTF-8?q?ull=20plan-conditioning=20integration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The K=3 head's input is now [h_s2_aux (256) || plan_params (6)] = 262- dim, matching the spec's intended architecture. Implements the input concat in both the trainer's replay-batch path and the collector's rollout-step path, with appropriate handling of the backward-side stride mismatch. NEW kernel strided_row_saxpy_kernel.cu: row-truncating SAXPY that accumulates first n_cols_copy columns per row of src [B, src_cols] into dst [B, dst_cols] scaled by alpha. Handles stride mismatch (src_cols != dst_cols). Needed because backward emits dh_s2_aux_to_buf [B, 262] but dh_s2_aux_accum [B, 256] only consumes first 256 cols. PLAN_PARAM_DIM = 6 constant in gpu_aux_heads.rs. Ops struct updates: - AuxTradeOutcomeForwardOps gains concat_kernel + launch_concat() - AuxTradeOutcomeBackwardOps gains strided_saxpy_kernel + launch_strided_row_saxpy() - Both load new cubins in new() Weight tensor resize: - sizes[163] = H × (SH2 + PLAN_PARAM_DIM) = 128 × 262 = 33,536 floats - fan_dims[163] = (H, SH2 + PLAN_PARAM_DIM) Trainer changes: - New aux_to_input_buf [B × 262] field - aux_dh_s2_to_buf resized to [B × 262] - aux_partial_to_w1 resized to [B × H × 262] - max_aux_tensor_len bumped for param_grad_final scratch Trainer forward (aux_heads_forward): - Concat h_s2_aux + plan_params_buf → aux_to_input_buf - forward() with SH2_TOTAL=262 Trainer backward (aux_heads_backward): - backward() with SH2_TOTAL=262; reads aux_to_input_buf - saxpy_f32_kernel SAXPY for dh_s2_aux REPLACED by launch_strided_row_saxpy: copies only first SH2=256 cols per row; trailing 6 cols (plan_params gradient) are dropped — STOP-GRAD on trade plan head from aux loss. Collector forward (rollout): - New exp_aux_to_input_buf [N × 262] field - Concat with plan_params_ptr = 0 (NULL) → zero-fill trailing 6 cols - forward() with SH2_TOTAL=262 Train/inference asymmetry (documented): - Trainer: real plan_params from trade plan head output - Collector: zeros (no trade plan launch in collector) - The head is trained on real plan-conditional outcomes but queried at rollout time with plan_params=0. Phase B5b-2 follow-up would add a trade plan launch to the collector to resolve. Deferred — current state is functional, head still receives plan signal during training. Stop-grad on plan_params: backward writes full [B, 262] gradient, but strided SAXPY only copies first 256 cols. Trade plan head weights NOT trained by K=3 aux loss in this commit. Verification: - cargo check -p ml clean. - cargo test -p ml --lib → 1016/0 green. Phase D next: 12-weight W atom-shift (4 actions × 3 outcomes). Audit: docs/dqn-wire-up-audit.md Phase B5b section. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 8 + crates/ml/src/cuda_pipeline/gpu_aux_heads.rs | 152 ++++++++++++++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 176 +++++++++++++----- .../cuda_pipeline/gpu_experience_collector.rs | 70 ++++++- .../cuda_pipeline/strided_row_saxpy_kernel.cu | 78 ++++++++ docs/dqn-wire-up-audit.md | 50 +++++ 6 files changed, 487 insertions(+), 47 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/strided_row_saxpy_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index d915b48a4..baecfb8b3 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -841,6 +841,14 @@ fn main() { // this commit — Phase B5b allocates buffers + wires launches + // flips SH2_TOTAL=262 in forward/backward. "aux_to_input_concat_kernel.cu", + // SP22 H6 vNext Phase B5b (2026-05-14): row-truncating SAXPY for + // backward dh_s2_aux accumulate. Adds first n_cols_copy columns + // per row of [B, src_cols] into [B, dst_cols] with row stride + // mismatch handling. Needed when B5b's plan-conditioned backward + // emits dh_s2_aux_to_buf [B, 262] but the aux trunk's accumulator + // dh_s2_aux_accum [B, 256] only consumes first 256 cols (plan_ + // params gradient is stop-grad). + "strided_row_saxpy_kernel.cu", // SP22 H6 vNext Phase C (2026-05-14): per-env 3-slot softmax // extractor for the K=3 trade-outcome aux head. Copies // aux_outcome_softmax [n_envs, K=3] → prev_aux_outcome_probs diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs index 74a7d3693..c4f539304 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -661,10 +661,19 @@ impl AuxHeadsBackwardOps { // `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, + AUX_TO_INPUT_CONCAT_CUBIN, AUX_TRADE_OUTCOME_BACKWARD_CUBIN, + AUX_TRADE_OUTCOME_FORWARD_CUBIN, AUX_TRADE_OUTCOME_LOSS_REDUCE_CUBIN, + STRIDED_ROW_SAXPY_CUBIN, TRADE_OUTCOME_LABEL_CUBIN, }; +/// SP22 H6 vNext Phase B5b (2026-05-14): trade-plan parameter dimension. +/// The trade plan head's output [B, P=6] gets concatenated with h_s2_aux +/// [B, 256] to form the K=3 trade-outcome head's plan-conditioned input +/// [B, SH2_TOTAL = 256 + 6 = 262]. Six features per the trade plan ABI: +/// {target_bars, profit_target, stop_loss, conviction, asymmetry, +/// trail_distance}. +pub(crate) const PLAN_PARAM_DIM: usize = 6; + /// 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 @@ -684,6 +693,11 @@ pub(crate) struct AuxTradeOutcomeForwardOps { forward_kernel: CudaFunction, loss_reduce_kernel: CudaFunction, label_kernel: CudaFunction, + /// SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioning concat + /// kernel. Writes `aux_to_input [B, SH2+P=262]` from `h_s2_aux + /// [B, SH2=256]` || `plan_params [B, P=6]`. NULL-tolerant on + /// plan_params (zeros for trailing P cols). + concat_kernel: CudaFunction, } impl AuxTradeOutcomeForwardOps { @@ -727,13 +741,76 @@ impl AuxTradeOutcomeForwardOps { "trade_outcome_label_kernel load: {e}" )))?; + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioning concat + // kernel handle. Loaded once at ops construction so production + // launches stay graph-capture-friendly. + let concat_module = context + .load_cubin(AUX_TO_INPUT_CONCAT_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "aux_to_input_concat cubin load: {e}" + )))?; + let concat_kernel = concat_module + .load_function("aux_to_input_concat") + .map_err(|e| MLError::ModelError(format!( + "aux_to_input_concat load: {e}" + )))?; + Ok(Self { forward_kernel, loss_reduce_kernel, label_kernel, + concat_kernel, }) } + /// SP22 H6 vNext Phase B5b (2026-05-14): launch the input concat + /// kernel to build `aux_to_input [B, SH2+P]` from `h_s2_aux + /// [B, SH2]` || `plan_params [B, P]`. Used by both the trainer's + /// `aux_heads_forward` (real plan_params from the trade plan head) + /// and the collector's rollout-step forward (plan_params_ptr = 0 + /// → zero-fill, the train/inference asymmetry documented in the + /// kernel doc). + /// + /// Block: 256 threads. Grid: `ceil((B * (SH2 + P)) / 256)` blocks + /// (one thread per output element). + pub(crate) fn launch_concat( + &self, + stream: &Arc, + h_s2_aux_ptr: u64, + plan_params_ptr: u64, + b: usize, + sh2: usize, + p: usize, + out_ptr: u64, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let sh2_i32 = sh2 as i32; + let p_i32 = p as i32; + let total_cols = sh2 + p; + let total = b * total_cols; + let block: u32 = 256; + let grid = (total as u32).div_ceil(block); + unsafe { + stream + .launch_builder(&self.concat_kernel) + .arg(&h_s2_aux_ptr) + .arg(&plan_params_ptr) + .arg(&b_i32) + .arg(&sh2_i32) + .arg(&p_i32) + .arg(&out_ptr) + .launch(LaunchConfig { + grid_dim: (grid.max(1), 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "aux_to_input_concat: {e}" + )))?; + } + Ok(()) + } + /// Launch `aux_trade_outcome_forward`: /// `Linear(h_s2_aux) → ELU → Linear → softmax` over K=3 trade outcomes. /// @@ -937,6 +1014,12 @@ impl AuxTradeOutcomeForwardOps { #[allow(missing_debug_implementations)] pub(crate) struct AuxTradeOutcomeBackwardOps { backward_kernel: CudaFunction, + /// SP22 H6 vNext Phase B5b (2026-05-14): row-truncating SAXPY for + /// dh_s2_aux accumulate when backward emits `[B, SH2_TOTAL=262]` + /// but the upstream consumer `dh_s2_aux_accum [B, SH2=256]` only + /// reads the first 256 cols (h_s2_aux gradient; plan_params + /// gradient is stop-grad). + strided_saxpy_kernel: CudaFunction, } impl AuxTradeOutcomeBackwardOps { @@ -953,7 +1036,70 @@ impl AuxTradeOutcomeBackwardOps { .map_err(|e| MLError::ModelError(format!( "aux_trade_outcome_backward load: {e}" )))?; - Ok(Self { backward_kernel }) + + // SP22 H6 vNext Phase B5b (2026-05-14): strided row SAXPY kernel + // handle. Used by `aux_heads_backward` to feed only the first + // SH2=256 cols of `dh_s2_aux_to_buf [B, 262]` into + // `dh_s2_aux_accum [B, 256]` — plan_params gradient (trailing + // 6 cols) is dropped (stop-grad). + let strided_module = context + .load_cubin(STRIDED_ROW_SAXPY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!( + "strided_row_saxpy cubin load: {e}" + )))?; + let strided_saxpy_kernel = strided_module + .load_function("strided_row_saxpy_f32") + .map_err(|e| MLError::ModelError(format!( + "strided_row_saxpy_f32 load: {e}" + )))?; + + Ok(Self { backward_kernel, strided_saxpy_kernel }) + } + + /// SP22 H6 vNext Phase B5b (2026-05-14): launch strided row SAXPY + /// to accumulate first `n_cols_copy` columns of `src [B, src_cols]` + /// into `dst [B, dst_cols]` scaled by `alpha`. Different row + /// strides on src and dst — element-wise SAXPY won't work when + /// src_cols != dst_cols. + #[allow(clippy::too_many_arguments)] + pub(crate) fn launch_strided_row_saxpy( + &self, + stream: &Arc, + dst_ptr: u64, + src_ptr: u64, + alpha: f32, + b: usize, + src_cols: usize, + dst_cols: usize, + n_cols_copy: usize, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let src_i32 = src_cols as i32; + let dst_i32 = dst_cols as i32; + let n_i32 = n_cols_copy as i32; + let total = b * n_cols_copy; + let block: u32 = 256; + let grid = (total as u32).div_ceil(block); + unsafe { + stream + .launch_builder(&self.strided_saxpy_kernel) + .arg(&dst_ptr) + .arg(&src_ptr) + .arg(&alpha) + .arg(&b_i32) + .arg(&src_i32) + .arg(&dst_i32) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (grid.max(1), 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "strided_row_saxpy_f32: {e}" + )))?; + } + Ok(()) } /// Launch `aux_trade_outcome_backward` — emits per-sample param-grad diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 31882d1f1..964483c0c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -693,6 +693,19 @@ pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = pub(crate) static SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_outcome_softmax_to_per_env_kernel.cubin")); +/// SP22 H6 vNext Phase B5b (2026-05-14): input concat kernel cubin for +/// plan-conditioned K=3 forward path. Loaded by +/// `AuxTradeOutcomeForwardOps` and the collector's rollout path. +#[allow(dead_code)] +pub(crate) static AUX_TO_INPUT_CONCAT_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_to_input_concat_kernel.cubin")); + +/// SP22 H6 vNext Phase B5b (2026-05-14): row-truncating SAXPY cubin for +/// backward dh_s2_aux accumulate. Loaded by `AuxTradeOutcomeBackwardOps`. +#[allow(dead_code)] +pub(crate) static STRIDED_ROW_SAXPY_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/strided_row_saxpy_kernel.cubin")); + /// SP22 H6 Phase 3 α structural-prior init (atom-shift design 2026-05-13). /// One-time GPU init kernel that writes the per-action prior /// `[-0.5, 0.0, +0.5, 0.0]` for `[Short, Hold, Long, Flat]` into the @@ -4783,10 +4796,15 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT { let h_dim = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; let k_to = crate::cuda_pipeline::gpu_aux_heads::AUX_OUTCOME_K; - sizes[163] = h_dim * sh2; // aux_to_w1 [H, SH2] - sizes[164] = h_dim; // aux_to_b1 [H] - sizes[165] = k_to * h_dim; // aux_to_w2 [K=3, H] - sizes[166] = k_to; // aux_to_b2 [K=3] + let p_dim = crate::cuda_pipeline::gpu_aux_heads::PLAN_PARAM_DIM; + // SP22 H6 vNext Phase B5b (2026-05-14): W1 input dim grew from + // SH2=256 to SH2_TOTAL = SH2 + PLAN_PARAM_DIM = 262, because the + // K=3 head's input is now `[h_s2_aux || plan_params]`. Adds + // 6 × H = 768 floats per W1 vs the pre-B5b shape. + sizes[163] = h_dim * (sh2 + p_dim); // aux_to_w1 [H, SH2+P=262] + sizes[164] = h_dim; // aux_to_b1 [H] + sizes[165] = k_to * h_dim; // aux_to_w2 [K=3, H] + sizes[166] = k_to; // aux_to_b2 [K=3] } debug_assert!(NUM_WEIGHT_TENSORS == 167, @@ -6407,6 +6425,18 @@ pub struct GpuDqnTrainer { /// `[B, AUX_OUTCOME_K=3]` per-sample db2 partial. Reduced into the /// flat `params_buf` Adam slot [166]. aux_partial_to_b2: CudaSlice, + /// SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input + /// buffer for the K=3 trade-outcome forward + backward. Sized + /// `[B, SH2 + PLAN_PARAM_DIM = 262]` f32 row-major. Populated each + /// step by `aux_to_fwd.launch_concat(h_s2_aux, plan_params, ...)` + /// before the forward call. Forward + backward both read this as + /// their "h_s2_aux" arg (now widened to include plan_params); the + /// kernel takes SH2 as a runtime arg so passing SH2_TOTAL=262 just + /// works. Backward writes `dh_s2_aux_to_buf [B, 262]` (matching + /// stride); the trainer's downstream SAXPY into `dh_s2_aux_accum + /// [B, 256]` uses the new strided-row-saxpy kernel to truncate + /// trailing 6 cols (plan_params gradient is stop-grad). + aux_to_input_buf: CudaSlice, /// SP20 Phase 5 (2026-05-10): per-sample aux confidence at the SAMPLED /// state — `[batch_size]` f32, range `[0, 0.5]` (K=2 peak-softmax-above- /// uniform; 0.0 = uniform / no information sentinel; 0.5 = peak softmax @@ -19478,11 +19508,26 @@ impl GpuDqnTrainer { let to_b1 = on_w_ptrs[164]; let to_w2 = on_w_ptrs[165]; let to_b2 = on_w_ptrs[166]; - self.aux_to_fwd.forward( + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input. + // Concat h_s2_aux [B, 256] || plan_params [B, 6] → aux_to_input + // [B, 262]. Trainer has real plan_params from the trade plan + // head's output. SH2_TOTAL = SH2 + PLAN_PARAM_DIM = 262. + let p_dim = super::gpu_aux_heads::PLAN_PARAM_DIM; + let sh2_total = sh2 + p_dim; + let aux_to_input_ptr = self.aux_to_input_buf.raw_ptr(); + let plan_params_ptr = self.plan_params_buf.raw_ptr(); + self.aux_to_fwd.launch_concat( &self.stream, h_s2_aux_ptr, + plan_params_ptr, + b, sh2, p_dim, + aux_to_input_ptr, + )?; + self.aux_to_fwd.forward( + &self.stream, + aux_to_input_ptr, to_w1, to_b1, to_w2, to_b2, - b, sh2, AUX_OUTCOME_K, + b, sh2_total, AUX_OUTCOME_K, self.aux_to_hidden_buf.raw_ptr(), self.aux_to_logits_buf.raw_ptr(), self.aux_to_softmax_buf.raw_ptr(), @@ -19612,15 +19657,26 @@ impl GpuDqnTrainer { let to_w1_bwd = on_w_ptrs[163]; let to_w2_bwd = on_w_ptrs[165]; let aux_kto = AUX_OUTCOME_K; + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned backward. + // Reads aux_to_input_buf (the [B, 262] concat tile produced by + // aux_heads_forward); passes SH2_TOTAL=262 so the kernel iterates + // the full plan-conditioned input dim. Writes aux_dh_s2_to_buf + // [B, 262] — full gradient including the plan_params columns. + // The downstream SAXPY into dh_s2_aux_accum [B, 256] uses the + // new strided-row-saxpy to copy only the first 256 cols + // (plan_params gradient = stop-grad). + let p_dim_bwd = super::gpu_aux_heads::PLAN_PARAM_DIM; + let sh2_total_bwd = sh2 + p_dim_bwd; + let aux_to_input_ptr = self.aux_to_input_buf.raw_ptr(); self.aux_to_bwd.backward( &self.stream, - h_s2_aux_ptr, + aux_to_input_ptr, to_w1_bwd, to_w2_bwd, self.aux_to_hidden_buf.raw_ptr(), self.aux_to_softmax_buf.raw_ptr(), self.aux_to_label_buf.raw_ptr(), self.aux_to_valid_count_buf.raw_ptr(), - b, sh2, aux_kto, + b, sh2_total_bwd, aux_kto, self.aux_partial_to_w1.raw_ptr(), self.aux_partial_to_b1.raw_ptr(), self.aux_partial_to_w2.raw_ptr(), @@ -19654,8 +19710,12 @@ impl GpuDqnTrainer { (124, self.aux_partial_rg_b1.raw_ptr(), aux_h), (125, self.aux_partial_rg_w2.raw_ptr(), aux_kr * aux_h), (126, self.aux_partial_rg_b2.raw_ptr(), aux_kr), - // SP22 H6 vNext Phase B4: trade-outcome head (K=3) param grads. - (163, self.aux_partial_to_w1.raw_ptr(), aux_h * sh2), + // SP22 H6 vNext Phase B4 + B5b: trade-outcome head (K=3) + // param grads. Phase B5b bumped W1 row size to SH2+P=262 to + // accommodate the plan-conditioned input concat. param_grad + // _reduce iterates per-element so the larger flat size is + // handled uniformly. + (163, self.aux_partial_to_w1.raw_ptr(), aux_h * (sh2 + super::gpu_aux_heads::PLAN_PARAM_DIM)), (164, self.aux_partial_to_b1.raw_ptr(), aux_h), (165, self.aux_partial_to_w2.raw_ptr(), aux_kto * aux_h), (166, self.aux_partial_to_b2.raw_ptr(), aux_kto), @@ -19748,30 +19808,33 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!( "aux_heads_backward saxpy dh_s2_aux regime: {e}" )))?; - // SP22 H6 vNext Phase B4 (2026-05-14): trade-outcome head's - // per-sample dh_s2_aux. SAXPYs into the SAME `dh_s2_aux_accum` - // buffer as the K=2 / K=5 heads — all three head's gradients - // flow through the aux trunk to its own weights, NOT into Q's - // encoder (structural stop-grad in `aux_trunk_backward`). - // Same `aux_w` scaling as the K=2/K=5 SAXPYs above; Phase E - // may differentiate per-head weights if magnitude balance - // requires it. - let dh_to_ptr = self.aux_dh_s2_to_buf.raw_ptr(); - self.stream - .launch_builder(&self.saxpy_f32_kernel) - .arg(&dh_s2_aux_accum_ptr) - .arg(&dh_to_ptr) - .arg(&aux_w) - .arg(&n_dh) - .launch(LaunchConfig { - grid_dim: (blocks_dh, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }) - .map_err(|e| MLError::ModelError(format!( - "aux_heads_backward saxpy dh_s2_aux trade_outcome: {e}" - )))?; + // SP22 H6 vNext Phase B4 + B5b (2026-05-14): trade-outcome + // head's per-sample dh_s2_aux. SAXPYs into the SAME + // `dh_s2_aux_accum [B, SH2=256]` as the K=2 / K=5 heads + // — all three heads' h_s2_aux gradients accumulate into + // the aux trunk's upstream accumulator (NOT Q's encoder). + // + // Phase B5b stride handling: aux_dh_s2_to_buf is now + // [B, SH2_TOTAL=262] (plan-conditioned backward output), + // but dh_s2_aux_accum is [B, SH2=256]. Use strided-row- + // saxpy to copy only the first SH2=256 cols per row. + // The trailing 6 cols per row are the plan_params + // gradient — DROPPED (stop-grad; we don't backprop + // through the trade plan head from aux loss in this + // commit). } + let p_dim_acc = super::gpu_aux_heads::PLAN_PARAM_DIM; + let sh2_total_acc = sh2 + p_dim_acc; + self.aux_to_bwd.launch_strided_row_saxpy( + &self.stream, + dh_s2_aux_accum_ptr, + self.aux_dh_s2_to_buf.raw_ptr(), + aux_w, // alpha + b, // batch dim + sh2_total_acc, // src_cols = 262 + sh2, // dst_cols = 256 + sh2, // n_cols_copy = 256 (first SH2, drops plan_params grad) + )?; // SP14 Layer C Phase C.5b: aux trunk backward — propagates the // accumulated upstream gradient `dh_s2_aux_accum` through the @@ -23124,9 +23187,16 @@ impl GpuDqnTrainer { let aux_partial_rg_b2 = alloc_f32(&stream, aux_b * aux_kr, "aux_partial_rg_b2")?; // SP13 B1.1a: include the K_NB-scaled tensor sizes in the max. // Both `aux_knb * aux_h` and `aux_knb` need coverage now. + // SP22 H6 vNext Phase B5b (2026-05-14): include K=3 trade-outcome + // head's plan-conditioned W1 size (H × SH2_TOTAL = H × 262). + // Computed inline so the constant flows down: aux_p_dim + + // aux_sh2 = SH2_TOTAL bumps the worst-case from 32,768 (K=2's + // W1) to 33,536 (K=3's plan-conditioned W1). + let aux_p_dim_max = super::gpu_aux_heads::PLAN_PARAM_DIM; let max_aux_tensor_len = (aux_h * aux_sh2) .max(aux_kr * aux_h) .max(aux_knb * aux_h) + .max(aux_h * (aux_sh2 + aux_p_dim_max)) // SP22 vNext B5b: K=3 plan-conditioned W1 .max(aux_h) .max(aux_kr) .max(aux_knb); @@ -23159,6 +23229,15 @@ impl GpuDqnTrainer { let aux_to_fwd = AuxTradeOutcomeForwardOps::new(&stream)?; let aux_to_bwd = AuxTradeOutcomeBackwardOps::new(&stream)?; let aux_kto = AUX_OUTCOME_K; + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input + // dim. SH2_TOTAL = SH2 + PLAN_PARAM_DIM = 262. Affects: + // - aux_dh_s2_to_buf [B, SH2_TOTAL] + // - aux_partial_to_w1 [B, H, SH2_TOTAL] + // - aux_to_input_buf [B, SH2_TOTAL] (NEW) + // - aux_param_grad_final_buf: max tensor size now H × SH2_TOTAL + // - sizes[163] (already bumped in compute_param_sizes) + let aux_p_dim = super::gpu_aux_heads::PLAN_PARAM_DIM; + let aux_sh2_total = aux_sh2 + aux_p_dim; let aux_to_hidden_buf = alloc_f32(&stream, aux_b * aux_h, "aux_to_hidden_buf")?; let aux_to_logits_buf = alloc_f32(&stream, aux_b * aux_kto, "aux_to_logits_buf")?; let aux_to_softmax_buf = alloc_f32(&stream, aux_b * aux_kto, "aux_to_softmax_buf")?; @@ -23168,15 +23247,19 @@ impl GpuDqnTrainer { )))?; let aux_to_loss_scalar_buf = alloc_f32(&stream, 1, "aux_to_loss_scalar_buf")?; let aux_to_valid_count_buf = alloc_f32(&stream, 1, "aux_to_valid_count_buf")?; - let aux_dh_s2_to_buf = alloc_f32(&stream, aux_b * aux_sh2, "aux_dh_s2_to_buf")?; - let aux_partial_to_w1 = alloc_f32(&stream, aux_b * aux_h * aux_sh2, "aux_partial_to_w1")?; + let aux_dh_s2_to_buf = alloc_f32(&stream, aux_b * aux_sh2_total, "aux_dh_s2_to_buf")?; + let aux_partial_to_w1 = alloc_f32(&stream, aux_b * aux_h * aux_sh2_total, "aux_partial_to_w1")?; let aux_partial_to_b1 = alloc_f32(&stream, aux_b * aux_h, "aux_partial_to_b1")?; let aux_partial_to_w2 = alloc_f32(&stream, aux_b * aux_kto * aux_h, "aux_partial_to_w2")?; let aux_partial_to_b2 = alloc_f32(&stream, aux_b * aux_kto, "aux_partial_to_b2")?; + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input + // buffer. Populated by aux_to_fwd.launch_concat before each + // forward call. + let aux_to_input_buf = alloc_f32(&stream, aux_b * aux_sh2_total, "aux_to_input_buf")?; // NOTE: aux_param_grad_final_buf is sized to handle the largest - // tensor across BOTH K=2 (next_bar) and K=3 (trade_outcome) - // heads — both heads' largest tensor is W1 [H, SH2] = 32,768 - // floats, identical. No re-sizing needed. + // tensor. Phase B5b grew aux_to_w1 to H × SH2_TOTAL = 33,536 + // floats — larger than the K=2/K=5 W1 (32,768 floats). Need to + // resize the scratch. // SP20 Phase 5: per-sample aux_conf at sampled state ([B] f32). PER's // direct-to-trainer `gather_f32_scalar` writes here on every sample // step once `set_trainer_aux_conf_ptr` is wired; until then, the @@ -25610,6 +25693,9 @@ impl GpuDqnTrainer { aux_partial_to_b1, aux_partial_to_w2, aux_partial_to_b2, + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned + // input buffer for the K=3 forward + backward. + aux_to_input_buf, aux_conf_at_state_buf, aux_weight, stochastic_depth_scale_buf, @@ -33657,10 +33743,16 @@ impl GpuDqnTrainer { { let h_dim = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; let k_to = crate::cuda_pipeline::gpu_aux_heads::AUX_OUTCOME_K; - fan_dims[163] = (h_dim, cfg.shared_h2); // aux_to_w1 [H, SH2] (Xavier) - fan_dims[164] = (0, 0); // aux_to_b1 (zero) - fan_dims[165] = (k_to, h_dim); // aux_to_w2 [K=3, H] (Xavier) - fan_dims[166] = (0, 0); // aux_to_b2 (zero) + let p_dim = crate::cuda_pipeline::gpu_aux_heads::PLAN_PARAM_DIM; + // SP22 H6 vNext Phase B5b (2026-05-14): Xavier fan_in is now + // SH2 + PLAN_PARAM_DIM = 262 (plan-conditioned input). The + // wider fan_in slightly reduces Xavier's `sqrt(6/(fan_in+ + // fan_out))` magnitude — fine, the head is still + // Linear→ELU→Linear. + fan_dims[163] = (h_dim, cfg.shared_h2 + p_dim); // aux_to_w1 [H, SH2+P=262] (Xavier) + fan_dims[164] = (0, 0); // aux_to_b1 (zero) + fan_dims[165] = (k_to, h_dim); // aux_to_w2 [K=3, H] (Xavier) + fan_dims[166] = (0, 0); // aux_to_b2 (zero) } // ── Phase 1 Task 1.6: MoE fan dims (36 tensors at [127..163)) ───────── diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index c1e91bd18..38d60f2fc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1297,6 +1297,27 @@ pub struct GpuExperienceCollector { /// consumer. exp_aux_to_label_per_sample: cudarc::driver::CudaSlice, + /// SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input + /// buffer for the collector-side K=3 forward. Sized + /// `[alloc_episodes × (SH2 + PLAN_PARAM_DIM) = N × 262]` f32 + /// row-major. + /// + /// **Train/inference asymmetry note**: trainer-side plan_params + /// come from the trade plan head's output (`plan_params_buf`). + /// Collector-side has NO trade plan head launch (the trade plan + /// runs only on the replay batch). The collector therefore passes + /// `plan_params_ptr = 0` (NULL) to the concat kernel, which fills + /// the trailing 6 columns with zeros. The K=3 forward at rollout + /// time evaluates "outcome prediction at plan_params=0", which + /// is a specific point — not the marginal-over-plans prediction. + /// + /// This mismatch means the head is trained on real plan_params + /// (replay) but queried at plan_params=0 (rollout's state bridge + /// to next-step state[121..124)). Phase B5b-2 (follow-up) adds a + /// trade plan head launch to the collector to populate real + /// plan_params at rollout time, resolving the asymmetry. + exp_aux_to_input_buf: cudarc::driver::CudaSlice, + /// SP22 H6 vNext Phase C (2026-05-14) — per-env K=3 trade-outcome /// softmax cache. Sized `[alloc_episodes × AUX_OUTCOME_K=3]` f32, /// row-major: env-major then class. Written end-of-step by @@ -2790,6 +2811,19 @@ impl GpuExperienceCollector { "sp22-vnext B4b: alloc exp_aux_to_label_per_sample ({} i32): {e}", alloc_episodes * alloc_timesteps )))?; + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned input + // buffer for the collector's K=3 forward. SH2_TOTAL = SH2 + + // PLAN_PARAM_DIM = 262. Populated by the concat kernel at the + // start of each rollout step's K=3 forward (plan_params source + // = NULL → zero-fill the trailing 6 cols; see field doc for + // train/inference asymmetry note). + let sh2_total_alloc = shared_h2 + super::gpu_aux_heads::PLAN_PARAM_DIM; + let exp_aux_to_input_buf = stream + .alloc_zeros::(alloc_episodes * sh2_total_alloc) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B5b: alloc exp_aux_to_input_buf ({} f32): {e}", + alloc_episodes * sh2_total_alloc + )))?; // SP22 H6 vNext Phase C (2026-05-14): K=3 per-env softmax cache // + producer kernel handle. Sized [alloc_episodes × AUX_OUTCOME_K=3] @@ -3405,6 +3439,11 @@ impl GpuExperienceCollector { exp_aux_to_softmax_buf, exp_aux_to_label_buf, exp_aux_to_label_per_sample, + // SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned + // input buffer [N × 262] for the collector-side K=3 + // forward. Populated by aux_to_fwd.launch_concat with + // plan_params=NULL (zero-fill the trailing 6 cols). + exp_aux_to_input_buf, // SP22 H6 vNext Phase C (2026-05-14): K=3 trade-outcome // per-env softmax cache + producer kernel handle (Phase C // scaffolding — populates `prev_aux_outcome_probs` every @@ -5943,12 +5982,39 @@ impl GpuExperienceCollector { 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( + // ── SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned ── + // input concat for the rollout-step K=3 forward. + // + // Collector has NO trade plan head launch (the trade plan + // runs only on replay batch); pass plan_params_ptr = 0 + // (NULL) to the concat kernel which zero-fills the + // trailing PLAN_PARAM_DIM=6 columns. Train/inference + // asymmetry documented in the `exp_aux_to_input_buf` + // struct field comment; Phase B5b-2 (follow-up) adds + // a real trade plan launch to the collector to resolve. + // + // SH2_TOTAL = SH2 + PLAN_PARAM_DIM = 262. Kernel takes + // SH2 as a runtime arg so passing sh2_total just works + // — no kernel surgery needed. + let p_dim = super::gpu_aux_heads::PLAN_PARAM_DIM; + let sh2_total = sh2 + p_dim; + let plan_params_null: u64 = 0; + let aux_to_input_ptr = self.exp_aux_to_input_buf.raw_ptr(); + self.exp_aux_to_fwd.launch_concat( &self.stream, self.exp_h_s2_aux.raw_ptr(), + plan_params_null, + n, + sh2, + p_dim, + aux_to_input_ptr, + )?; + self.exp_aux_to_fwd.forward( + &self.stream, + aux_to_input_ptr, to_w1, to_b1, to_w2, to_b2, n, // batch dim = n_episodes - sh2, + sh2_total, // SH2_TOTAL = 262 AUX_OUTCOME_K, self.exp_aux_to_hidden_buf.raw_ptr(), self.exp_aux_to_logits_buf.raw_ptr(), diff --git a/crates/ml/src/cuda_pipeline/strided_row_saxpy_kernel.cu b/crates/ml/src/cuda_pipeline/strided_row_saxpy_kernel.cu new file mode 100644 index 000000000..80b4d8dd7 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/strided_row_saxpy_kernel.cu @@ -0,0 +1,78 @@ +// crates/ml/src/cuda_pipeline/strided_row_saxpy_kernel.cu +// +// SP22 H6 vNext Phase B5b (2026-05-14) — strided row-truncating SAXPY. +// +// Adds the first `n_cols_copy` columns per row of a `[B, src_cols]` +// f32 tensor into a `[B, dst_cols]` f32 tensor, scaling by `alpha`. +// Required when `src_cols > dst_cols` (truncate trailing cols) OR when +// `src_cols < dst_cols` (zero-pad trailing — not used here; this kernel +// only supports `dst_cols ≤ src_cols`). +// +// Operation per row b ∈ [0, B): +// dst[b * dst_cols + j] += alpha × src[b * src_cols + j] for j ∈ [0, n_cols_copy) +// +// Element layout: row-major. The dst row stride (`dst_cols`) and src +// row stride (`src_cols`) can differ; only the first `n_cols_copy` +// columns of each row participate in the SAXPY. +// +// SP22 H6 vNext use case +// ────────────────────── +// Phase B5b's plan-conditioned aux head emits backward `dh_s2_aux_to +// _buf [B, SH2_TOTAL=262]` (gradient for both h_s2_aux and plan_params +// inputs). The aux trunk's upstream gradient accumulator +// `dh_s2_aux_accum [B, SH2=256]` only consumes the first 256 columns +// (the h_s2_aux gradient); the trailing 6 columns are the plan_params +// gradient which gets dropped (stop-grad — the trade plan head's +// weights are NOT trained by aux loss in this commit). +// +// The standard `dqn_saxpy_f32_kernel` is element-wise with a single +// stride, which would corrupt rows when src and dst strides differ. +// This kernel handles the stride mismatch by computing `(b, j)` from +// the flat thread index against the COPY dims (`B × n_cols_copy`) and +// reading/writing at the proper row strides. +// +// Discipline +// ────────── +// - Pure GPU map; no reduction, no atomicAdd +// - GPU-only per `feedback_cpu_is_read_only.md` +// +// Launch +// ────── +// grid = (ceil((B * n_cols_copy) / 256)), block = (256). One thread per +// (row, col) pair within the copy region. + +#include + +extern "C" __global__ void strided_row_saxpy_f32( + /* `[B, dst_cols]` row-major destination. Accumulated into. */ + float* __restrict__ dst, + /* `[B, src_cols]` row-major source. Read-only. */ + const float* __restrict__ src, + /* SAXPY scalar: `dst += alpha × src`. */ + float alpha, + /* Batch dim (number of rows). */ + int B, + /* Source row stride (columns per row in src). MUST be >= + * `n_cols_copy`; trailing cols are ignored. */ + int src_cols, + /* Destination row stride (columns per row in dst). MUST be >= + * `n_cols_copy`; trailing cols (beyond `n_cols_copy`) are + * untouched. */ + int dst_cols, + /* Number of leading columns to SAXPY per row. Must satisfy + * `n_cols_copy <= src_cols` AND `n_cols_copy <= dst_cols`. */ + int n_cols_copy +) { + const long long total = (long long)B * (long long)n_cols_copy; + long long idx = (long long)blockIdx.x * (long long)blockDim.x + + (long long)threadIdx.x; + if (idx >= total) return; + + /* Decompose flat idx → (b, j) against the COPY dims. */ + const int b = (int)(idx / (long long)n_cols_copy); + const int j = (int)(idx % (long long)n_cols_copy); + + /* Read at src row stride, write at dst row stride. */ + const float s = src[(long long)b * (long long)src_cols + (long long)j]; + dst[(long long)b * (long long)dst_cols + (long long)j] += alpha * s; +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 0c1d3edaa..fac9aef62 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -18101,3 +18101,53 @@ The K=3 head now closes the loop: it learns from real labels via the replay path - **Phase E**: dW backward + Adam for W[4, 3] — gradient for the 12 weights. - **Phase F**: Validation smoke at structural prior (W populated, β=0.5). - **B5b** (deferred): Plan-conditioning input concat — refinement, not on critical path. + +#### Phase B5b — Full plan-conditioning integration (2026-05-14) + +The K=3 head's input is now [h_s2_aux (256) || plan_params (6)] = 262-dim, matching the spec's intended architecture. Implements the input concat in both the trainer's replay-batch path and the collector's rollout-step path, with appropriate handling of the backward-side stride mismatch. + +**NEW kernel `strided_row_saxpy_kernel.cu`**: row-truncating SAXPY that accumulates first `n_cols_copy` columns per row of `src [B, src_cols]` into `dst [B, dst_cols]` scaled by `alpha`. Handles stride mismatch (`src_cols ≠ dst_cols`). Needed because backward emits `dh_s2_aux_to_buf [B, 262]` but `dh_s2_aux_accum [B, 256]` only consumes the first 256 cols. ~30 lines, cubin 5.3 KB. + +**`PLAN_PARAM_DIM = 6` constant** in `gpu_aux_heads.rs`: the 6 trade plan parameters per the trade plan head's ABI (`{target_bars, profit_target, stop_loss, conviction, asymmetry, trail_distance}`). + +**Ops struct updates** (`gpu_aux_heads.rs`): +- `AuxTradeOutcomeForwardOps` gains `concat_kernel` handle + `launch_concat()` method. +- `AuxTradeOutcomeBackwardOps` gains `strided_saxpy_kernel` handle + `launch_strided_row_saxpy()` method. +- Both ops structs load their new cubins in `new()` (`AUX_TO_INPUT_CONCAT_CUBIN`, `STRIDED_ROW_SAXPY_CUBIN`). + +**Weight tensor resize** (`compute_param_sizes` + `xavier_init_params_buf`): +- `sizes[163] = H * (SH2 + PLAN_PARAM_DIM) = 128 × 262 = 33,536 floats` (was `H × SH2 = 32,768`). Adam SAXPY iterates per-element so the larger flat size updates uniformly. +- `fan_dims[163] = (H, SH2 + PLAN_PARAM_DIM)` — Xavier sees wider fan_in. + +**Trainer struct changes** (`gpu_dqn_trainer.rs`): +- New `aux_to_input_buf: CudaSlice` sized `[B × SH2_TOTAL=262]`. Populated by `aux_to_fwd.launch_concat(h_s2_aux, plan_params_buf, ...)` before each forward call. +- `aux_dh_s2_to_buf` resized to `[B × SH2_TOTAL=262]` (backward output stride). +- `aux_partial_to_w1` resized to `[B × H × SH2_TOTAL=262]` (dW1 partial dim). +- `max_aux_tensor_len` bumped to include the new W1 size for the param_grad_final scratch buffer. + +**Trainer forward path** (`aux_heads_forward` in `gpu_dqn_trainer.rs`): +- Before the K=3 `.forward()` call: launch concat with real `plan_params_buf` (trainer-side has the trade plan head's output). +- `.forward()` called with `SH2_TOTAL=262` instead of `SH2=256`. + +**Trainer backward path** (`aux_heads_backward`): +- `.backward()` called with `SH2_TOTAL=262`; reads `aux_to_input_buf` instead of `h_s2_aux_ptr` directly. +- The legacy `saxpy_f32_kernel` SAXPY for `aux_dh_s2_to_buf` is **replaced** by `aux_to_bwd.launch_strided_row_saxpy(...)` which copies only the first `SH2=256` cols per row into `dh_s2_aux_accum`. The trailing 6 cols (plan_params gradient) are dropped — stop-grad. Phase E would add trade plan training via aux loss if needed. + +**Collector forward path** (rollout loop in `gpu_experience_collector.rs`): +- New `exp_aux_to_input_buf [N × 262]` field on the collector. +- Before the K=3 `.forward()` call: launch concat with `plan_params_ptr = 0` (NULL) → kernel zero-fills the trailing 6 cols. +- `.forward()` called with `SH2_TOTAL=262`. + +**Train/inference asymmetry** (documented): +- Trainer: input = `[h_s2_aux || plan_params]` where plan_params comes from the trade plan head's output on the replay batch. +- Collector: input = `[h_s2_aux || ZEROS]` — plan_params are zero-filled because the collector has no trade plan head launch. +- Consequence: the head is trained on real plan-conditional outcome predictions but queried at rollout time with plan_params=0. The K=3 state-bridge prediction fed to Phase C state slots [121..124) is therefore "outcome at plan_params=0", a specific evaluation rather than the plan-marginal prediction. +- **Phase B5b-2 (follow-up)** would add a trade plan head launch to the collector to populate real `plan_params` at rollout time, fully resolving the asymmetry. Deferred — current asymmetry is functional (no NaN, no kernel crash) and the head still receives plan signal during training. + +**Stop-grad on plan_params**: backward writes a full `[B, 262]` gradient tensor, but the strided SAXPY only copies the first 256 cols into the aux trunk accumulator. The trailing 6 cols (plan_params gradient) are dropped. This means the trade plan head's weights are NOT trained by the K=3 aux loss in this commit. Phase E or a separate follow-up could plumb that gradient back through if the K=3 head's training signal benefits from co-adapting the trade plan head. + +**Verification**: +- `cargo check -p ml` clean (21 warnings, none new). +- `cargo test -p ml --lib` → **1016/0 green**. + +**Phase D next**: 12-weight W atom-shift (4 actions × 3 outcomes). Extends the existing SP22 Phase 3 atom-shift mechanism from 4 weights (per-action) to 12 weights (per-action × per-outcome). The K=3 head's softmax probs (now plan-conditioned at training time) modulate Q-target atom positions per action.