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

feat(ml-alpha): raise Mamba2 state cap 16→32 + add auc_h6000 early-stop

Browse Source 
Three correlated changes for the next CV round:

1. Mamba2 state_dim cap: 16 → 32
   cuda/mamba2_alpha_kernel.cu: MAMBA2_ALPHA_MAX_STATE_D 16 → 32.
   Per-thread state register `float x[32]` (128 B/thread) and
   per-thread x_hist replay cache `float x_hist[K*32]` (up to
   12 KiB/thread of local memory at K=96). L40S/H100 register file
   (256 KiB/SM) absorbs this without occupancy collapse for our
   block dims (32-128 threads). Update Rust-side
   MAMBA2_KERNEL_STATE_MAX + validation message + test name. Kernel
   header doc updated.

2. New early-stop option: auc_h6000
   examples/alpha_train.rs: add the long-horizon AUC as a third
   early-stop metric. The ISV CV (3a196382f, 5d42ab0e9, 0171c8c0e)
   showed mean_auc-best-epoch and h6000-best-epoch can differ by
   1-2 epochs and the h6000 gap can be 5-6pt within a single run
   (fblb2 fold-1: saved E10 h6000=0.681, but E11 h6000=0.739 — we
   threw away the deployment-better checkpoint). For multi-minute
   trading deployment we want the h6000-best checkpoint directly.

3. Summary JSON: best_auc_h6000_epoch / best_auc_h6000 /
   best_auc_h6000_per_horizon
   So the analysis tooling can see the h6000-best checkpoint
   independently of mean_auc / val_loss bests.

Test rename: test_mamba2_config_rejects_state_over_16 →
test_mamba2_config_rejects_state_over_32 (tests now reject state_dim=33).

build.rs cache-bust v4 forces cluster nodes to recompile the kernel
against the new MAMBA2_ALPHA_MAX_STATE_D — old cubins from previous
SHAs were sized for state_d=16 and would silently truncate state_d=32
state arrays.

Validation: 26 lib + 23 integration ml-alpha tests pass. Mamba2 block
tests use state_dim=8 or 16 (well below the new cap), exercise both
the forward + backward + AdamW paths.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-17 20:25:47 +02:00
parent d220755309
commit a45fd85986
4 changed files with 79 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ml-alpha/build.rs
View File

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

13
crates/ml-alpha/cuda/mamba2_alpha_kernel.cu
View File

@@ -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 <cuda_runtime.h>
#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
/* ---------------------------------------------------------------------

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

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

21
crates/ml-alpha/src/mamba2_block.rs
View File

@@ -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]