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

feat(ml-alpha): wire batched kernels through trainer + CLI (#8)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 10:42:35 +02:00
parent 829ddfa62c
commit c3ee5e165a
3 changed files with 425 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
crates/ml-alpha/examples/alpha_train.rs
View File

@@ -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<Mbp10RawInput>> = Vec::with_capacity(cli.batch_size);
let mut label_batch: Vec<Vec<[f32; N_HORIZONS]>> = 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<f32>; N_HORIZONS] = Default::default();
let mut val_loss_sum = 0.0_f32;
let mut val_steps = 0usize;
let mut val_snap_batch: Vec<Vec<Mbp10RawInput>> = Vec::with_capacity(cli.batch_size);
let mut val_label_batch: Vec<Vec<[f32; N_HORIZONS]>> = 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 {

599
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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<CudaModule>,
_bce_module: Arc<CudaModule>,
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<f32>,
/// 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<f32>,
/// Single-position zero h_old buffer (used only at k=0 — first
/// position of the sequence has no preceding state).
zero_h_d: CudaSlice<f32>,
@@ -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::<f32>(10)?,
regime_d: stream.alloc_zeros::<f32>(REGIME_DIM)?,
snap_feat_d: stream.alloc_zeros::<f32>(FEATURE_DIM)?,
h_new_per_k_d: stream.alloc_zeros::<f32>(k * n_hid)?,
probs_per_k_d: stream.alloc_zeros::<f32>(k * N_HORIZONS)?,
labels_per_k_d: stream.alloc_zeros::<f32>(k * N_HORIZONS)?,
grad_probs_per_k_d: stream.alloc_zeros::<f32>(k * N_HORIZONS)?,
h_new_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * n_hid)?,
probs_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
labels_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
grad_probs_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
loss_weights_d: upload(&stream, &cfg.horizon_weights)?,
loss_d: stream.alloc_zeros::<f32>(1)?,
valid_d: stream.alloc_zeros::<i32>(1)?,
grad_h_carry_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_h_new_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_x_d: stream.alloc_zeros::<f32>(n_in)?,
zero_h_d: stream.alloc_zeros::<f32>(n_hid)?,
grad_h_carry_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
grad_h_new_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
zero_h_d: stream.alloc_zeros::<f32>(cfg.n_batch * n_hid)?,
grad_w_in_d: stream.alloc_zeros::<f32>(n_hid * n_in)?,
grad_w_rec_d: stream.alloc_zeros::<f32>(n_hid * n_hid)?,
grad_b_d: stream.alloc_zeros::<f32>(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<f32> {
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<f32> {
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::<f32>();
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::<f32>()) 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::<f32>()) 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::<f32>()) 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::<f32>();
let nh_bytes = N_HORIZONS * std::mem::size_of::<f32>();
// 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::<f32>();
let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::<f32>();
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<f32>)> {
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<f32>)> {
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::<f32>();
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::<f32>(), 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::<f32>();
let nh_bytes = N_HORIZONS * std::mem::size_of::<f32>();
let kb_hid_bytes = b_sz * HIDDEN_DIM * std::mem::size_of::<f32>();
let kb_nh_bytes = b_sz * N_HORIZONS * std::mem::size_of::<f32>();
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")?;

1
crates/ml-alpha/tests/perception_overfit.rs
View File

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