feat(ml-alpha): lift Mamba2 kernel seq_len cap from 32 → 256
Previous BPTT-unroll run had val_loss trending (0.6941→0.6922 over 5
epochs) but AUCs oscillating around 0.50 — the architecture lacked
context for medium/long horizons (h300, h1000, h6000 ≫ seq_len=32).
Phase 1d.2 validated the SSM at seq_len=6000; this is a step toward
restoring useful sequence depth.
Bumps:
- MAMBA2_ALPHA_MAX_K constant: 32 → 256
- x_hist per-thread replay buffer: 2KB → 16KB (spills to
DRAM-backed per-thread local memory; L2-cached, acceptable
perf cost vs the 8x context gain)
- Mamba2BlockConfig::validate updates the cap
Backward compat: legacy Phase E.3 callers (alpha_baseline,
alpha_dqn_h600_smoke) only use K=12 / K=32 — unaffected by the
larger compile-time max.
Synthetic overfit still converges 0.7135 → 0.2079 in 250 steps.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -30,7 +30,8 @@
|
||||
*
|
||||
* Kernel constraints (must hold at every launch):
|
||||
* state_d ≤ 16 (register array `float x[16]`)
|
||||
* K ≤ 32 (local-memory array `float x_hist[32 * 16]` ≤ 2 KiB / thread)
|
||||
* K ≤ 256 (local-memory array `float x_hist[256 * 16]` = 16 KiB / thread —
|
||||
* spills to per-thread DRAM-backed local memory but L2-cached)
|
||||
*
|
||||
* Rust constructor enforces both via Mamba2BlockConfig::validate.
|
||||
* ===================================================================== */
|
||||
@@ -38,7 +39,7 @@
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define MAMBA2_ALPHA_MAX_STATE_D 16
|
||||
#define MAMBA2_ALPHA_MAX_K 32
|
||||
#define MAMBA2_ALPHA_MAX_K 256
|
||||
|
||||
/* ---------------------------------------------------------------------
|
||||
* Forward scan: one thread per (sample i, hidden-channel j).
|
||||
|
||||
@@ -47,7 +47,7 @@ use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLin
|
||||
/// `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;
|
||||
pub const MAMBA2_KERNEL_SEQ_MAX: usize = 32;
|
||||
pub const MAMBA2_KERNEL_SEQ_MAX: usize = 256;
|
||||
|
||||
/// Configuration for a Mamba2 sequence block.
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
Reference in New Issue
Block a user