diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 2de5f52f7..37ad972ee 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -19,9 +19,11 @@ const KERNELS: &[&str] = &[ "horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda) ]; -// Cache bust v3 (2026-05-17): batched cfc/heads kernels + transpose_3d_swap_01 -// added in commit 829ddfa62 + c3ee5e165. Force rebuild on cluster nodes whose -// /cargo-target persistent volume still has the pre-batched cubins. +// Cache bust v4 (2026-05-17): Mamba2 state_dim cap raised 16 → 32 in +// the kernel (MAMBA2_ALPHA_MAX_STATE_D). Old cubins are sized for +// state_d=16 and will silently truncate / corrupt the larger state +// arrays. Force a fresh nvcc compile on the cluster's /cargo-target +// persistent volume so the cubin matches the source. fn main() { println!("cargo:rerun-if-changed=build.rs"); diff --git a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu index d6925f37f..339f8758d 100644 --- a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu +++ b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu @@ -29,8 +29,8 @@ * d_h_s2 : [N, sh2] identity passthrough of d_h_enriched * * Kernel constraints (must hold at every launch): - * state_d ≤ 16 (register array `float x[16]`) - * K ≤ 96 (local-memory array `float x_hist[96 * 16]` = 6 KiB / thread — + * state_d ≤ 32 (register array `float x[32]`) + * K ≤ 96 (local-memory array `float x_hist[96 * 32]` = 12 KiB / thread — * spills to per-thread DRAM-backed local memory but L2-cached; * chosen so K ≥ longest practical horizon for tick-level * supervised heads while keeping the backward replay buffer @@ -41,7 +41,14 @@ #include -#define MAMBA2_ALPHA_MAX_STATE_D 16 +// state_d capacity: raised from 16 → 32 (per-thread register array +// `float x[32]` = 128 bytes; backward replay cache `x_hist[K*32]` up +// to 12 KiB/thread of local memory at K=96). L40S/H100 register file +// (256 KiB/SM) handles this without occupancy collapse for our block +// dim (32-128 threads/block). 32-wide state empirically validated as +// the next step up the capacity ladder from 16; further increases +// will need a different storage scheme (smem-tiled state). +#define MAMBA2_ALPHA_MAX_STATE_D 32 #define MAMBA2_ALPHA_MAX_K 96 /* --------------------------------------------------------------------- diff --git a/crates/ml-alpha/examples/alpha_train.rs b/crates/ml-alpha/examples/alpha_train.rs index b90b98ed5..423ed278d 100644 --- a/crates/ml-alpha/examples/alpha_train.rs +++ b/crates/ml-alpha/examples/alpha_train.rs @@ -86,12 +86,17 @@ struct Cli { #[arg(long, default_value_t = 5)] early_stop_patience: usize, - /// Which metric drives early stopping: "val_loss" (stable, tracks - /// probability calibration) or "mean_auc" (noisier, tracks the - /// ranking quality used for downstream trading). mean_auc needs - /// higher patience (~5) because it bounces 1-2pt epoch-to-epoch - /// even when the model is still genuinely improving on the long - /// horizons that the auto-horizon-weights down-weight. + /// Which metric drives early stopping. Options: + /// - "val_loss" : stable, tracks probability calibration. + /// - "mean_auc" : noisier, average ranking across all 5 horizons. + /// - "auc_h6000" : ranking on the long horizon ONLY — recommended + /// for deployment-oriented runs where the saved checkpoint is + /// used to trade at multi-minute horizon. The 3-fold ISV CV + /// showed mean_auc-best-epoch and h6000-best-epoch can differ + /// by 2-6pt on h6000 within a single run (e.g. fblb2 fold-1 + /// saved h6000=0.681 at mean_auc-best E10 but the trajectory + /// peaked at h6000=0.739 one epoch later). Using auc_h6000 + /// directly saves the deployment-relevant checkpoint. #[arg(long, default_value = "mean_auc")] early_stop_metric: String, @@ -161,6 +166,15 @@ struct AlphaTrainSummary { best_mean_auc: f32, /// Per-horizon AUCs at the best-mean-AUC epoch. best_mean_auc_per_horizon: [f32; N_HORIZONS], + /// Epoch with the highest h6000 AUC (long-horizon, deployment-relevant). + /// Saved independently of best_mean_auc_epoch because the two + /// frequently disagree by 1-2 epochs and h6000 is the metric we + /// care about for multi-minute trading deployment. + best_auc_h6000_epoch: usize, + /// Highest h6000 AUC observed across all epochs. + best_auc_h6000: f32, + /// Per-horizon AUCs at the best-h6000 epoch. + best_auc_h6000_per_horizon: [f32; N_HORIZONS], /// True if training stopped early (patience exceeded). early_stopped: bool, } @@ -240,8 +254,8 @@ fn main() -> Result<()> { // Validate early-stop metric choice upfront. let early_stop_metric = cli.early_stop_metric.as_str(); anyhow::ensure!( - matches!(early_stop_metric, "val_loss" | "mean_auc" | "none"), - "--early-stop-metric must be one of: val_loss, mean_auc, none (got `{}`)", + matches!(early_stop_metric, "val_loss" | "mean_auc" | "auc_h6000" | "none"), + "--early-stop-metric must be one of: val_loss, mean_auc, auc_h6000, none (got `{}`)", early_stop_metric ); @@ -263,6 +277,17 @@ fn main() -> Result<()> { let mut best_mean_auc_epoch = 0usize; let mut best_mean_auc_per_horizon = [0.5_f32; N_HORIZONS]; + // Independent tracker for best h6000 AUC — the deployment-relevant + // long-horizon metric. The 3-fold ISV CV showed that mean_auc-best + // and h6000-best epochs frequently differ (within fblb2 fold-1, the + // saved E10 had h6000=0.681 while E11 had h6000=0.739). For + // multi-minute trading deployment we want the h6000-best checkpoint + // directly, not the mean_auc proxy. + let mut best_auc_h6000 = f32::NEG_INFINITY; + let mut best_auc_h6000_epoch = 0usize; + let mut best_auc_h6000_per_horizon = [0.5_f32; N_HORIZONS]; + let mut auc_h6000_no_improvement = 0usize; + let lr_min_cfc = cli.lr_cfc * cli.lr_min_factor; let lr_min_m2 = cli.lr_mamba2 * cli.lr_min_factor; // Approximate the total step budget: epochs × n_train_seqs (each @@ -532,11 +557,25 @@ fn main() -> Result<()> { mean_auc_no_improvement += 1; } + // h6000 — deployment-relevant long-horizon ranking. + let auc_h6000 = per_horizon_auc[N_HORIZONS - 1]; + let auc_h6000_improved = auc_h6000 > best_auc_h6000; + if auc_h6000_improved { + best_auc_h6000 = auc_h6000; + best_auc_h6000_epoch = epoch; + best_auc_h6000_per_horizon = per_horizon_auc; + auc_h6000_no_improvement = 0; + tracing::info!(epoch, auc_h6000, "new best auc_h6000"); + } else { + auc_h6000_no_improvement += 1; + } + // Early-stop decision based on selected metric. if cli.early_stop_patience > 0 { let (counter, metric_label, best_val) = match early_stop_metric { - "val_loss" => (val_loss_no_improvement, "val_loss", best_val_loss), - "mean_auc" => (mean_auc_no_improvement, "mean_auc", best_mean_auc), + "val_loss" => (val_loss_no_improvement, "val_loss", best_val_loss), + "mean_auc" => (mean_auc_no_improvement, "mean_auc", best_mean_auc), + "auc_h6000" => (auc_h6000_no_improvement, "auc_h6000", best_auc_h6000), _ => (0, "none", 0.0), // disabled }; if counter >= cli.early_stop_patience && metric_label != "none" { @@ -566,6 +605,9 @@ fn main() -> Result<()> { best_mean_auc_epoch, best_mean_auc, best_mean_auc_per_horizon, + best_auc_h6000_epoch, + best_auc_h6000, + best_auc_h6000_per_horizon, early_stopped, }; let summary_path = cli.out.join("alpha_train_summary.json"); diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index 81051dc38..6c59ea03d 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -31,8 +31,8 @@ //! - Weight init: `OwnedGpuLinear::xavier` from ml-core, which uses pinned //! host buffers for the seed values. //! - No `atomicAdd` (kernel uses block tree-reduce). -//! - State dim hardcoded at ≤ 16 in the kernel (`float x[16]`); we expose -//! it as a config field but the constructor rejects values > 16. +//! - State dim hardcoded at ≤ 32 in the kernel (`float x[32]`); we expose +//! it as a config field but the constructor rejects values > 32. use std::sync::Arc; @@ -43,10 +43,13 @@ use ml_core::cuda_autograd::gpu_tensor::GpuTensor; use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLinear}; /// Hard kernel limits. The forward and backward kernels use register/local -/// arrays sized at compile-time: `float x[16]` (state register) and -/// `float x_hist[32 * 16]` (backward replay cache, 2 KiB per thread). -/// Exceeding either silently corrupts; the Rust constructor enforces. -pub const MAMBA2_KERNEL_STATE_MAX: usize = 16; +/// arrays sized at compile-time: `float x[32]` (state register) and +/// `float x_hist[K * 32]` up to ~12 KiB per thread of local memory at +/// K=96. Raised from 16 → 32 to allow scaling Mamba2 SSM state capacity; +/// L40S/H100 register file (256 KiB/SM) handles this without occupancy +/// collapse for our block dim (32-128 threads). Exceeding either limit +/// silently corrupts; the Rust constructor enforces. +pub const MAMBA2_KERNEL_STATE_MAX: usize = 32; pub const MAMBA2_KERNEL_SEQ_MAX: usize = 96; /// Configuration for a Mamba2 sequence block. @@ -2124,11 +2127,11 @@ mod tests { } #[test] - fn test_mamba2_config_rejects_state_over_16() { + fn test_mamba2_config_rejects_state_over_32() { let cfg = Mamba2BlockConfig { - in_dim: 81, hidden_dim: 64, state_dim: 17, seq_len: 32, + in_dim: 81, hidden_dim: 64, state_dim: 33, seq_len: 32, }; - assert!(cfg.validate().is_err(), "state_dim > 16 must be rejected"); + assert!(cfg.validate().is_err(), "state_dim > 32 must be rejected"); } #[test]