From 485150c7b7976c31c4f08d5b55f11078d474420c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 09:13:48 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20lift=20Mamba2=20kernel=20seq?= =?UTF-8?q?=5Flen=20cap=20from=2032=20=E2=86=92=20256?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/cuda/mamba2_alpha_kernel.cu | 5 +++-- crates/ml-alpha/src/mamba2_block.rs | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu index 09aa122ad..98cd3b714 100644 --- a/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu +++ b/crates/ml-alpha/cuda/mamba2_alpha_kernel.cu @@ -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 #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). diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs index 3c16ca839..b3d707500 100644 --- a/crates/ml-alpha/src/mamba2_block.rs +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -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)]