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:
jgrusewski
2026-05-17 09:13:48 +02:00
parent 16f5febf27
commit 485150c7b7
2 changed files with 4 additions and 3 deletions

View File

@@ -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).

View File

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