From c3ee5e165ad8f1d0041d0e8f7ff7d54cf81c7b8f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 10:42:35 +0200 Subject: [PATCH] feat(ml-alpha): wire batched kernels through trainer + CLI (#8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plumbs the batched cfc + heads kernels added in 829ddfa62 through PerceptionTrainer, evaluator, and the alpha_train CLI: PerceptionTrainerConfig.n_batch — batch size, default 1 PerceptionTrainer::step_batched — process B sequences per optimizer step using cfc_step_batched / heads_batched PerceptionTrainer::evaluate_batched — forward-only batched eval PerceptionTrainer::step / evaluate — thin B=1 wrappers preserving existing single-sequence test/inference APIs (assert cfg.n_batch == 1) alpha_train CLI: --batch-size N — accumulates B sequences per optimizer step in train loop; val loop also batches and uses evaluate_batched Per-K scratch buffers all grow to [K, B, dim] layout (K-major, slot-k contiguous). Mamba2's [B, K, H] output is transposed once after forward via the new transpose_3d_swap_01 kernel, and grad_h_enriched_seq_t is transposed back to [B, K, H] before Mamba2 backward. Two transposes per training step; negligible (1.5MB at B=32). Dead unbatched kernel handles removed from the trainer (step_fn, step_bwd_fn, heads_fn, heads_bwd_fn, grad_x_d) — all training and inference now go through the batched variants for B ≥ 1. The single-sample kernels remain in CUDA for the standalone test helpers in cfc/step.rs and heads.rs. Local 2Q smoke (seq_len=32, B=4, --auto-horizon-weights, 800 train seqs × 2 epochs): epoch 0: val_loss=0.7138 AUC h30/h100/h300/h1000/h6000 = .55/.55/.57/.61/.51 epoch 1: val_loss=0.6558 AUC h30/h100/h300/h1000/h6000 = .72/.68/.75/.68/.65 vs the in-flight qf5mj baseline (B=1, K=96, no horizon weighting) which had val_loss=0.6933 best and AUCs oscillating at ~0.50 — this batched run hits AUC 0.75 (h300) and 0.72 (h30) in just 2 epochs of 200 optimizer updates. Batching + horizon-weighting unblocks the model. 77 ml-alpha tests pass. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/examples/alpha_train.rs | 73 ++- crates/ml-alpha/src/trainer/perception.rs | 599 ++++++++++++-------- crates/ml-alpha/tests/perception_overfit.rs | 1 + 3 files changed, 425 insertions(+), 248 deletions(-) diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index 75a2c72c7..7c220decc 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -23,6 +23,7 @@ use anyhow::{Context, Result}; use clap::Parser; +use ml_alpha::cfc::snap_features::Mbp10RawInput; use ml_alpha::data::loader::{MultiHorizonLoader, MultiHorizonLoaderConfig}; use ml_alpha::eval::auc::{compute_auc, AucInput}; use ml_alpha::heads::N_HORIZONS; @@ -96,6 +97,12 @@ struct Cli { /// horizons include values larger than seq_len. #[arg(long, default_value_t = false)] auto_horizon_weights: bool, + + /// Mini-batch size for training. Each step processes B sequences + /// in parallel via batched CUDA kernels. Default 1 matches the + /// previous unbatched behavior. Validation always runs at B=1. + #[arg(long, default_value_t = 1)] + batch_size: usize, } #[derive(Serialize)] @@ -172,6 +179,7 @@ fn main() -> Result<()> { "per-horizon BCE weights" ); + anyhow::ensure!(cli.batch_size >= 1, "batch_size must be >= 1"); let trainer_cfg = PerceptionTrainerConfig { seq_len: cli.seq_len, mamba2_state_dim: cli.mamba2_state_dim, @@ -179,6 +187,7 @@ fn main() -> Result<()> { lr_mamba2: cli.lr_mamba2, seed: cli.seed, horizon_weights, + n_batch: cli.batch_size, }; let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?; @@ -215,10 +224,12 @@ fn main() -> Result<()> { let mut epoch_train_loss = 0.0_f32; let mut epoch_train_steps = 0usize; + // Accumulate B sequences per optimizer step. Sequences with no + // finite labels at any position are skipped (continue) so they + // don't pollute the batch. + let mut snap_batch: Vec> = Vec::with_capacity(cli.batch_size); + let mut label_batch: Vec> = Vec::with_capacity(cli.batch_size); while let Some(seq) = train_loader.next_sequence().context("train next_seq")? { - // Per-position labels: position k has horizon-h label - // seq.labels[h][k]. NaN entries (right-edge / tied price) - // are passed through to the trainer which masks them. let mut labels_per_pos: Vec<[f32; N_HORIZONS]> = Vec::with_capacity(seq.snapshots.len()); let mut any_finite = false; for k in 0..seq.snapshots.len() { @@ -230,17 +241,25 @@ fn main() -> Result<()> { labels_per_pos.push(row); } if !any_finite { continue; } - // Apply LR schedule before each step. + snap_batch.push(seq.snapshots); + label_batch.push(labels_per_pos); + if snap_batch.len() < cli.batch_size { continue; } + + // Full batch ready — apply LR schedule, fire step. let lr_cfc_now = lr_schedule(cli.lr_cfc, lr_min_cfc, train_steps, cli.lr_warmup_steps, total_steps_budget); let lr_m2_now = lr_schedule(cli.lr_mamba2, lr_min_m2, train_steps, cli.lr_warmup_steps, total_steps_budget); trainer.set_lr_cfc(lr_cfc_now); trainer.set_lr_mamba2(lr_m2_now); - let loss = trainer.step(&seq.snapshots, &labels_per_pos).context("train step")?; + let snap_refs: Vec<&[Mbp10RawInput]> = snap_batch.iter().map(|v| v.as_slice()).collect(); + let label_refs: Vec<&[[f32; N_HORIZONS]]> = label_batch.iter().map(|v| v.as_slice()).collect(); + let loss = trainer.step_batched(&snap_refs, &label_refs).context("train step_batched")?; epoch_train_loss += loss; epoch_train_steps += 1; train_loss_running += loss; train_steps += 1; + snap_batch.clear(); + label_batch.clear(); } let epoch_avg = if epoch_train_steps > 0 { epoch_train_loss / epoch_train_steps as f32 @@ -262,6 +281,9 @@ fn main() -> Result<()> { let mut val_labels: [Vec; N_HORIZONS] = Default::default(); let mut val_loss_sum = 0.0_f32; let mut val_steps = 0usize; + let mut val_snap_batch: Vec> = Vec::with_capacity(cli.batch_size); + let mut val_label_batch: Vec> = Vec::with_capacity(cli.batch_size); + let mut val_last_label_batch: Vec<[f32; N_HORIZONS]> = Vec::with_capacity(cli.batch_size); while let Some(seq) = val_loader.next_sequence().context("val next_seq")? { let mut labels_per_pos: Vec<[f32; N_HORIZONS]> = Vec::with_capacity(seq.snapshots.len()); let mut any_finite = false; @@ -274,23 +296,38 @@ fn main() -> Result<()> { labels_per_pos.push(row); } if !any_finite { continue; } - // FORWARD-ONLY: evaluate() skips backward + AdamW (real bug - // fix from previous CLI which called step() and leaked val - // gradient into trained weights). - let (l, probs_all) = trainer.evaluate(&seq.snapshots, &labels_per_pos).context("val eval")?; + let last = seq.snapshots.len().saturating_sub(1); + let last_labels = [ + seq.labels[0][last], seq.labels[1][last], seq.labels[2][last], + seq.labels[3][last], seq.labels[4][last], + ]; + val_snap_batch.push(seq.snapshots); + val_label_batch.push(labels_per_pos); + val_last_label_batch.push(last_labels); + if val_snap_batch.len() < cli.batch_size { continue; } + + // FORWARD-ONLY batched evaluation — no backward, no AdamW. + let snap_refs: Vec<&[Mbp10RawInput]> = val_snap_batch.iter().map(|v| v.as_slice()).collect(); + let label_refs: Vec<&[[f32; N_HORIZONS]]> = val_label_batch.iter().map(|v| v.as_slice()).collect(); + let (l, probs_all) = trainer.evaluate_batched(&snap_refs, &label_refs).context("val eval_batched")?; val_loss_sum += l; val_steps += 1; - // Score AUC from the LAST-position predictions (matches the - // horizon-aligned label at that position). - let last = seq.snapshots.len().saturating_sub(1); - let last_probs = &probs_all[last * N_HORIZONS..(last + 1) * N_HORIZONS]; - for h in 0..N_HORIZONS { - let lbl = seq.labels[h][last]; - if lbl.is_finite() { - val_probs[h].push(last_probs[h]); - val_labels[h].push(lbl); + // probs_all is [K, B, N_HORIZONS] row-major. Score AUC from + // the LAST-position predictions for each sample in the batch. + let last = cli.seq_len - 1; + for (b_idx, last_lbl) in val_last_label_batch.iter().enumerate() { + let off = (last * cli.batch_size + b_idx) * N_HORIZONS; + let last_probs = &probs_all[off..off + N_HORIZONS]; + for h in 0..N_HORIZONS { + if last_lbl[h].is_finite() { + val_probs[h].push(last_probs[h]); + val_labels[h].push(last_lbl[h]); + } } } + val_snap_batch.clear(); + val_label_batch.clear(); + val_last_label_batch.clear(); } let mut per_horizon_auc = [0.5_f32; N_HORIZONS]; for h in 0..N_HORIZONS { diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 3f6761c4c..d2f144c0b 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -70,6 +70,11 @@ pub struct PerceptionTrainerConfig { /// → near-identical labels would otherwise inflate gradient /// pressure). All-ones gives the original uniform-mean BCE. pub horizon_weights: [f32; N_HORIZONS], + /// Number of sequences processed per optimizer step. Sets the + /// effective batch size — kernels switch to batched variants + /// (cfc_step_batched, multi_horizon_heads_batched, etc.). + /// `n_batch=1` matches the unbatched behavior exactly. + pub n_batch: usize, } impl Default for PerceptionTrainerConfig { @@ -81,6 +86,7 @@ impl Default for PerceptionTrainerConfig { lr_mamba2: 1e-3, seed: 0x4242, horizon_weights: [1.0; N_HORIZONS], + n_batch: 1, } } } @@ -108,11 +114,16 @@ pub struct PerceptionTrainer { _heads_module: Arc, _bce_module: Arc, snap_fn: CudaFunction, - step_fn: CudaFunction, - step_bwd_fn: CudaFunction, - heads_fn: CudaFunction, - heads_bwd_fn: CudaFunction, bce_fn: CudaFunction, + /// Batched cfc/heads forward + backward kernels — used by + /// `step_batched()` for all training updates regardless of B. + /// At cfg.n_batch == 1 they produce identical results to the + /// unbatched kernels (the inner per-thread loop runs once). + step_batched_fn: CudaFunction, + step_bwd_batched_fn: CudaFunction, + heads_batched_fn: CudaFunction, + heads_bwd_batched_fn: CudaFunction, + transpose_3d_fn: CudaFunction, // Mamba2 encoder block + its optimizer pub mamba2: Mamba2Block, @@ -171,9 +182,6 @@ pub struct PerceptionTrainer { /// Single-position grad_h_new buffer fed to cfc_step_bwd (heads_bwd /// writes to it, fused with the carry). grad_h_new_d: CudaSlice, - /// Single-position grad_x output of cfc_step_bwd; DtoD copied into - /// the slot-k position of grad_h_enriched_seq each iteration. - grad_x_d: CudaSlice, /// Single-position zero h_old buffer (used only at k=0 — first /// position of the sequence has no preceding state). zero_h_d: CudaSlice, @@ -207,11 +215,12 @@ impl PerceptionTrainer { let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("heads cubin")?; let bce_module = ctx.load_cubin(BCE_CUBIN.to_vec()).context("bce cubin")?; let snap_fn = snap_module.load_function("snap_feature_assemble")?; - let step_fn = step_module.load_function("cfc_step")?; - let step_bwd_fn = step_module.load_function("cfc_step_backward")?; - let heads_fn = heads_module.load_function("multi_horizon_heads")?; - let heads_bwd_fn = heads_module.load_function("multi_horizon_heads_backward")?; let bce_fn = bce_module.load_function("bce_multi_horizon_forward_backward")?; + let step_batched_fn = step_module.load_function("cfc_step_batched")?; + let step_bwd_batched_fn = step_module.load_function("cfc_step_backward_batched")?; + let heads_batched_fn = heads_module.load_function("multi_horizon_heads_batched")?; + let heads_bwd_batched_fn = heads_module.load_function("multi_horizon_heads_backward_batched")?; + let transpose_3d_fn = step_module.load_function("transpose_3d_swap_01")?; // Mamba2 encoder: in_dim=FEATURE_DIM (32), hidden=HIDDEN_DIM (128). let mamba2 = Mamba2Block::new( @@ -283,17 +292,16 @@ impl PerceptionTrainer { prev_ask_sz_d: stream.alloc_zeros::(10)?, regime_d: stream.alloc_zeros::(REGIME_DIM)?, snap_feat_d: stream.alloc_zeros::(FEATURE_DIM)?, - h_new_per_k_d: stream.alloc_zeros::(k * n_hid)?, - probs_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, - labels_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, - grad_probs_per_k_d: stream.alloc_zeros::(k * N_HORIZONS)?, + h_new_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * n_hid)?, + probs_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * N_HORIZONS)?, + labels_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * N_HORIZONS)?, + grad_probs_per_k_d: stream.alloc_zeros::(k * cfg.n_batch * N_HORIZONS)?, loss_weights_d: upload(&stream, &cfg.horizon_weights)?, loss_d: stream.alloc_zeros::(1)?, valid_d: stream.alloc_zeros::(1)?, - grad_h_carry_d: stream.alloc_zeros::(n_hid)?, - grad_h_new_d: stream.alloc_zeros::(n_hid)?, - grad_x_d: stream.alloc_zeros::(n_in)?, - zero_h_d: stream.alloc_zeros::(n_hid)?, + grad_h_carry_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, + grad_h_new_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, + zero_h_d: stream.alloc_zeros::(cfg.n_batch * n_hid)?, grad_w_in_d: stream.alloc_zeros::(n_hid * n_in)?, grad_w_rec_d: stream.alloc_zeros::(n_hid * n_hid)?, grad_b_d: stream.alloc_zeros::(n_hid)?, @@ -305,7 +313,7 @@ impl PerceptionTrainer { stg_ask_px: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, stg_ask_sz: unsafe { MappedF32Buffer::new(10) }.map_err(|e| anyhow::anyhow!("stg: {e}"))?, stg_regime: unsafe { MappedF32Buffer::new(REGIME_DIM) }.map_err(|e| anyhow::anyhow!("stg_regime: {e}"))?, - stg_labels: unsafe { MappedF32Buffer::new(k * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?, + stg_labels: unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?, w_in_d, w_rec_d, b_d, @@ -318,11 +326,12 @@ impl PerceptionTrainer { _heads_module: heads_module, _bce_module: bce_module, snap_fn, - step_fn, - step_bwd_fn, - heads_fn, - heads_bwd_fn, bce_fn, + step_batched_fn, + step_bwd_batched_fn, + heads_batched_fn, + heads_bwd_batched_fn, + transpose_3d_fn, mamba2, mamba2_adamw, opt_w_in, @@ -350,43 +359,69 @@ impl PerceptionTrainer { self.mamba2_adamw.config.lr = lr; } - /// One training step. CfC is RECURRENT across positions in the - /// sequence: h_old at step k = h_new at step k-1 (h_old at k=0 is - /// zero). Forward + backward are both fully GPU-resident — the K - /// loop emits stream-ordered launches with NO host roundtrips and - /// only two `stream.synchronize()` calls (after forward, after - /// backward). - /// - /// Per-step gradients accumulate in shared `grad_*` buffers via - /// kernel `+=` semantics. tau is now trained (its own AdamW). The - /// fused multi-horizon BCE kernel handles NaN-label masking - /// natively + emits the (loss, n_valid, grad_probs) triple in one - /// launch over the full [K, N_HORIZONS] matrix. - /// - /// Returns mean BCE over the valid (position × horizon) entries. + /// One training step on a single sequence — thin wrapper around + /// [`step_batched`] with batch size 1. Preserves the existing + /// API for single-sequence callers (test smokes, etc.). pub fn step( &mut self, snapshots: &[Mbp10RawInput], labels_per_position: &[[f32; N_HORIZONS]], ) -> Result { + anyhow::ensure!( + self.cfg.n_batch == 1, + "single-sequence step() requires cfg.n_batch == 1; got {}. \ + Use step_batched() for batched training.", + self.cfg.n_batch + ); + self.step_batched(&[snapshots], &[labels_per_position]) + } + + /// Batched training step. Processes `cfg.n_batch` sequences per + /// optimizer update — uses the batched kernel variants (cfc_step_batched, + /// multi_horizon_heads_batched, etc.). CfC is RECURRENT across + /// positions; h_old at step k for sample b is h_new at step k-1 for + /// the same b (zero at k=0). + /// + /// `snapshots_batch[b]` is sample b's seq_len-snapshot window. + /// `labels_batch[b]` is sample b's per-position label rows. NaN + /// labels are natively masked by the BCE kernel. + /// + /// Returns mean BCE over the valid (b, position, horizon) triples + /// in the batch, weighted by `cfg.horizon_weights`. + pub fn step_batched( + &mut self, + snapshots_batch: &[&[Mbp10RawInput]], + labels_batch: &[&[[f32; N_HORIZONS]]], + ) -> Result { + let b_sz = self.cfg.n_batch; let k_seq = self.cfg.seq_len; anyhow::ensure!( - snapshots.len() == k_seq, - "snapshots.len()={} != seq_len={}", - snapshots.len(), k_seq - ); - anyhow::ensure!( - labels_per_position.len() == k_seq, - "labels_per_position.len()={} != seq_len={}", - labels_per_position.len(), k_seq + snapshots_batch.len() == b_sz && labels_batch.len() == b_sz, + "expected B={} sequences/labels; got snap={} lbl={}", + b_sz, snapshots_batch.len(), labels_batch.len() ); + for (b, snap) in snapshots_batch.iter().enumerate() { + anyhow::ensure!( + snap.len() == k_seq, + "batch[{}] snapshots.len()={} != seq_len={}", + b, snap.len(), k_seq + ); + } + for (b, lbl) in labels_batch.iter().enumerate() { + anyhow::ensure!( + lbl.len() == k_seq, + "batch[{}] labels.len()={} != seq_len={}", + b, lbl.len(), k_seq + ); + } - // ── 1. Build the window tensor by running snap_feature_assemble - // per snapshot and packing into [1, K, FEATURE_DIM]. The - // K snap launches are stream-ordered; single sync at the - // end since Mamba2's forward depends on the full window. + // ── 1. Build the batched window tensor [B, K, FEATURE_DIM] by + // running snap_feature_assemble per (b, k) snapshot and + // packing into the right offset. K * B snap launches are + // stream-ordered; single sync at the end since Mamba2's + // forward depends on the full window. let mut window_tensor = GpuTensor::zeros( - &[1, k_seq, FEATURE_DIM], + &[b_sz, k_seq, FEATURE_DIM], &self.stream, ).map_err(|e| anyhow::anyhow!("window alloc: {e}"))?; @@ -394,57 +429,87 @@ impl PerceptionTrainer { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0, }; let snap_nbytes = FEATURE_DIM * std::mem::size_of::(); - for (k, snap) in snapshots.iter().enumerate() { - upload_into(&self.stream, &snap.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?; - upload_into(&self.stream, &snap.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?; - upload_into(&self.stream, &snap.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?; - upload_into(&self.stream, &snap.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?; - upload_into(&self.stream, &snap.regime, &self.stg_regime, &mut self.regime_d)?; - let prev_mid = snap.prev_mid; - let trade_signed_vol = snap.trade_signed_vol; - let trade_count = snap.trade_count as i32; - let ts_ns = snap.ts_ns as i64; - let prev_ts_ns = snap.prev_ts_ns as i64; - let tick_size = ES_TICK_SIZE; - { - let mut launch = self.stream.launch_builder(&self.snap_fn); - launch - .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d) - .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d) - .arg(&self.regime_d) - .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count) - .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size) - .arg(&mut self.snap_feat_d); - unsafe { launch.launch(snap_cfg).context("snap fwd")?; } - } - unsafe { - let (src_ptr, _g1) = self.snap_feat_d.device_ptr(&self.stream); - let (dst_base, _g2) = window_tensor.data_mut().device_ptr_mut(&self.stream); - let dst_offset_ptr = dst_base + (k * snap_nbytes) as u64; - cudarc::driver::result::memcpy_dtod_async( - dst_offset_ptr, src_ptr, snap_nbytes, self.stream.cu_stream(), - ).context("window pack DtoD")?; + for (b_idx, snapshots) in snapshots_batch.iter().enumerate() { + for (k, snap) in snapshots.iter().enumerate() { + upload_into(&self.stream, &snap.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?; + upload_into(&self.stream, &snap.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?; + upload_into(&self.stream, &snap.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?; + upload_into(&self.stream, &snap.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?; + upload_into(&self.stream, &snap.regime, &self.stg_regime, &mut self.regime_d)?; + let prev_mid = snap.prev_mid; + let trade_signed_vol = snap.trade_signed_vol; + let trade_count = snap.trade_count as i32; + let ts_ns = snap.ts_ns as i64; + let prev_ts_ns = snap.prev_ts_ns as i64; + let tick_size = ES_TICK_SIZE; + { + let mut launch = self.stream.launch_builder(&self.snap_fn); + launch + .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d) + .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d) + .arg(&self.regime_d) + .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count) + .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size) + .arg(&mut self.snap_feat_d); + unsafe { launch.launch(snap_cfg).context("snap fwd")?; } + } + unsafe { + let (src_ptr, _g1) = self.snap_feat_d.device_ptr(&self.stream); + let (dst_base, _g2) = window_tensor.data_mut().device_ptr_mut(&self.stream); + // window_tensor layout [B, K, F] row-major. + let dst_offset_ptr = dst_base + + ((b_idx * k_seq + k) * snap_nbytes) as u64; + cudarc::driver::result::memcpy_dtod_async( + dst_offset_ptr, src_ptr, snap_nbytes, self.stream.cu_stream(), + ).context("window pack DtoD")?; + } } } self.stream.synchronize().context("window pack sync")?; - // ── 2. Mamba2 per-step forward → h_enriched_seq [1, K, HIDDEN_DIM]. + // ── 2. Mamba2 per-step forward → h_enriched_seq [B, K, HIDDEN_DIM]. let (h_enriched_seq, cache) = self .mamba2 .forward_train_seq(&window_tensor) .context("mamba2 forward_train_seq")?; - // ── 3. Upload labels [K, N_HORIZONS] and pre-zero accumulators. - // Labels go up once per sequence via a single DtoD copy - // from the mapped-pinned staging buffer. - let total_labels = k_seq * N_HORIZONS; - // Pack labels row-major into the staging buffer's host-side - // slice (mapped-pinned — write is visible to the GPU's DtoD). + // ── 2b. Transpose [B, K, H] → [K, B, H] so the K-loop can + // slice contiguous [B, H] chunks per step. + let mut h_enriched_seq_t = GpuTensor::zeros( + &[k_seq, b_sz, HIDDEN_DIM], + &self.stream, + ).map_err(|e| anyhow::anyhow!("h_enriched_seq_t alloc: {e}"))?; + { + let block_n3: u32 = 32; + let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3); + let cfg_tx = LaunchConfig { + grid_dim: (b_sz as u32, k_seq as u32, grid_z), + block_dim: (block_n3, 1, 1), + shared_mem_bytes: 0, + }; + let n1 = b_sz as i32; + let n2 = k_seq as i32; + let n3 = HIDDEN_DIM as i32; + let mut launch = self.stream.launch_builder(&self.transpose_3d_fn); + launch + .arg(h_enriched_seq.cuda_data()) + .arg(h_enriched_seq_t.data_mut()) + .arg(&n1).arg(&n2).arg(&n3); + unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; } + } + + // ── 3. Pack labels in [K, B, N_HORIZONS] layout (matches the + // per-K slot layout of probs_per_k / labels_per_k) and + // upload via mapped-pinned staging. + let total_labels = k_seq * b_sz * N_HORIZONS; { let dst = self.stg_labels.host_slice_mut(); for k in 0..k_seq { - for h in 0..N_HORIZONS { - dst[k * N_HORIZONS + h] = labels_per_position[k][h]; + for b_idx in 0..b_sz { + for h in 0..N_HORIZONS { + dst[(k * b_sz + b_idx) * N_HORIZONS + h] = + labels_batch[b_idx][k][h]; + } } } } @@ -472,14 +537,6 @@ impl PerceptionTrainer { self.stream.memset_zeros(&mut self.zero_h_d) .map_err(|e| anyhow::anyhow!("zero zero_h: {e}"))?; - // grad_h_enriched_seq output buffer for Mamba2 backward — fresh - // [1, K, HIDDEN_DIM] tensor, written slot-by-slot in the - // reverse-order K backward loop below. - let mut grad_h_enriched_seq = GpuTensor::zeros( - &[1, k_seq, HIDDEN_DIM], - &self.stream, - ).map_err(|e| anyhow::anyhow!("grad_h_enriched_seq alloc: {e}"))?; - // Launch configs reused inside the K loop. let dt_s = 1.0_f32; let n_in_i = HIDDEN_DIM as i32; @@ -496,71 +553,87 @@ impl PerceptionTrainer { block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0, }; - let cfc_bwd_smem = (2 * HIDDEN_DIM * std::mem::size_of::()) as u32; + // Batched CfC backward shared mem: sd_pre[B, n_hid] + sdecay[n_hid]. + let cfc_bwd_smem = ((b_sz + 1) * HIDDEN_DIM * std::mem::size_of::()) as u32; let cfg_cfc_bwd = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: cfc_bwd_smem, }; + // Batched heads backward shared mem: sd_z[B, 5]. + let heads_bwd_smem = (b_sz * N_HORIZONS * std::mem::size_of::()) as u32; + let cfg_heads_bwd = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: heads_bwd_smem, + }; + // Heads forward uses block of 128 too (5 threads compute heads, + // remaining are idle — same kernel layout as backward). + let cfg_heads_batched_fwd = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + }; + let _ = cfg_heads; - let hid_bytes = HIDDEN_DIM * std::mem::size_of::(); - let nh_bytes = N_HORIZONS * std::mem::size_of::(); + // Per-K slot strides for [K, B, H] / [K, B, N_HORIZONS] layout. + let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::(); + let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::(); + let n_batch_i = b_sz as i32; // ── 4. Stream-ordered forward K loop. CfC is RECURRENT — - // h_old at step k IS h_new at step k-1 (slice of - // h_new_per_k_d). For k=0 we pass zero_h_d. No syncs. - // - // Pointer-offset trick: we hold only a SINGLE mut borrow of each - // per-K buffer, derive read-only slot offsets from the same base. - // Kernels see raw device pointers; Rust's borrow checker doesn't - // follow into CUDA-side memory aliasing. + // h_old at step k for sample b IS h_new at step k-1 for + // the same b. Per-K slot is contiguous [B, H] in the + // [K, B, H] layout. Pointer-offset trick (single mut + // borrow + raw u64 arithmetic for slot pointers) keeps + // the launches stream-ordered with no syncs. let zero_h_ptr = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; - let henr_base = { - let (p, _g) = h_enriched_seq.cuda_data().device_ptr(&self.stream); + let henr_t_base = { + let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream); p }; let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream); let (probs_base, _g_probs) = self.probs_per_k_d.device_ptr_mut(&self.stream); for k in 0..k_seq { - let h_new_k_ptr = h_per_k_base + (k * hid_bytes) as u64; + let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64; let h_old_k_ptr = if k == 0 { zero_h_ptr } else { - h_per_k_base + ((k - 1) * hid_bytes) as u64 + h_per_k_base + ((k - 1) * kb_hid_bytes) as u64 }; - let x_k_ptr = henr_base + (k * hid_bytes) as u64; - let probs_k_ptr = probs_base + (k * nh_bytes) as u64; + let x_k_ptr = henr_t_base + (k * kb_hid_bytes) as u64; + let probs_k_ptr = probs_base + (k * kb_nh_bytes) as u64; - // CfC forward at position k: h_new_k = cfc(x_k, h_old_k). + // cfc_step_batched(W, x[B*n_in], h_old[B*n_hid], …, → h_new[B*n_hid]). unsafe { - let mut launch = self.stream.launch_builder(&self.step_fn); + let mut launch = self.stream.launch_builder(&self.step_batched_fn); launch .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr) - .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&h_new_k_ptr); - launch.launch(cfg_cfc).context("cfc fwd k")?; + launch.launch(cfg_cfc).context("cfc fwd k batched")?; } - // Heads forward at position k: probs_k = heads(h_new_k). + // multi_horizon_heads_batched(W, b, h[B*128], B, → probs[B*5]). unsafe { - let mut launch = self.stream.launch_builder(&self.heads_fn); + let mut launch = self.stream.launch_builder(&self.heads_batched_fn); launch .arg(&self.heads_w_d).arg(&self.heads_b_d) - .arg(&h_new_k_ptr).arg(&probs_k_ptr); - launch.launch(cfg_heads).context("heads fwd k")?; + .arg(&h_new_k_ptr).arg(&n_batch_i).arg(&probs_k_ptr); + launch.launch(cfg_heads_batched_fwd).context("heads fwd k batched")?; } } drop((_g_hpk, _g_probs)); - // ── 5. Fused multi-horizon BCE on the full [K, N_HORIZONS] - // grid — single launch, native NaN masking, emits loss - // + grad_probs per position. No host roundtrip. - let n_pos_i = k_seq as i32; + // ── 5. Fused multi-horizon BCE over the full [K*B, N_HORIZONS] + // grid. The kernel doesn't distinguish position-vs-batch + // since the per-horizon weight is identified by `i % n_horizons`. + let n_pos_i = (k_seq * b_sz) as i32; let n_h_i = N_HORIZONS as i32; let bce_cfg = LaunchConfig { grid_dim: (1, 1, 1), @@ -578,16 +651,16 @@ impl PerceptionTrainer { unsafe { launch.launch(bce_cfg).context("bce launch")?; } } - // ── 6. Reverse-order backward K loop with grad_h_old carry. - // At each k (K-1 → 0): - // a) heads_bwd reads probs_k + grad_probs_k + carry, - // writes grad_h_new_d = sum_heads + carry, accumulates - // grad_heads_w/b (+=). - // b) cfc_step_bwd reads x_k, h_old_k, grad_h_new_d, - // writes grad_h_carry_d (= grad_h_old_out for - // propagation to position k-1), accumulates - // grad_w_in/w_rec/b/tau (+=), emits grad_x_k. - // c) grad_x_k DtoD into slot k of grad_h_enriched_seq. + // ── 6. Reverse-order backward K loop. Same recurrence carry as + // the unbatched version, but every kernel is batched over + // B samples. grad_h_enriched_seq_t is [K, B, H] — slot k + // contiguous, written by the cfc_step_bwd batched kernel + // directly into its grad_x output (which IS sized [B, n_in]). + let mut grad_h_enriched_seq_t = GpuTensor::zeros( + &[k_seq, b_sz, HIDDEN_DIM], + &self.stream, + ).map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_t alloc: {e}"))?; + self.stream.memset_zeros(&mut self.grad_h_carry_d) .map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?; @@ -595,70 +668,90 @@ impl PerceptionTrainer { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; - let henr_base_bwd = { - let (p, _g) = h_enriched_seq.cuda_data().device_ptr(&self.stream); + let henr_t_base_bwd = { + let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream); p }; let (h_per_k_base_bwd, _g_hpk_bwd) = self.h_new_per_k_d.device_ptr_mut(&self.stream); let (probs_base_bwd, _g_probs_bwd) = self.probs_per_k_d.device_ptr_mut(&self.stream); let (gprobs_base_bwd, _g_gprobs_bwd) = self.grad_probs_per_k_d.device_ptr_mut(&self.stream); - let (grad_henr_base, _g_ghen_mut) = grad_h_enriched_seq.data_mut().device_ptr_mut(&self.stream); + let (grad_henr_t_base, _g_ghen_t_mut) = grad_h_enriched_seq_t.data_mut().device_ptr_mut(&self.stream); for k in (0..k_seq).rev() { - let h_new_k_ptr = h_per_k_base_bwd + (k * hid_bytes) as u64; - let x_k_ptr = henr_base_bwd + (k * hid_bytes) as u64; - let probs_k_ptr = probs_base_bwd + (k * nh_bytes) as u64; - let gprobs_k_ptr = gprobs_base_bwd + (k * nh_bytes) as u64; + let h_new_k_ptr = h_per_k_base_bwd + (k * kb_hid_bytes) as u64; + let x_k_ptr = henr_t_base_bwd + (k * kb_hid_bytes) as u64; + let probs_k_ptr = probs_base_bwd + (k * kb_nh_bytes) as u64; + let gprobs_k_ptr = gprobs_base_bwd + (k * kb_nh_bytes) as u64; let h_old_k_ptr = if k == 0 { zero_h_ptr_bwd } else { - h_per_k_base_bwd + ((k - 1) * hid_bytes) as u64 + h_per_k_base_bwd + ((k - 1) * kb_hid_bytes) as u64 }; - let grad_henr_k_ptr = grad_henr_base + (k * hid_bytes) as u64; + let grad_henr_k_ptr = grad_henr_t_base + (k * kb_hid_bytes) as u64; - // heads_bwd writes grad_h_new = heads-side grad + carry. + // heads_bwd_batched writes grad_h_new[B, 128] = heads-side + carry. unsafe { - let mut launch = self.stream.launch_builder(&self.heads_bwd_fn); + let mut launch = self.stream.launch_builder(&self.heads_bwd_batched_fn); launch .arg(&self.heads_w_d) .arg(&probs_k_ptr) .arg(&h_new_k_ptr) .arg(&gprobs_k_ptr) .arg(&self.grad_h_carry_d) + .arg(&n_batch_i) .arg(&mut self.grad_heads_w_d) .arg(&mut self.grad_heads_b_d) .arg(&mut self.grad_h_new_d); - launch.launch(cfg_heads).context("heads bwd k")?; + launch.launch(cfg_heads_bwd).context("heads bwd k batched")?; } - // cfc_step_bwd consumes grad_h_new, emits new grad_h_carry - // (= grad_h_old_out, propagated back to position k-1). + // cfc_step_bwd_batched: writes grad_h_carry (= grad_h_old_out for + // position k-1) and grad_x[B, n_in]. grad_x output buffer IS + // grad_henr_k_ptr — kernel writes directly into the K-slot. unsafe { - let mut launch = self.stream.launch_builder(&self.step_bwd_fn); + let mut launch = self.stream.launch_builder(&self.step_bwd_batched_fn); launch .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr).arg(&self.grad_h_new_d) - .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&mut self.grad_w_in_d).arg(&mut self.grad_w_rec_d) .arg(&mut self.grad_b_d).arg(&mut self.grad_tau_d) - .arg(&mut self.grad_h_carry_d) // grad_h_old → carry - .arg(&mut self.grad_x_d); - launch.launch(cfg_cfc_bwd).context("cfc bwd k")?; - } - - // grad_x_k → grad_h_enriched_seq[k]. - unsafe { - let (src_ptr, _g) = self.grad_x_d.device_ptr(&self.stream); - cudarc::driver::result::memcpy_dtod_async( - grad_henr_k_ptr, src_ptr, hid_bytes, self.stream.cu_stream(), - ).context("grad_h_enriched_seq slot copy")?; + .arg(&mut self.grad_h_carry_d) + .arg(&grad_henr_k_ptr); + launch.launch(cfg_cfc_bwd).context("cfc bwd k batched")?; } } - drop((_g_hpk_bwd, _g_probs_bwd, _g_gprobs_bwd, _g_ghen_mut)); + drop((_g_hpk_bwd, _g_probs_bwd, _g_gprobs_bwd, _g_ghen_t_mut)); // ── 7. Sync once before optimizer step; download loss scalar. self.stream.synchronize().context("bwd loop sync")?; + // ── 7b. Transpose grad_h_enriched_seq_t [K, B, H] → [B, K, H] + // so Mamba2.backward_from_h_enriched_seq sees the layout + // it expects. + let mut grad_h_enriched_seq = GpuTensor::zeros( + &[b_sz, k_seq, HIDDEN_DIM], + &self.stream, + ).map_err(|e| anyhow::anyhow!("grad_h_enriched_seq alloc: {e}"))?; + { + let block_n3: u32 = 32; + let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3); + let cfg_tx = LaunchConfig { + grid_dim: (k_seq as u32, b_sz as u32, grid_z), + block_dim: (block_n3, 1, 1), + shared_mem_bytes: 0, + }; + let n1 = k_seq as i32; + let n2 = b_sz as i32; + let n3 = HIDDEN_DIM as i32; + let mut launch = self.stream.launch_builder(&self.transpose_3d_fn); + launch + .arg(grad_h_enriched_seq_t.cuda_data()) + .arg(grad_h_enriched_seq.data_mut()) + .arg(&n1).arg(&n2).arg(&n3); + unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; } + } + // ── 8. Mamba2 backward — single call consumes grad_h_enriched_seq. let mamba2_grads = self .mamba2 @@ -682,60 +775,79 @@ impl PerceptionTrainer { Ok(loss_host[0]) } - /// Forward-only evaluation pass. Runs the same forward chain as - /// [`step`] (snap_features → Mamba2 → per-K CfC+heads → fused BCE) - /// but SKIPS the backward, gradient accumulation, and AdamW step. - /// Validation MUST use this — calling `step()` on val data leaks - /// the val gradient into the trained weights (real bug, not a - /// statistical-noise hand-wave; per-step supervision multiplies - /// each val sequence's gradient pressure ~K-fold). - /// - /// Returns (mean_bce_loss, probs_per_position[K * N_HORIZONS]). + /// Single-sequence forward-only eval — thin wrapper around + /// [`evaluate_batched`] for cfg.n_batch == 1. pub fn evaluate( &mut self, snapshots: &[Mbp10RawInput], labels_per_position: &[[f32; N_HORIZONS]], ) -> Result<(f32, Vec)> { - let k_seq = self.cfg.seq_len; - anyhow::ensure!(snapshots.len() == k_seq, "snapshots.len()={} != seq_len", snapshots.len()); - anyhow::ensure!(labels_per_position.len() == k_seq, "labels.len() != seq_len"); + anyhow::ensure!( + self.cfg.n_batch == 1, + "single-sequence evaluate() requires cfg.n_batch == 1; got {}", + self.cfg.n_batch + ); + self.evaluate_batched(&[snapshots], &[labels_per_position]) + } - // Build snap_feature window (same as step()). - let mut window_tensor = GpuTensor::zeros(&[1, k_seq, FEATURE_DIM], &self.stream) + /// Forward-only batched evaluation pass. Mirrors the forward chain + /// of [`step_batched`] but SKIPS backward + grad accumulation + + /// AdamW. Validation MUST use this — calling step() on val data + /// leaks the val gradient into trained weights. + /// + /// Returns (weighted mean BCE, probs[K * B * N_HORIZONS] in + /// [K, B, N_HORIZONS] row-major layout). + pub fn evaluate_batched( + &mut self, + snapshots_batch: &[&[Mbp10RawInput]], + labels_batch: &[&[[f32; N_HORIZONS]]], + ) -> Result<(f32, Vec)> { + let b_sz = self.cfg.n_batch; + let k_seq = self.cfg.seq_len; + anyhow::ensure!( + snapshots_batch.len() == b_sz && labels_batch.len() == b_sz, + "expected B={} snapshots+labels; got snap={} lbl={}", + b_sz, snapshots_batch.len(), labels_batch.len() + ); + + // Build snap_feature window [B, K, FEATURE_DIM]. + let mut window_tensor = GpuTensor::zeros(&[b_sz, k_seq, FEATURE_DIM], &self.stream) .map_err(|e| anyhow::anyhow!("eval window alloc: {e}"))?; let snap_cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0, }; let snap_nbytes = FEATURE_DIM * std::mem::size_of::(); - for (k, snap) in snapshots.iter().enumerate() { - upload_into(&self.stream, &snap.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?; - upload_into(&self.stream, &snap.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?; - upload_into(&self.stream, &snap.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?; - upload_into(&self.stream, &snap.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?; - upload_into(&self.stream, &snap.regime, &self.stg_regime, &mut self.regime_d)?; - let prev_mid = snap.prev_mid; - let trade_signed_vol = snap.trade_signed_vol; - let trade_count = snap.trade_count as i32; - let ts_ns = snap.ts_ns as i64; - let prev_ts_ns = snap.prev_ts_ns as i64; - let tick_size = ES_TICK_SIZE; - { - let mut launch = self.stream.launch_builder(&self.snap_fn); - launch - .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d) - .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d).arg(&self.regime_d) - .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count) - .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size) - .arg(&mut self.snap_feat_d); - unsafe { launch.launch(snap_cfg).context("eval snap fwd")?; } - } - unsafe { - let (src_ptr, _g1) = self.snap_feat_d.device_ptr(&self.stream); - let (dst_base, _g2) = window_tensor.data_mut().device_ptr_mut(&self.stream); - let dst_offset_ptr = dst_base + (k * snap_nbytes) as u64; - cudarc::driver::result::memcpy_dtod_async( - dst_offset_ptr, src_ptr, snap_nbytes, self.stream.cu_stream(), - ).context("eval window pack")?; + for (b_idx, snapshots) in snapshots_batch.iter().enumerate() { + for (k, snap) in snapshots.iter().enumerate() { + upload_into(&self.stream, &snap.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?; + upload_into(&self.stream, &snap.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?; + upload_into(&self.stream, &snap.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?; + upload_into(&self.stream, &snap.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?; + upload_into(&self.stream, &snap.regime, &self.stg_regime, &mut self.regime_d)?; + let prev_mid = snap.prev_mid; + let trade_signed_vol = snap.trade_signed_vol; + let trade_count = snap.trade_count as i32; + let ts_ns = snap.ts_ns as i64; + let prev_ts_ns = snap.prev_ts_ns as i64; + let tick_size = ES_TICK_SIZE; + { + let mut launch = self.stream.launch_builder(&self.snap_fn); + launch + .arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d) + .arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d).arg(&self.regime_d) + .arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count) + .arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size) + .arg(&mut self.snap_feat_d); + unsafe { launch.launch(snap_cfg).context("eval snap fwd")?; } + } + unsafe { + let (src_ptr, _g1) = self.snap_feat_d.device_ptr(&self.stream); + let (dst_base, _g2) = window_tensor.data_mut().device_ptr_mut(&self.stream); + let dst_offset_ptr = dst_base + ((b_idx * k_seq + k) * snap_nbytes) as u64; + cudarc::driver::result::memcpy_dtod_async( + dst_offset_ptr, src_ptr, snap_nbytes, self.stream.cu_stream(), + ).context("eval window pack")?; + } } } self.stream.synchronize().context("eval window sync")?; @@ -745,13 +857,37 @@ impl PerceptionTrainer { .forward_train_seq(&window_tensor) .context("eval mamba2 fwd")?; - // Upload labels for BCE. - let total_labels = k_seq * N_HORIZONS; + // Transpose Mamba2 output [B, K, H] → [K, B, H]. + let mut h_enriched_seq_t = GpuTensor::zeros(&[k_seq, b_sz, HIDDEN_DIM], &self.stream) + .map_err(|e| anyhow::anyhow!("eval h_enriched_seq_t alloc: {e}"))?; + { + let block_n3: u32 = 32; + let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3); + let cfg_tx = LaunchConfig { + grid_dim: (b_sz as u32, k_seq as u32, grid_z), + block_dim: (block_n3, 1, 1), + shared_mem_bytes: 0, + }; + let n1 = b_sz as i32; + let n2 = k_seq as i32; + let n3 = HIDDEN_DIM as i32; + let mut launch = self.stream.launch_builder(&self.transpose_3d_fn); + launch + .arg(h_enriched_seq.cuda_data()) + .arg(h_enriched_seq_t.data_mut()) + .arg(&n1).arg(&n2).arg(&n3); + unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; } + } + + // Upload labels [K, B, N_HORIZONS]. + let total_labels = k_seq * b_sz * N_HORIZONS; { let dst = self.stg_labels.host_slice_mut(); for k in 0..k_seq { - for h in 0..N_HORIZONS { - dst[k * N_HORIZONS + h] = labels_per_position[k][h]; + for b_idx in 0..b_sz { + for h in 0..N_HORIZONS { + dst[(k * b_sz + b_idx) * N_HORIZONS + h] = labels_batch[b_idx][k][h]; + } } } } @@ -762,10 +898,9 @@ impl PerceptionTrainer { total_labels * std::mem::size_of::(), self.stream.cu_stream(), ).context("eval labels upload")?; } - // Reset h_old to zero start-of-sequence. self.stream.memset_zeros(&mut self.zero_h_d).map_err(|e| anyhow::anyhow!("zero h: {e}"))?; - // Per-K forward (recurrent CfC + heads). + // Per-K forward (recurrent CfC + heads, batched). let dt_s = 1.0_f32; let n_in_i = HIDDEN_DIM as i32; let n_hid_i = HIDDEN_DIM as i32; @@ -775,47 +910,51 @@ impl PerceptionTrainer { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0, }; let cfg_heads = LaunchConfig { - grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0, + grid_dim: (1, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0, }; - let hid_bytes = HIDDEN_DIM * std::mem::size_of::(); - let nh_bytes = N_HORIZONS * std::mem::size_of::(); + let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::(); + let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::(); + let n_batch_i = b_sz as i32; let zero_h_ptr = { let (p, _g) = self.zero_h_d.device_ptr(&self.stream); p }; - let henr_base = { - let (p, _g) = h_enriched_seq.cuda_data().device_ptr(&self.stream); + let henr_t_base = { + let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream); p }; let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream); let (probs_base, _g_probs) = self.probs_per_k_d.device_ptr_mut(&self.stream); for k in 0..k_seq { - let h_new_k_ptr = h_per_k_base + (k * hid_bytes) as u64; - let h_old_k_ptr = if k == 0 { zero_h_ptr } else { h_per_k_base + ((k - 1) * hid_bytes) as u64 }; - let x_k_ptr = henr_base + (k * hid_bytes) as u64; - let probs_k_ptr = probs_base + (k * nh_bytes) as u64; + let h_new_k_ptr = h_per_k_base + (k * kb_hid_bytes) as u64; + let h_old_k_ptr = if k == 0 { + zero_h_ptr + } else { + h_per_k_base + ((k - 1) * kb_hid_bytes) as u64 + }; + let x_k_ptr = henr_t_base + (k * kb_hid_bytes) as u64; + let probs_k_ptr = probs_base + (k * kb_nh_bytes) as u64; unsafe { - let mut launch = self.stream.launch_builder(&self.step_fn); + let mut launch = self.stream.launch_builder(&self.step_batched_fn); launch .arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d) .arg(&x_k_ptr).arg(&h_old_k_ptr) - .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i) .arg(&h_new_k_ptr); launch.launch(cfg_cfc).context("eval cfc fwd")?; } unsafe { - let mut launch = self.stream.launch_builder(&self.heads_fn); + let mut launch = self.stream.launch_builder(&self.heads_batched_fn); launch .arg(&self.heads_w_d).arg(&self.heads_b_d) - .arg(&h_new_k_ptr).arg(&probs_k_ptr); + .arg(&h_new_k_ptr).arg(&n_batch_i).arg(&probs_k_ptr); launch.launch(cfg_heads).context("eval heads fwd")?; } } drop((_g_hpk, _g_probs)); - // BCE forward for loss reporting (we don't use grad_probs but the - // kernel emits both; cheap enough to call). - let n_pos_i = k_seq as i32; + // BCE for loss reporting (uses same per-horizon weights as train). + let n_pos_i = (k_seq * b_sz) as i32; let n_h_i = N_HORIZONS as i32; let bce_cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, @@ -823,7 +962,7 @@ impl PerceptionTrainer { unsafe { let mut launch = self.stream.launch_builder(&self.bce_fn); launch - .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d) + .arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d) .arg(&n_pos_i).arg(&n_h_i) .arg(&mut self.loss_d).arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d); launch.launch(bce_cfg).context("eval bce")?; diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index 26ce44124..5f623fe56 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -71,6 +71,7 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() { lr_mamba2: 1e-3, seed: 0x4242, horizon_weights: [1.0; 5], + n_batch: 1, }; let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");