629ebd667c
Two same-seed runs now produce bit-equal eval_summary.json, alpha_rl_train_summary.json,
and diag.jsonl (modulo wall-clock elapsed_s). The 5-phase falsification chain landed:
Phase 2 PER tree-rebuild: __threadfence is NOT a grid-wide barrier; multiple blocks
raced across sum-tree levels. Fix: Grid=(1) Block=(1024) + __syncthreads
in rl_per_tree_rebuild.cu.
Phase 2.3 cuBLAS GEMM_DFALT + TF32 default-math allowed split-K non-deterministic
accumulation at 3 sites. New crates/ml-alpha/src/cublas_determinism.rs
applies CUBLAS_PEDANTIC_MATH via FOXHUNT_DETERMINISTIC env toggle
(0=TF32 prod, 1=PEDANTIC dev default, 2=DEFAULT_MATH control).
Phase 2.6 Two bugs surfaced sequentially in the backward kernel chain:
(1) rl_iqn_tau_cos_features had a multi-block r/w race on prng_state[batch]
— all N_TAU=32 blocks read seed; only tau_idx==0 wrote back; no
inter-block barrier. Fix: split into READ-ONLY rl_iqn_tau_cos_features
+ new sibling rl_iqn_advance_prng_state launched on same stream
(kernel-launch ordering = grid-wide barrier).
(2) OutcomeHead::new called near_zero_xavier without scoped_init_seed,
falling back to time+thread-id RNG. Stayed dormant until first done
event activated non-sentinel labels and divergent weights flowed via
grad_h_t_outcome into encoder gradient. Fix: add seed param + install
scoped_init_seed(dqn_seed.wrapping_add(0x0CE0)) guard.
Validation (./scripts/determinism-check.sh --quick, RTX 3050, b=128, 200+50 steps):
- All 200 rows of checksums.* leaves match (rel-tol 1e-5, abs-tol 1e-7)
- eval_summary.json, alpha_rl_train_summary.json byte-equal between runs
- diag.jsonl byte-equal modulo elapsed_s
- Eval pnl identical run-A vs run-B at seed 42
Pre-fix baseline (Phase 2.5 measurement): same-seed eval pnl spread $450k
($187k vs -$261k). Post-fix: $0 spread.
Speed cost: ~1.5ms/step amortised; ~10-15% slower than TF32 production
(PEDANTIC tax — acceptable in dev, toggle to FOXHUNT_DETERMINISTIC=0 for prod).
Mapped-pinned discipline: all 11 NEW memcpy_dtoh sites in diagnostic dump methods
+ per-step checksum readback use a new pub(crate) helper
read_slice_d_into<T: Copy>(stream, src, dst) — MappedRecordBuffer + raw
memcpy_dtod_async + raw_stream_sync + volatile read. Generic over T (f32, f64,
i32, u32, u8). Satisfies feedback_no_htod_htoh_only_mapped_pinned + hook guard.
Bundled Tier 1.5 fast-dev-cycle infrastructure (spec
docs/superpowers/specs/2026-06-02-fast-dev-cycle.md):
- scripts/local-mid-smoke.sh b=128, 2000+500, ~10min on RTX 3050
- scripts/determinism-check.sh runs mid-smoke twice, diffs checksums
- scripts/tier1_5_verdict.py behavioral kill verdict
- AdamW checkpoint save/load (crates/ml-alpha/src/trainer/optim.rs)
- IntegratedTrainer checkpoint save/load (resume from checkpoint)
- 15 Phase 1 checksum leaves in build_diag_value
- Env-gated dump methods (FOXHUNT_DETERMINISM_DEBUG_PER/MAMBA2/RL/BACKWARD)
for future divergence-chasing — never run in production
Documentation:
- docs/superpowers/specs/2026-06-02-determinism-foundation.md
- docs/superpowers/specs/2026-06-02-fast-dev-cycle.md
- docs/superpowers/plans/2026-06-02-determinism-foundation-implementation.md
- docs/superpowers/notes/2026-06-02-determinism-phase{1,2,2.2,2.5,2.6}-*.md
- Adjacent specs/plans/notes from the analytical chain that surfaced determinism
as the load-bearing blocker (eval-summary, eval-boundary, regime-observer,
multi-head policy, regime-invariance, Phase 3 IQN-complement post-mortem)
Unlocks: every controller / architecture / reward-shaping A/B from this commit
onward attributes outcome differences to the change, not random-init kernel-race
drift cascading through training x eval LOB-sim trajectories. The eval-collapse
investigation (pearl_reward_signal_anti_aligned_with_pnl, multi-head spec,
regime-invariance spec) is now testable with trustworthy verdicts.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2828 lines
127 KiB
Rust
2828 lines
127 KiB
Rust
//! Phase 1d.1 — From-scratch Mamba2 sequence block for ml-alpha.
|
||
//!
|
||
//! GPU-pure stateful encoder over snapshot streams. Architecture:
|
||
//!
|
||
//! ```text
|
||
//! input [B, K, in_dim]
|
||
//! │ cuBLAS GEMM (W_in)
|
||
//! ▼
|
||
//! x [B, K, hidden_dim]
|
||
//! ├─ cuBLAS GEMM (W_a) ──▶ a_proj [B, K, state_dim]
|
||
//! ├─ cuBLAS GEMM (W_b) ──▶ b_proj [B, K, state_dim]
|
||
//! ▼
|
||
//! mamba2_scan_projected_fwd kernel
|
||
//! (selective SSM scan over K timesteps; sigmoid gating; W_c mixes
|
||
//! state into per-position hidden output)
|
||
//! ▼
|
||
//! h_enriched [B, hidden_dim] (residual added in kernel via h_s2; we
|
||
//! pass zeros for the SSM-residual term)
|
||
//! │ cuBLAS GEMM (W_out)
|
||
//! ▼
|
||
//! logit [B, 1]
|
||
//! ```
|
||
//!
|
||
//! Backward: analytical via `mamba2_scan_projected_bwd` kernel; cuBLAS
|
||
//! transposed-GEMMs for the projection backward passes. Lands in a follow-up
|
||
//! session — this module establishes the skeleton: weight allocation,
|
||
//! cubin loading, kernel-symbol resolution.
|
||
//!
|
||
//! Discipline (per project memory):
|
||
//! - All weights / state / activations live on GPU (`CudaSlice<f32>`).
|
||
//! - 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 ≤ 32 in the kernel (`float x[32]`); we expose
|
||
//! it as a config field but the constructor rejects values > 32.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use anyhow::{anyhow, Result};
|
||
use cudarc::cublas::CudaBlas;
|
||
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
|
||
use cudarc::driver::sys::CUstream;
|
||
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
|
||
use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLinear};
|
||
|
||
use crate::trainer::raw_launch::{RawArgs, raw_launch};
|
||
|
||
/// Hard kernel limits. The forward and backward kernels use register/local
|
||
/// 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.
|
||
#[derive(Debug, Clone)]
|
||
pub struct Mamba2BlockConfig {
|
||
/// Per-snapshot feature dim (e.g. 81 for the snapshot pipeline).
|
||
pub in_dim: usize,
|
||
/// Hidden / model dim (=`sh2` in the kernel signature).
|
||
pub hidden_dim: usize,
|
||
/// SSM state dim (≤ 16).
|
||
pub state_dim: usize,
|
||
/// Sequence length per batch (K in kernel; number of timesteps).
|
||
pub seq_len: usize,
|
||
}
|
||
|
||
impl Mamba2BlockConfig {
|
||
pub fn validate(&self) -> Result<()> {
|
||
if self.in_dim == 0 || self.hidden_dim == 0 || self.state_dim == 0 || self.seq_len < 2 {
|
||
return Err(anyhow!(
|
||
"Mamba2BlockConfig: invalid dims (in={}, hidden={}, state={}, seq_len={})",
|
||
self.in_dim, self.hidden_dim, self.state_dim, self.seq_len
|
||
));
|
||
}
|
||
if self.state_dim > MAMBA2_KERNEL_STATE_MAX {
|
||
return Err(anyhow!(
|
||
"Mamba2BlockConfig: state_dim {} exceeds kernel max {} (mamba2_alpha_kernel.cu \
|
||
hardcodes `float x[{}]` register array per thread)",
|
||
self.state_dim, MAMBA2_KERNEL_STATE_MAX, MAMBA2_KERNEL_STATE_MAX
|
||
));
|
||
}
|
||
if self.seq_len > MAMBA2_KERNEL_SEQ_MAX {
|
||
return Err(anyhow!(
|
||
"Mamba2BlockConfig: seq_len {} exceeds kernel max {} (backward kernel's \
|
||
x_hist replay cache is `float x_hist[{}*{}]` per thread)",
|
||
self.seq_len, MAMBA2_KERNEL_SEQ_MAX,
|
||
MAMBA2_KERNEL_SEQ_MAX, MAMBA2_KERNEL_STATE_MAX
|
||
));
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// All parameter gradients produced by [`Mamba2Block::backward`].
|
||
///
|
||
/// Tensor shapes match the corresponding forward parameters:
|
||
/// - `dw_in / db_in` : `[hidden_dim, in_dim] / [hidden_dim]`
|
||
/// - `dw_a / db_a` : `[state_dim, hidden_dim] / [state_dim]`
|
||
/// - `dw_b / db_b` : `[state_dim, hidden_dim] / [state_dim]`
|
||
/// - `dw_c` : `[hidden_dim, state_dim]` (stored as flat slice;
|
||
/// layout matches the `W_c` weight)
|
||
/// - `dw_out / db_out` : `[1, hidden_dim] / [1]`
|
||
pub struct Mamba2BackwardGrads {
|
||
pub dw_in: GpuTensor,
|
||
pub db_in: GpuTensor,
|
||
pub dw_a: GpuTensor,
|
||
pub db_a: GpuTensor,
|
||
pub dw_b: GpuTensor,
|
||
pub db_b: GpuTensor,
|
||
pub dw_c: CudaSlice<f32>,
|
||
pub dw_out: GpuTensor,
|
||
pub db_out: GpuTensor,
|
||
}
|
||
|
||
/// Saved activations + per-stage tensors that the backward pass needs. Returned
|
||
/// by [`Mamba2Block::forward_train`]; consumed by [`Mamba2Block::backward`].
|
||
/// For inference-only ([`Mamba2Block::forward`]) the cache is computed but
|
||
/// immediately dropped.
|
||
pub struct Mamba2ForwardCache {
|
||
/// Input as reshaped 2-D `[N*K, in_dim]` — needed to backprop into W_in.
|
||
pub input_2d: GpuTensor,
|
||
/// Output of W_in projection, `[N*K, hidden_dim]`. Needed to backprop into
|
||
/// W_a / W_b and to compose dx/dx_prev for the W_in gradient.
|
||
pub x: GpuTensor,
|
||
/// A projection output `[N*K, state_dim]`. Backward kernel reads this to
|
||
/// replay the forward state path.
|
||
pub a_proj: GpuTensor,
|
||
/// B projection output `[N*K, state_dim]`. Same role as a_proj.
|
||
pub b_proj: GpuTensor,
|
||
/// Output of the scan kernel, `[N, hidden_dim]`. Needed to backprop into
|
||
/// W_out (and is also the input to the W_out forward GEMM).
|
||
pub h_enriched: GpuTensor,
|
||
}
|
||
|
||
/// Cache produced by [`Mamba2Block::forward_train_seq`] — variant of
|
||
/// [`Mamba2ForwardCache`] that exposes the SSM output at EVERY timestep
|
||
/// instead of just the final position. Used by ml-alpha when supervising
|
||
/// the model at every snapshot in a sequence (per-step heads + BCE).
|
||
pub struct Mamba2ForwardCacheSeq {
|
||
pub input_2d: GpuTensor,
|
||
pub x: GpuTensor,
|
||
pub a_proj: GpuTensor,
|
||
pub b_proj: GpuTensor,
|
||
/// Per-step enriched state, `[N, K, hidden_dim]`. Slot `[i, t, j]` is
|
||
/// `h_s2[i, j] + sum_s w_c[j, s] * x[i, t, s]` where x is the SSM
|
||
/// state AFTER step t (post-gate update).
|
||
pub h_enriched_seq: GpuTensor,
|
||
}
|
||
|
||
/// Pre-allocated forward intermediates for [`Mamba2Block::forward_train_seq_into`].
|
||
/// Holds all GpuTensors that the regular `forward_train_seq` allocates per
|
||
/// call (input_2d projection output `x`, A/B projections, h_s2 residual,
|
||
/// h_enriched_seq scan output). Constructed once per (n_batch, seq_len,
|
||
/// in_dim, hidden_dim, state_dim) tuple and reused — eliminates the
|
||
/// per-step `cudaMalloc` calls that block CUDA Graph capture.
|
||
///
|
||
/// `h_s2` is zero-initialised once and never written (no residual carry
|
||
/// from a prior chunk in the supervised path). The scan kernel reads it
|
||
/// as a constant addition to h_enriched_seq.
|
||
pub struct Mamba2BlockForwardScratch {
|
||
pub x: GpuTensor, // [n_rows, hidden_dim] (W_in output)
|
||
pub a_proj: GpuTensor, // [n_rows, state_dim]
|
||
pub b_proj: GpuTensor, // [n_rows, state_dim]
|
||
pub h_s2: GpuTensor, // [n_batch, hidden_dim] (zero residual)
|
||
pub h_enriched_seq: GpuTensor, // [n_batch, seq_len, hidden_dim]
|
||
pub n_batch: usize,
|
||
pub seq_len: usize,
|
||
pub in_dim: usize,
|
||
pub hidden_dim: usize,
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
impl Mamba2BlockForwardScratch {
|
||
pub fn new(
|
||
stream: &Arc<CudaStream>,
|
||
n_batch: usize,
|
||
seq_len: usize,
|
||
in_dim: usize,
|
||
hidden_dim: usize,
|
||
state_dim: usize,
|
||
) -> Result<Self> {
|
||
let n_rows = n_batch * seq_len;
|
||
Ok(Self {
|
||
x: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("fwd scratch x: {e}"))?,
|
||
a_proj: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||
.map_err(|e| anyhow!("fwd scratch a_proj: {e}"))?,
|
||
b_proj: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||
.map_err(|e| anyhow!("fwd scratch b_proj: {e}"))?,
|
||
h_s2: GpuTensor::zeros(&[n_batch, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("fwd scratch h_s2: {e}"))?,
|
||
h_enriched_seq: GpuTensor::zeros(&[n_batch, seq_len, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("fwd scratch h_enriched_seq: {e}"))?,
|
||
n_batch, seq_len, in_dim, hidden_dim, state_dim,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// Pre-allocated scratch for [`Mamba2Block::step_into`] (CRT Phase A0.5).
|
||
/// Sized for a single-snapshot forward pass (K=1) and holds the
|
||
/// PERSISTENT recurrent SSM register state across calls so the encoder
|
||
/// can advance by one event per call instead of re-running over a
|
||
/// K-window every time.
|
||
///
|
||
/// Layout differences vs [`Mamba2BlockForwardScratch`]:
|
||
/// - `x`, `a_proj`, `b_proj` are `[n_batch, ...]` instead of `[n_batch*K, ...]`
|
||
/// - `x_state` is new: `[n_batch, hidden_dim, state_dim]` — the per-(i, j)
|
||
/// thread register state read at start of each step kernel and written
|
||
/// back at the end. Zero-initialised on construction; updated in-place
|
||
/// by [`Mamba2Block::step_into`]. Use [`Self::reset_state`] to zero it
|
||
/// between sessions.
|
||
/// - `h_out` replaces `h_enriched_seq`: shape `[n_batch, hidden_dim]`,
|
||
/// the single-step enriched output (no K dimension).
|
||
///
|
||
/// `h_s2` stays zero-initialised (matches the supervised path's
|
||
/// residual-zero convention). Construct once per (n_batch, in_dim,
|
||
/// hidden_dim, state_dim) tuple; reuse for every forward_step call.
|
||
pub struct Mamba2BlockStepScratch {
|
||
/// W_in output `[n_batch, hidden_dim]` (single-row projection).
|
||
pub x: GpuTensor,
|
||
/// W_a output `[n_batch, state_dim]`.
|
||
pub a_proj: GpuTensor,
|
||
/// W_b output `[n_batch, state_dim]`.
|
||
pub b_proj: GpuTensor,
|
||
/// SSM register state `[n_batch, hidden_dim, state_dim]` — PERSISTENT
|
||
/// across step_into calls. Zero-initialised on `new()`. The step
|
||
/// kernel reads this at entry, advances by one step, writes back.
|
||
pub x_state: CudaSlice<f32>,
|
||
/// h_s2 residual `[n_batch, hidden_dim]` — kept zero by convention,
|
||
/// matches the supervised forward_train_seq_into path.
|
||
pub h_s2: GpuTensor,
|
||
/// Single-step enriched output `[n_batch, hidden_dim]`.
|
||
pub h_out: GpuTensor,
|
||
pub n_batch: usize,
|
||
pub in_dim: usize,
|
||
pub hidden_dim: usize,
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
impl Mamba2BlockStepScratch {
|
||
pub fn new(
|
||
stream: &Arc<CudaStream>,
|
||
n_batch: usize,
|
||
in_dim: usize,
|
||
hidden_dim: usize,
|
||
state_dim: usize,
|
||
) -> Result<Self> {
|
||
Ok(Self {
|
||
x: GpuTensor::zeros(&[n_batch, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("step scratch x: {e}"))?,
|
||
a_proj: GpuTensor::zeros(&[n_batch, state_dim], stream)
|
||
.map_err(|e| anyhow!("step scratch a_proj: {e}"))?,
|
||
b_proj: GpuTensor::zeros(&[n_batch, state_dim], stream)
|
||
.map_err(|e| anyhow!("step scratch b_proj: {e}"))?,
|
||
x_state: stream
|
||
.alloc_zeros::<f32>(n_batch * hidden_dim * state_dim)
|
||
.map_err(|e| anyhow!("step scratch x_state: {e}"))?,
|
||
h_s2: GpuTensor::zeros(&[n_batch, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("step scratch h_s2: {e}"))?,
|
||
h_out: GpuTensor::zeros(&[n_batch, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("step scratch h_out: {e}"))?,
|
||
n_batch, in_dim, hidden_dim, state_dim,
|
||
})
|
||
}
|
||
|
||
/// Zero the persistent SSM register state. Used for session-gap
|
||
/// resets — drop accumulated context and restart from a clean slate.
|
||
/// h_s2 stays at its construction zero; only x_state needs an
|
||
/// explicit reset because step_into writes back to it every call.
|
||
pub fn reset_state(&mut self, stream: &Arc<CudaStream>) -> Result<()> {
|
||
stream
|
||
.memset_zeros(&mut self.x_state)
|
||
.map_err(|e| anyhow!("step scratch reset: {e}"))?;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
impl Mamba2BlockStepScratch {
|
||
/// Test-only readback of the persistent x_state register buffer.
|
||
/// One-shot DtoH — acceptable in tests; NEVER call on the hot path.
|
||
pub fn read_x_state(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>> {
|
||
let n = self.x_state.len();
|
||
let mut host = vec![0.0_f32; n];
|
||
stream
|
||
.memcpy_dtoh(&self.x_state, &mut host)
|
||
.map_err(|e| anyhow!("read_x_state dtoh: {e}"))?;
|
||
stream
|
||
.synchronize()
|
||
.map_err(|e| anyhow!("read_x_state sync: {e}"))?;
|
||
Ok(host)
|
||
}
|
||
}
|
||
|
||
/// Pre-allocated outputs + intermediates for the full Mamba2 seq
|
||
/// backward — see [`Mamba2Block::backward_from_h_enriched_seq_full_into`].
|
||
/// Holds the cuBLAS linear-backward outputs (dw_in/db_in/dw_a/db_a/
|
||
/// dw_b/db_b plus the reshaped d_a_proj_2d / d_b_proj_2d / d_x_from_a
|
||
/// / d_x_from_b / d_x intermediates) so the backward pass does zero
|
||
/// per-step allocations.
|
||
///
|
||
/// W_out is unused in the seq path — dw_out / db_out are zero-init
|
||
/// shells of the right shape so the Mamba2AdamW step is a no-op for
|
||
/// those parameters.
|
||
pub struct Mamba2BackwardGradsBuffers {
|
||
pub d_a_proj_2d: GpuTensor, // [n_rows, state_dim] (cuBLAS-reshape of d_a_proj_flat)
|
||
pub d_b_proj_2d: GpuTensor, // [n_rows, state_dim]
|
||
pub d_x_from_a: GpuTensor, // [n_rows, hidden_dim]
|
||
pub d_x_from_b: GpuTensor, // [n_rows, hidden_dim]
|
||
pub d_x: GpuTensor, // [n_rows, hidden_dim]
|
||
pub dw_in: GpuTensor, // [hidden_dim, in_dim]
|
||
pub db_in: GpuTensor, // [hidden_dim]
|
||
pub dw_a: GpuTensor, // [state_dim, hidden_dim]
|
||
pub db_a: GpuTensor, // [state_dim]
|
||
pub dw_b: GpuTensor, // [state_dim, hidden_dim]
|
||
pub db_b: GpuTensor, // [state_dim]
|
||
pub dw_c: GpuTensor, // [hidden_dim, state_dim]
|
||
pub d_x_from_in: GpuTensor, // [n_rows, in_dim] (unused but allocated for w_in_into's dx_out)
|
||
pub dw_out: GpuTensor, // [1, hidden_dim] (zero)
|
||
pub db_out: GpuTensor, // [1] (zero)
|
||
pub n_batch: usize,
|
||
pub seq_len: usize,
|
||
pub in_dim: usize,
|
||
pub hidden_dim: usize,
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
impl Mamba2BackwardGradsBuffers {
|
||
pub fn new(
|
||
stream: &Arc<CudaStream>,
|
||
n_batch: usize,
|
||
seq_len: usize,
|
||
in_dim: usize,
|
||
hidden_dim: usize,
|
||
state_dim: usize,
|
||
) -> Result<Self> {
|
||
let n_rows = n_batch * seq_len;
|
||
Ok(Self {
|
||
d_a_proj_2d: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_a_proj_2d: {e}"))?,
|
||
d_b_proj_2d: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_b_proj_2d: {e}"))?,
|
||
d_x_from_a: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_x_from_a: {e}"))?,
|
||
d_x_from_b: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_x_from_b: {e}"))?,
|
||
d_x: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_x: {e}"))?,
|
||
dw_in: GpuTensor::zeros(&[hidden_dim, in_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads dw_in: {e}"))?,
|
||
db_in: GpuTensor::zeros(&[hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads db_in: {e}"))?,
|
||
dw_a: GpuTensor::zeros(&[state_dim, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads dw_a: {e}"))?,
|
||
db_a: GpuTensor::zeros(&[state_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads db_a: {e}"))?,
|
||
dw_b: GpuTensor::zeros(&[state_dim, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads dw_b: {e}"))?,
|
||
db_b: GpuTensor::zeros(&[state_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads db_b: {e}"))?,
|
||
dw_c: GpuTensor::zeros(&[hidden_dim, state_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads dw_c: {e}"))?,
|
||
d_x_from_in: GpuTensor::zeros(&[n_rows, in_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads d_x_from_in: {e}"))?,
|
||
dw_out: GpuTensor::zeros(&[1, hidden_dim], stream)
|
||
.map_err(|e| anyhow!("bwd grads dw_out: {e}"))?,
|
||
db_out: GpuTensor::zeros(&[1], stream)
|
||
.map_err(|e| anyhow!("bwd grads db_out: {e}"))?,
|
||
n_batch, seq_len, in_dim, hidden_dim, state_dim,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// Pre-allocated scratch for [`Mamba2Block::backward_from_h_enriched_seq_into`].
|
||
/// Holds the 4 BIG per-step scratch buffers (the ones at ~6MB each that
|
||
/// dominated the per-call `alloc_zeros` churn). Smaller buffers
|
||
/// (`d_a_proj_flat`, `d_b_proj_flat`, `dw_c`) remain per-call —
|
||
/// they're tiny and feed directly into ownership-transferred
|
||
/// LinearGrads outputs, where pre-allocation would force a
|
||
/// refactor of ml-core's cuBLAS wrappers without proportional gain.
|
||
///
|
||
/// Construct ONCE per (n_batch, seq_len, hidden_dim, state_dim) tuple.
|
||
pub struct Mamba2BackwardScratch {
|
||
pub d_a_per_channel: CudaSlice<f32>, // [N, sh2, K, state_d]
|
||
pub d_b_per_channel: CudaSlice<f32>, // [N, sh2, K, state_d]
|
||
pub d_w_c_per_sample: CudaSlice<f32>, // [N, sh2, state_d]
|
||
pub d_h_s2: CudaSlice<f32>, // [N, sh2]
|
||
pub n_batch: usize,
|
||
pub seq_len: usize,
|
||
pub hidden_dim: usize,
|
||
pub state_dim: usize,
|
||
}
|
||
|
||
impl Mamba2BackwardScratch {
|
||
pub fn new(
|
||
stream: &Arc<CudaStream>,
|
||
n_batch: usize,
|
||
seq_len: usize,
|
||
hidden_dim: usize,
|
||
state_dim: usize,
|
||
) -> Result<Self> {
|
||
let per_chan_n = n_batch * hidden_dim * seq_len * state_dim;
|
||
let per_sample_n = n_batch * hidden_dim * state_dim;
|
||
Ok(Self {
|
||
d_a_per_channel: stream.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("scratch d_a_per_channel: {e}"))?,
|
||
d_b_per_channel: stream.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("scratch d_b_per_channel: {e}"))?,
|
||
d_w_c_per_sample: stream.alloc_zeros::<f32>(per_sample_n)
|
||
.map_err(|e| anyhow!("scratch d_w_c_per_sample: {e}"))?,
|
||
d_h_s2: stream.alloc_zeros::<f32>(n_batch * hidden_dim)
|
||
.map_err(|e| anyhow!("scratch d_h_s2: {e}"))?,
|
||
n_batch, seq_len, hidden_dim, state_dim,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// GPU-resident Mamba2 sequence block.
|
||
///
|
||
/// Owns all parameter tensors (`W_in`, `W_a`, `W_b`, `W_c`, `W_out`) on the
|
||
/// CUDA device. The forward + backward kernels are loaded once at
|
||
/// construction from the precompiled cubin.
|
||
pub struct Mamba2Block {
|
||
pub config: Mamba2BlockConfig,
|
||
pub stream: Arc<CudaStream>,
|
||
/// Cached raw CUstream handle for [`raw_launch`] — avoids the
|
||
/// Arc deref + method call on every kernel launch.
|
||
pub raw_stream: CUstream,
|
||
|
||
// ── Parameters ────────────────────────────────────────────────────
|
||
/// Input projection: `in_dim → hidden_dim`. Cuts the snapshot row down
|
||
/// to model space before the selective scan.
|
||
pub w_in: OwnedGpuLinear,
|
||
/// A-projection (gate): `hidden_dim → state_dim`. The kernel applies
|
||
/// sigmoid to produce per-state gates (`a_proj`).
|
||
pub w_a: OwnedGpuLinear,
|
||
/// B-projection (input contribution): `hidden_dim → state_dim`. Added
|
||
/// to the carried state each step (`b_proj`).
|
||
pub w_b: OwnedGpuLinear,
|
||
/// Context-mix weights `W_c` in the kernel signature, shape
|
||
/// `[hidden_dim, state_dim]`. Multiplies the final state to produce the
|
||
/// per-position context that's added to `h_s2`. Stored as a raw
|
||
/// `CudaSlice<f32>` (no bias) because the kernel reads it directly.
|
||
pub w_c: CudaSlice<f32>,
|
||
/// Output projection: `hidden_dim → 1`. Reduces the enriched hidden
|
||
/// state to the binary direction logit.
|
||
pub w_out: OwnedGpuLinear,
|
||
|
||
// ── Kernel handles ────────────────────────────────────────────────
|
||
_module: Arc<CudaModule>,
|
||
pub kernel_fwd: CudaFunction,
|
||
pub kernel_bwd: CudaFunction,
|
||
/// Per-step variant of the forward scan — writes h_enriched at every
|
||
/// timestep (used by per-position supervision in PerceptionTrainer).
|
||
pub kernel_fwd_seq: CudaFunction,
|
||
/// SINGLE-STEP variant of the forward scan — reads and writes the
|
||
/// recurrent SSM register state to/from device memory so the SSM can
|
||
/// be advanced one snapshot per launch across many forward_step calls.
|
||
/// Used by [`Mamba2Block::step_into`] for event-rate inference
|
||
/// (CRT Phase A0.5). The supervised path keeps using
|
||
/// `kernel_fwd_seq` over a full K-window.
|
||
pub kernel_fwd_step: CudaFunction,
|
||
/// Per-step variant of the backward scan — accepts d_h_enriched_seq
|
||
/// and injects gradient contributions into d_state at every step.
|
||
pub kernel_bwd_seq: CudaFunction,
|
||
/// Reduces `d_a_per_channel` or `d_b_per_channel` (same kernel, two
|
||
/// call sites with different I/O pointers).
|
||
pub kernel_reduce_d_proj: CudaFunction,
|
||
/// Reduces `d_w_c_per_sample[N, sh2, state_d]` across N.
|
||
pub kernel_reduce_d_w_c: CudaFunction,
|
||
|
||
/// AdamW step kernel — applied per parameter tensor.
|
||
pub kernel_adamw: CudaFunction,
|
||
|
||
/// cuBLAS handle for the four linear projections (input / A / B / output).
|
||
pub cublas: CudaBlas,
|
||
/// Pre-allocated cuBLAS workspace. cuBLAS lazily allocates internal
|
||
/// workspace per gemm call when none is set, which is incompatible
|
||
/// with CUDA Graph stream capture (cudaMalloc inside the captured
|
||
/// region trips `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`). Setting a
|
||
/// user-managed workspace via `cublasSetWorkspace_v2` removes the
|
||
/// per-call alloc; the buffer stays alive for the lifetime of the
|
||
/// block, large enough for all gemm sizes used by the projections.
|
||
_cublas_workspace: CudaSlice<u8>,
|
||
}
|
||
|
||
impl Mamba2Block {
|
||
/// Construct a fresh Mamba2 block with Xavier-initialised projections
|
||
/// and the SSM scan kernels loaded from the precompiled cubin.
|
||
pub fn new(config: Mamba2BlockConfig, stream: Arc<CudaStream>) -> Result<Self> {
|
||
config.validate()?;
|
||
|
||
// ── Load the precompiled cubin and resolve kernel symbols. ────
|
||
// The cubin is produced by `build.rs` at compile time; no nvrtc.
|
||
static CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_alpha_kernel.cubin"));
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(CUBIN.to_vec())
|
||
.map_err(|e| anyhow!("Mamba2Block: cubin load failed: {e}"))?;
|
||
let kernel_fwd = module
|
||
.load_function("mamba2_alpha_scan_fwd")
|
||
.map_err(|e| anyhow!("Mamba2Block: forward kernel symbol resolve: {e}"))?;
|
||
let kernel_bwd = module
|
||
.load_function("mamba2_alpha_scan_bwd")
|
||
.map_err(|e| anyhow!("Mamba2Block: backward kernel symbol resolve: {e}"))?;
|
||
let kernel_fwd_seq = module
|
||
.load_function("mamba2_alpha_scan_fwd_seq")
|
||
.map_err(|e| anyhow!("Mamba2Block: per-step forward kernel resolve: {e}"))?;
|
||
let kernel_fwd_step = module
|
||
.load_function("mamba2_alpha_scan_fwd_step")
|
||
.map_err(|e| anyhow!("Mamba2Block: single-step forward kernel resolve: {e}"))?;
|
||
let kernel_bwd_seq = module
|
||
.load_function("mamba2_alpha_scan_bwd_seq")
|
||
.map_err(|e| anyhow!("Mamba2Block: per-step backward kernel resolve: {e}"))?;
|
||
let kernel_reduce_d_proj = module
|
||
.load_function("mamba2_alpha_reduce_d_proj")
|
||
.map_err(|e| anyhow!("Mamba2Block: d_proj reduction kernel resolve: {e}"))?;
|
||
let kernel_reduce_d_w_c = module
|
||
.load_function("mamba2_alpha_reduce_d_w_c")
|
||
.map_err(|e| anyhow!("Mamba2Block: d_w_c reduction kernel resolve: {e}"))?;
|
||
let kernel_adamw = module
|
||
.load_function("mamba2_alpha_adamw_step")
|
||
.map_err(|e| anyhow!("Mamba2Block: AdamW kernel resolve: {e}"))?;
|
||
|
||
// cuBLAS handle for the four GEMM projections, with a pre-
|
||
// allocated user-managed workspace so subsequent gemm calls
|
||
// don't trigger internal allocations during CUDA Graph capture.
|
||
// NVIDIA recommends 4 MiB minimum for general use, larger for
|
||
// batched ops on Ampere+; 8 MiB is comfortable headroom for
|
||
// our gemm shapes (largest is hidden_dim×in_dim ≈ 128×32).
|
||
let cublas = CudaBlas::new(Arc::clone(&stream))
|
||
.map_err(|e| anyhow!("Mamba2Block: cuBLAS init failed: {e}"))?;
|
||
// Determinism foundation (spec 2026-06-02 §2.B):
|
||
// Disable non-deterministic GEMM algorithms. Per the Phase 2.2
|
||
// mamba2 investigation, the 3 backward GEMMs in
|
||
// backward_from_h_enriched_seq_full_into were the root cause of
|
||
// the eval-pnl drift between same-seed runs.
|
||
let cublas = crate::cublas_determinism::apply_deterministic_math_mode(cublas)
|
||
.map_err(|e| anyhow!("Mamba2Block: cublas determinism setup: {e}"))?;
|
||
const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;
|
||
let cublas_workspace = stream
|
||
.alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
|
||
.map_err(|e| anyhow!("Mamba2Block: cuBLAS workspace alloc: {e}"))?;
|
||
unsafe {
|
||
let ws_ptr = cublas_workspace.raw_ptr();
|
||
cudarc::cublas::sys::cublasSetWorkspace_v2(
|
||
*cublas.handle(),
|
||
ws_ptr as *mut std::ffi::c_void,
|
||
CUBLAS_WORKSPACE_BYTES,
|
||
)
|
||
.result()
|
||
.map_err(|e| anyhow!("Mamba2Block: cublasSetWorkspace_v2: {e:?}"))?;
|
||
}
|
||
|
||
// ── Parameter allocation + initialisation. ──────────────────────
|
||
// All on GPU; init helpers in ml-core::cuda_autograd::init use
|
||
// pinned host buffers for the Glorot/Xavier seed transfer.
|
||
let w_in = OwnedGpuLinear::xavier("mamba2.w_in", config.in_dim, config.hidden_dim, &stream)
|
||
.map_err(|e| anyhow!("init w_in: {e}"))?;
|
||
let w_a = OwnedGpuLinear::xavier(
|
||
"mamba2.w_a", config.hidden_dim, config.state_dim, &stream,
|
||
)
|
||
.map_err(|e| anyhow!("init w_a: {e}"))?;
|
||
let w_b = OwnedGpuLinear::xavier(
|
||
"mamba2.w_b", config.hidden_dim, config.state_dim, &stream,
|
||
)
|
||
.map_err(|e| anyhow!("init w_b: {e}"))?;
|
||
|
||
// W_c is a raw kernel-fed weight (no bias). Initialise via
|
||
// OwnedGpuLinear::xavier then keep only the weight slice.
|
||
let w_c_linear = OwnedGpuLinear::xavier(
|
||
"mamba2.w_c", config.state_dim, config.hidden_dim, &stream,
|
||
)
|
||
.map_err(|e| anyhow!("init w_c: {e}"))?;
|
||
let w_c = w_c_linear.weight; // discard bias for kernel-direct use
|
||
|
||
let w_out = OwnedGpuLinear::xavier("mamba2.w_out", config.hidden_dim, 1, &stream)
|
||
.map_err(|e| anyhow!("init w_out: {e}"))?;
|
||
|
||
let raw_stream = stream.cu_stream();
|
||
Ok(Self {
|
||
config,
|
||
stream,
|
||
raw_stream,
|
||
w_in,
|
||
w_a,
|
||
w_b,
|
||
w_c,
|
||
w_out,
|
||
_module: module,
|
||
kernel_fwd,
|
||
kernel_bwd,
|
||
kernel_fwd_seq,
|
||
kernel_fwd_step,
|
||
kernel_bwd_seq,
|
||
kernel_reduce_d_proj,
|
||
kernel_reduce_d_w_c,
|
||
kernel_adamw,
|
||
cublas,
|
||
_cublas_workspace: cublas_workspace,
|
||
})
|
||
}
|
||
|
||
// ── Forward pass ───────────────────────────────────────────────────
|
||
|
||
/// GPU-native forward pass with full activation cache returned alongside
|
||
/// the output logits. Use this when training (backward needs the cache).
|
||
///
|
||
/// Steps (all on GPU; no host roundtrip):
|
||
/// 1. `x = input @ Wᵢₙ.T + bᵢₙ` (cuBLAS sgemm + bias add)
|
||
/// 2. `a = x @ Wₐ.T + bₐ` (cuBLAS sgemm + bias add)
|
||
/// 3. `b = x @ Wᵦ.T + bᵦ` (cuBLAS sgemm + bias add)
|
||
/// 4. zero-init `h_s2[B, hidden_dim]`, scratch `h_enriched[B, hidden_dim]`
|
||
/// 5. launch `mamba2_alpha_scan_fwd(a, b, w_c, h_s2, h_enriched, B, K, H, S)`
|
||
/// 6. `logit = h_enriched @ Wₒᵤₜ.T + bₒᵤₜ` (cuBLAS sgemm + bias add)
|
||
pub fn forward_train(&self, input: &GpuTensor) -> Result<(GpuTensor, Mamba2ForwardCache)> {
|
||
let c = &self.config;
|
||
let n_batch = match input.shape() {
|
||
[b, k, d] if *k == c.seq_len && *d == c.in_dim => *b,
|
||
shape => {
|
||
return Err(anyhow!(
|
||
"Mamba2Block::forward_train: expected input shape [B, {}, {}], got {:?}",
|
||
c.seq_len, c.in_dim, shape
|
||
));
|
||
}
|
||
};
|
||
let n_rows = n_batch * c.seq_len; // flattened over time for the per-step projections
|
||
|
||
// Reshape input to [N*K, in_dim] (logical reshape — same storage).
|
||
let input_2d = GpuTensor::new(
|
||
input.cuda_data().clone(),
|
||
vec![n_rows, c.in_dim],
|
||
)
|
||
.map_err(|e| anyhow!("reshape input → 2D: {e}"))?;
|
||
|
||
// (1) input projection: x = input @ W_in.T + b_in → [N*K, hidden_dim]
|
||
let (x, _) = self.w_in.inner.forward_with_slices(
|
||
&input_2d, &self.w_in.weight, &self.w_in.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_in forward: {e}"))?;
|
||
|
||
// (2) a_proj = x @ W_a.T + b_a → [N*K, state_dim]
|
||
let (a_proj, _) = self.w_a.inner.forward_with_slices(
|
||
&x, &self.w_a.weight, &self.w_a.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_a forward: {e}"))?;
|
||
|
||
// (3) b_proj = x @ W_b.T + b_b → [N*K, state_dim]
|
||
let (b_proj, _) = self.w_b.inner.forward_with_slices(
|
||
&x, &self.w_b.weight, &self.w_b.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_b forward: {e}"))?;
|
||
|
||
// (4) initial hidden h_s2 (zero residual; no carry from previous chunk)
|
||
// and output buffer h_enriched, both [N, hidden_dim].
|
||
let h_s2 = GpuTensor::zeros(&[n_batch, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc h_s2: {e}"))?;
|
||
let mut h_enriched = GpuTensor::zeros(&[n_batch, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc h_enriched: {e}"))?;
|
||
|
||
// (5) launch the scan kernel. Grid: (N, ceil(H/32)), Block: 32.
|
||
let block_threads: u32 = 32;
|
||
let grid_x: u32 = n_batch as u32;
|
||
let grid_y: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(h_s2.cuda_data().raw_ptr());
|
||
args.push_ptr(h_enriched.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_fwd.cu_function(),
|
||
(grid_x, grid_y, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("mamba2_alpha_scan_fwd launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// (6) output projection: logit = h_enriched @ W_out.T + b_out → [N, 1]
|
||
let (logit, _) = self.w_out.inner.forward_with_slices(
|
||
&h_enriched, &self.w_out.weight, &self.w_out.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_out forward: {e}"))?;
|
||
|
||
let cache = Mamba2ForwardCache { input_2d, x, a_proj, b_proj, h_enriched };
|
||
Ok((logit, cache))
|
||
}
|
||
|
||
/// Forward inference only — runs `forward_train` and discards the cache.
|
||
/// Use when you only need predictions (eval, calibration set scoring).
|
||
#[inline]
|
||
pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor> {
|
||
self.forward_train(input).map(|(logit, _cache)| logit)
|
||
}
|
||
|
||
// ── Backward pass ─────────────────────────────────────────────────
|
||
|
||
/// GPU-native backward pass. `d_logit` is the upstream gradient
|
||
/// `[n_batch, 1]` flowing back from the loss; `cache` is the
|
||
/// `Mamba2ForwardCache` returned by [`forward_train`]. Returns all
|
||
/// nine parameter gradients in a [`Mamba2BackwardGrads`].
|
||
///
|
||
/// Chain (mirrors forward in reverse):
|
||
/// 6′. W_out backward : d_logit + h_enriched_act → d_h_enriched, dw_out, db_out
|
||
/// 5′. scan backward : (a_proj, b_proj, d_h_enriched, W_c)
|
||
/// → d_a_per_channel, d_b_per_channel,
|
||
/// d_w_c_per_sample, d_h_s2 (=d_h_enriched)
|
||
/// ↳ reduce d_a_per_channel across j → d_a_proj [N, K, state]
|
||
/// ↳ reduce d_b_per_channel across j → d_b_proj [N, K, state]
|
||
/// ↳ reduce d_w_c_per_sample across i → dw_c [sh2, state]
|
||
/// 3′. W_b backward : d_b_proj[N*K, state] + x_act → d_x_from_b, dw_b, db_b
|
||
/// 2′. W_a backward : d_a_proj[N*K, state] + x_act → d_x_from_a, dw_a, db_a
|
||
/// ↳ d_x = d_x_from_a + d_x_from_b
|
||
/// 1′. W_in backward : d_x[N*K, hidden] + input_act → (d_input dropped), dw_in, db_in
|
||
pub fn backward(
|
||
&self,
|
||
cache: &Mamba2ForwardCache,
|
||
d_logit: &GpuTensor,
|
||
) -> Result<Mamba2BackwardGrads> {
|
||
let c = &self.config;
|
||
let n_batch = cache.h_enriched.shape()[0];
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
if d_logit.shape() != [n_batch, 1] {
|
||
return Err(anyhow!(
|
||
"Mamba2Block::backward: d_logit shape {:?} != [{}, 1]",
|
||
d_logit.shape(), n_batch
|
||
));
|
||
}
|
||
|
||
// ── 6′. Output projection backward ────────────────────────────
|
||
let h_enriched_act = LinearActivations { input: cache.h_enriched.clone() };
|
||
let LinearGrads { dw: dw_out, db: db_out, dx: d_h_enriched } = self
|
||
.w_out
|
||
.inner
|
||
.backward_with_slices(d_logit, &h_enriched_act, &self.w_out.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_out backward: {e}"))?;
|
||
|
||
// ── 5′. Allocate scan-backward scratch buffers ────────────────
|
||
let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim;
|
||
let per_sample_n = n_batch * c.hidden_dim * c.state_dim;
|
||
let d_a_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_a_per_channel ({} floats): {e}", per_chan_n))?;
|
||
let d_b_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?;
|
||
let d_w_c_per_sample = self.stream
|
||
.alloc_zeros::<f32>(per_sample_n)
|
||
.map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?;
|
||
let d_h_s2 = self.stream
|
||
.alloc_zeros::<f32>(n_batch * c.hidden_dim)
|
||
.map_err(|e| anyhow!("alloc d_h_s2: {e}"))?;
|
||
|
||
// ── Launch scan backward kernel ───────────────────────────────
|
||
let block_threads: u32 = 32;
|
||
let grid_y_h: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(cache.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(cache.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(d_h_enriched.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(d_h_s2.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_bwd.cu_function(),
|
||
(n_batch as u32, grid_y_h, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("mamba2_alpha_scan_bwd launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Reduce d_a_per_channel and d_b_per_channel across j ───────
|
||
let red_grid_z: u32 =
|
||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let red_cfg = (
|
||
(n_batch as u32, c.seq_len as u32, red_grid_z),
|
||
(block_threads, 1, 1),
|
||
);
|
||
let d_a_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?;
|
||
let d_b_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_a_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce_d_proj (a) launch: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce_d_proj (b) launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Reduce d_w_c_per_sample across i → dw_c[sh2, state] ───────
|
||
let red_w_c_grid_y: u32 = red_grid_z;
|
||
let red_w_c_cfg = (
|
||
(c.hidden_dim as u32, red_w_c_grid_y, 1),
|
||
(block_threads, 1, 1),
|
||
);
|
||
let dw_c: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc dw_c: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(dw_c.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_w_c.cu_function(),
|
||
red_w_c_cfg.0, red_w_c_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce_d_w_c launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── 3′. W_b backward (uses x_act = cache.x) ───────────────────
|
||
let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_b_proj: {e}"))?;
|
||
let x_act = LinearActivations { input: cache.x.clone() };
|
||
let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
|
||
.w_b
|
||
.inner
|
||
.backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_b backward: {e}"))?;
|
||
|
||
// ── 2′. W_a backward ─────────────────────────────────────────
|
||
let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_a_proj: {e}"))?;
|
||
let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
|
||
.w_a
|
||
.inner
|
||
.backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_a backward: {e}"))?;
|
||
|
||
// d_x = d_x_from_a + d_x_from_b
|
||
let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
|
||
.map_err(|e| anyhow!("sum d_x branches: {e}"))?;
|
||
|
||
// ── 1′. W_in backward ─────────────────────────────────────────
|
||
let input_act = LinearActivations { input: cache.input_2d.clone() };
|
||
let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
|
||
.w_in
|
||
.inner
|
||
.backward_with_slices(&d_x, &input_act, &self.w_in.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_in backward: {e}"))?;
|
||
|
||
Ok(Mamba2BackwardGrads {
|
||
dw_in, db_in,
|
||
dw_a, db_a,
|
||
dw_b, db_b,
|
||
dw_c,
|
||
dw_out, db_out,
|
||
})
|
||
}
|
||
|
||
/// Phase E.4.A T10 (2026-05-15): backward variant that BYPASSES the
|
||
/// W_out projection. Phase E uses Mamba2's `h_enriched` directly as
|
||
/// input to a downstream classifier (C51 head); the `d_logit`-based
|
||
/// `backward` is unsuitable because Phase E never computes a
|
||
/// `d_logit` (there is no scalar logit in its loss path).
|
||
///
|
||
/// `d_h_enriched` MUST have shape `[B, hidden_dim]` matching
|
||
/// `cache.h_enriched`. Returns `Mamba2BackwardGrads` where
|
||
/// `dw_out` and `db_out` are zero-initialised tensors of the
|
||
/// correct shapes (AdamW step on zero grad is a no-op with the
|
||
/// correct moment-decay; W_out parameters effectively freeze —
|
||
/// correct semantics since Phase E never uses them).
|
||
pub fn backward_from_h_enriched(
|
||
&self,
|
||
cache: &Mamba2ForwardCache,
|
||
d_h_enriched: &GpuTensor,
|
||
) -> Result<Mamba2BackwardGrads> {
|
||
let c = &self.config;
|
||
let n_batch = cache.h_enriched.shape()[0];
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
if d_h_enriched.shape() != [n_batch, c.hidden_dim] {
|
||
return Err(anyhow!(
|
||
"Mamba2Block::backward_from_h_enriched: \
|
||
d_h_enriched shape {:?} != [{}, {}]",
|
||
d_h_enriched.shape(), n_batch, c.hidden_dim
|
||
));
|
||
}
|
||
|
||
// ── 6′ SKIPPED. Phase E doesn't go through W_out. ────────────
|
||
// (Zero w_out grads are constructed at return.)
|
||
|
||
// ── 5′. Allocate scan-backward scratch buffers ────────────────
|
||
let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim;
|
||
let per_sample_n = n_batch * c.hidden_dim * c.state_dim;
|
||
let d_a_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_a_per_channel: {e}"))?;
|
||
let d_b_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?;
|
||
let d_w_c_per_sample = self.stream
|
||
.alloc_zeros::<f32>(per_sample_n)
|
||
.map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?;
|
||
let d_h_s2 = self.stream
|
||
.alloc_zeros::<f32>(n_batch * c.hidden_dim)
|
||
.map_err(|e| anyhow!("alloc d_h_s2: {e}"))?;
|
||
|
||
let block_threads: u32 = 32;
|
||
let grid_y_h: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(cache.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(cache.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(d_h_enriched.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(d_h_s2.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_bwd.cu_function(),
|
||
(n_batch as u32, grid_y_h, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("mamba2_alpha_scan_bwd launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Reduce d_a_per_channel and d_b_per_channel ───────────────
|
||
let red_grid_z: u32 =
|
||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let red_cfg = (
|
||
(n_batch as u32, c.seq_len as u32, red_grid_z),
|
||
(block_threads, 1, 1),
|
||
);
|
||
let d_a_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?;
|
||
let d_b_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_a_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_a_proj: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_b_proj: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Reduce d_w_c_per_sample → dw_c ──────────────────────────
|
||
let dw_c: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc dw_c: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(dw_c.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_w_c.cu_function(),
|
||
(c.hidden_dim as u32, red_grid_z, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce dw_c: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── 3′. W_b backward ─────────────────────────────────────────
|
||
let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_b_proj: {e}"))?;
|
||
let x_act = LinearActivations { input: cache.x.clone() };
|
||
let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
|
||
.w_b.inner
|
||
.backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_b backward: {e}"))?;
|
||
|
||
// ── 2′. W_a backward ─────────────────────────────────────────
|
||
let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_a_proj: {e}"))?;
|
||
let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
|
||
.w_a.inner
|
||
.backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_a backward: {e}"))?;
|
||
|
||
let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
|
||
.map_err(|e| anyhow!("sum d_x branches: {e}"))?;
|
||
|
||
// ── 1′. W_in backward ────────────────────────────────────────
|
||
let input_act = LinearActivations { input: cache.input_2d.clone() };
|
||
let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
|
||
.w_in.inner
|
||
.backward_with_slices(&d_x, &input_act, &self.w_in.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_in backward: {e}"))?;
|
||
|
||
// ── 6′ replacement: zero W_out grads ────────────────────────
|
||
// AdamW will see zero grad → effectively freezes w_out
|
||
// parameters. Correct semantics — Phase E doesn't use them.
|
||
let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero dw_out: {e}"))?;
|
||
let db_out = GpuTensor::zeros(&[1], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero db_out: {e}"))?;
|
||
|
||
Ok(Mamba2BackwardGrads {
|
||
dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c, dw_out, db_out,
|
||
})
|
||
}
|
||
|
||
/// Per-step variant of [`Mamba2Block::forward_train`] — emits the SSM
|
||
/// output at EVERY timestep instead of only the final position. Used
|
||
/// by ml-alpha PerceptionTrainer to supervise the model at every
|
||
/// snapshot in a sequence (denser gradient signal than final-step
|
||
/// only).
|
||
///
|
||
/// Skips the W_out projection — callers route `h_enriched_seq`
|
||
/// directly into a downstream head. Caller pairs this with
|
||
/// [`Mamba2Block::backward_from_h_enriched_seq`] for the backward
|
||
/// chain.
|
||
pub fn forward_train_seq(
|
||
&self,
|
||
input: &GpuTensor,
|
||
) -> Result<(GpuTensor, Mamba2ForwardCacheSeq)> {
|
||
let c = &self.config;
|
||
let n_batch = match input.shape() {
|
||
[b, k, d] if *k == c.seq_len && *d == c.in_dim => *b,
|
||
shape => {
|
||
return Err(anyhow!(
|
||
"Mamba2Block::forward_train_seq: expected [B, {}, {}], got {:?}",
|
||
c.seq_len, c.in_dim, shape
|
||
));
|
||
}
|
||
};
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
let input_2d = GpuTensor::new(input.cuda_data().clone(), vec![n_rows, c.in_dim])
|
||
.map_err(|e| anyhow!("reshape input → 2D: {e}"))?;
|
||
|
||
let (x, _) = self.w_in.inner.forward_with_slices(
|
||
&input_2d, &self.w_in.weight, &self.w_in.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_in forward: {e}"))?;
|
||
|
||
let (a_proj, _) = self.w_a.inner.forward_with_slices(
|
||
&x, &self.w_a.weight, &self.w_a.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_a forward: {e}"))?;
|
||
|
||
let (b_proj, _) = self.w_b.inner.forward_with_slices(
|
||
&x, &self.w_b.weight, &self.w_b.bias, &self.cublas, &self.stream,
|
||
).map_err(|e| anyhow!("w_b forward: {e}"))?;
|
||
|
||
let h_s2 = GpuTensor::zeros(&[n_batch, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc h_s2: {e}"))?;
|
||
let mut h_enriched_seq = GpuTensor::zeros(
|
||
&[n_batch, c.seq_len, c.hidden_dim],
|
||
&self.stream,
|
||
)
|
||
.map_err(|e| anyhow!("alloc h_enriched_seq: {e}"))?;
|
||
|
||
let block_threads: u32 = 32;
|
||
let grid_y: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(h_s2.cuda_data().raw_ptr());
|
||
args.push_ptr(h_enriched_seq.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_fwd_seq.cu_function(),
|
||
(n_batch as u32, grid_y, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_fwd_seq launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
let cache = Mamba2ForwardCacheSeq {
|
||
input_2d, x, a_proj, b_proj,
|
||
h_enriched_seq: h_enriched_seq.clone(),
|
||
};
|
||
Ok((h_enriched_seq, cache))
|
||
}
|
||
|
||
/// Zero-allocation variant of [`forward_train_seq`] — writes all
|
||
/// intermediates into caller-provided pre-allocated buffers
|
||
/// ([`Mamba2BlockForwardScratch`]). No `GpuTensor::zeros` / cuBLAS
|
||
/// alloc per call, no `cudaMalloc`. The scratch's `h_enriched_seq`
|
||
/// IS the output (no fresh tensor returned).
|
||
///
|
||
/// Designed for the PerceptionTrainer hot path and CUDA Graph
|
||
/// capture. The scratch's contents are valid until the next
|
||
/// `forward_train_seq_into` call.
|
||
pub fn forward_train_seq_into(
|
||
&self,
|
||
input: &GpuTensor,
|
||
scratch: &mut Mamba2BlockForwardScratch,
|
||
) -> Result<()> {
|
||
let c = &self.config;
|
||
let n_batch = match input.shape() {
|
||
[b, k, d] if *k == c.seq_len && *d == c.in_dim => *b,
|
||
shape => {
|
||
return Err(anyhow!(
|
||
"forward_train_seq_into: expected [B, {}, {}], got {:?}",
|
||
c.seq_len, c.in_dim, shape
|
||
));
|
||
}
|
||
};
|
||
anyhow::ensure!(
|
||
scratch.n_batch == n_batch
|
||
&& scratch.seq_len == c.seq_len
|
||
&& scratch.in_dim == c.in_dim
|
||
&& scratch.hidden_dim == c.hidden_dim
|
||
&& scratch.state_dim == c.state_dim,
|
||
"fwd scratch shape mismatch: expected ({},{},{},{},{}) got ({},{},{},{},{})",
|
||
n_batch, c.seq_len, c.in_dim, c.hidden_dim, c.state_dim,
|
||
scratch.n_batch, scratch.seq_len, scratch.in_dim,
|
||
scratch.hidden_dim, scratch.state_dim
|
||
);
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
// 1. x = input @ W_in.T + b_in — input is [B, K, in_dim], viewed
|
||
// as [B*K, in_dim] via raw slice + batch dim (no GpuTensor wrap,
|
||
// no CudaSlice::clone — cuMemAlloc inside capture is forbidden).
|
||
self.w_in.inner.forward_with_slices_into(
|
||
input.cuda_data(), n_rows,
|
||
&self.w_in.weight, &self.w_in.bias,
|
||
&self.cublas, &self.stream, &mut scratch.x,
|
||
).map_err(|e| anyhow!("w_in fwd_into: {e}"))?;
|
||
|
||
// 2. a_proj = x @ W_a.T + b_a.
|
||
self.w_a.inner.forward_with_slices_into(
|
||
scratch.x.cuda_data(), n_rows,
|
||
&self.w_a.weight, &self.w_a.bias,
|
||
&self.cublas, &self.stream, &mut scratch.a_proj,
|
||
).map_err(|e| anyhow!("w_a fwd_into: {e}"))?;
|
||
|
||
// 3. b_proj = x @ W_b.T + b_b.
|
||
self.w_b.inner.forward_with_slices_into(
|
||
scratch.x.cuda_data(), n_rows,
|
||
&self.w_b.weight, &self.w_b.bias,
|
||
&self.cublas, &self.stream, &mut scratch.b_proj,
|
||
).map_err(|e| anyhow!("w_b fwd_into: {e}"))?;
|
||
|
||
// 4. scan_fwd_seq → scratch.h_enriched_seq.
|
||
// h_s2 stays zero from construction (never written in supervised path).
|
||
let block_threads: u32 = 32;
|
||
let grid_y: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(scratch.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(scratch.h_s2.cuda_data().raw_ptr());
|
||
args.push_ptr(scratch.h_enriched_seq.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_fwd_seq.cu_function(),
|
||
(n_batch as u32, grid_y, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_fwd_seq_into launch: {e:?}"))?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// CRT Phase A0.5: single-step incremental forward pass.
|
||
///
|
||
/// Advances the SSM register state (`scratch.x_state`) by exactly one
|
||
/// timestep using the snapshot held in `input` (shape `[n_batch, in_dim]`).
|
||
/// Writes the enriched output to `scratch.h_out` (shape
|
||
/// `[n_batch, hidden_dim]`). The recurrent state in `scratch.x_state`
|
||
/// is updated in-place so the next call continues from the new state.
|
||
///
|
||
/// Sequence equivalence: calling step_into N times on snapshots
|
||
/// `(x_0, x_1, … x_{N-1})` produces the same h_out at step N-1 as
|
||
/// `forward_train_seq_into` over the same N snapshots with K=N would
|
||
/// produce at the final timestep position (within float-summation
|
||
/// tolerance — the per-step kernel does the same arithmetic in the
|
||
/// same order, the only differences are GEMM-batch granularity and
|
||
/// the absence of intermediate per-step h_enriched writes).
|
||
pub fn step_into(
|
||
&self,
|
||
input: &GpuTensor,
|
||
scratch: &mut Mamba2BlockStepScratch,
|
||
) -> Result<()> {
|
||
let c = &self.config;
|
||
let n_batch = match input.shape() {
|
||
[b, d] if *d == c.in_dim => *b,
|
||
// Accept [B, 1, in_dim] too — keeps callers free to use the
|
||
// same staging tensor shape as forward_train_seq_into.
|
||
[b, k, d] if *k == 1 && *d == c.in_dim => *b,
|
||
shape => {
|
||
return Err(anyhow!(
|
||
"step_into: expected input shape [B, {0}] or [B, 1, {0}], got {1:?}",
|
||
c.in_dim, shape
|
||
));
|
||
}
|
||
};
|
||
anyhow::ensure!(
|
||
scratch.n_batch == n_batch
|
||
&& scratch.in_dim == c.in_dim
|
||
&& scratch.hidden_dim == c.hidden_dim
|
||
&& scratch.state_dim == c.state_dim,
|
||
"step scratch shape mismatch: expected ({},{},{},{}) got ({},{},{},{})",
|
||
n_batch, c.in_dim, c.hidden_dim, c.state_dim,
|
||
scratch.n_batch, scratch.in_dim, scratch.hidden_dim, scratch.state_dim
|
||
);
|
||
|
||
// 1. x = input @ W_in.T + b_in — single row per batch entry.
|
||
self.w_in.inner.forward_with_slices_into(
|
||
input.cuda_data(), n_batch,
|
||
&self.w_in.weight, &self.w_in.bias,
|
||
&self.cublas, &self.stream, &mut scratch.x,
|
||
).map_err(|e| anyhow!("w_in step fwd_into: {e}"))?;
|
||
|
||
// 2. a_proj = x @ W_a.T + b_a.
|
||
self.w_a.inner.forward_with_slices_into(
|
||
scratch.x.cuda_data(), n_batch,
|
||
&self.w_a.weight, &self.w_a.bias,
|
||
&self.cublas, &self.stream, &mut scratch.a_proj,
|
||
).map_err(|e| anyhow!("w_a step fwd_into: {e}"))?;
|
||
|
||
// 3. b_proj = x @ W_b.T + b_b.
|
||
self.w_b.inner.forward_with_slices_into(
|
||
scratch.x.cuda_data(), n_batch,
|
||
&self.w_b.weight, &self.w_b.bias,
|
||
&self.cublas, &self.stream, &mut scratch.b_proj,
|
||
).map_err(|e| anyhow!("w_b step fwd_into: {e}"))?;
|
||
|
||
// 4. scan_fwd_step — reads + writes scratch.x_state in-place,
|
||
// writes scratch.h_out. h_s2 stays zero (matches supervised
|
||
// forward_train_seq_into convention).
|
||
let block_threads: u32 = 32;
|
||
let grid_y: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.x_state.raw_ptr());
|
||
args.push_ptr(scratch.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(scratch.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(scratch.h_s2.cuda_data().raw_ptr());
|
||
args.push_ptr(scratch.h_out.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_fwd_step.cu_function(),
|
||
(n_batch as u32, grid_y, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_fwd_step launch: {e:?}"))?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// CRT Phase A0.5 (corrective) — bit-identical seeding helper.
|
||
/// Advances `scratch.x_state` by one timestep reading per-step
|
||
/// `a_proj` / `b_proj` row pointers from an EXTERNAL seq-path scratch
|
||
/// (typically the trainer's `mamba2_fwd_scratch.a_proj/b_proj` after
|
||
/// a `forward_only` call). Same scan kernel as `step_into`, but skips
|
||
/// the `w_in`/`w_a`/`w_b` GEMMs because the caller already has those
|
||
/// projections from a one-shot batched cuBLAS call.
|
||
///
|
||
/// This is the key to bit-identity vs forward_only: the seq path's
|
||
/// W_in→W_a/W_b GEMMs are computed once over all K rows in one cuBLAS
|
||
/// call. Re-running them per-row through `step_into` would not match
|
||
/// bit-for-bit (cuBLAS picks different algorithms for different
|
||
/// matrix shapes). By replaying scan_fwd_step over the seq path's
|
||
/// already-computed a_proj/b_proj, the x_state evolution is bit-
|
||
/// identical to scan_fwd_seq's internal register-state trajectory.
|
||
///
|
||
/// Args:
|
||
/// - `a_proj_row_ptr` / `b_proj_row_ptr`: device pointers into
|
||
/// `[N=n_batch, state_d]` slices of the seq scratch at the k-th row
|
||
/// (caller computes the byte offset).
|
||
/// - `scratch`: step scratch; `x_state` is read + written in place.
|
||
pub fn step_advance_from_seq_row(
|
||
&self,
|
||
a_proj_row_ptr: u64,
|
||
b_proj_row_ptr: u64,
|
||
scratch: &mut Mamba2BlockStepScratch,
|
||
) -> Result<()> {
|
||
let c = &self.config;
|
||
anyhow::ensure!(
|
||
scratch.hidden_dim == c.hidden_dim && scratch.state_dim == c.state_dim,
|
||
"step scratch shape mismatch: hidden_dim {} vs {}, state_dim {} vs {}",
|
||
scratch.hidden_dim, c.hidden_dim, scratch.state_dim, c.state_dim
|
||
);
|
||
let n_batch = scratch.n_batch;
|
||
let block_threads: u32 = 32;
|
||
let grid_y: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.x_state.raw_ptr());
|
||
args.push_ptr(a_proj_row_ptr);
|
||
args.push_ptr(b_proj_row_ptr);
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(scratch.h_s2.cuda_data().raw_ptr());
|
||
args.push_ptr(scratch.h_out.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_fwd_step.cu_function(),
|
||
(n_batch as u32, grid_y, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_fwd_step (seed) launch: {e:?}"))?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Backward chain paired with [`forward_train_seq`]. `d_h_enriched_seq`
|
||
/// has shape `[N, K, hidden_dim]` matching `cache.h_enriched_seq`.
|
||
/// Returns all nine parameter gradients (`dw_out` / `db_out` zeroed,
|
||
/// since this path skips W_out entirely).
|
||
pub fn backward_from_h_enriched_seq(
|
||
&self,
|
||
cache: &Mamba2ForwardCacheSeq,
|
||
d_h_enriched_seq: &GpuTensor,
|
||
) -> Result<Mamba2BackwardGrads> {
|
||
let c = &self.config;
|
||
let n_batch = cache.h_enriched_seq.shape()[0];
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
if d_h_enriched_seq.shape() != [n_batch, c.seq_len, c.hidden_dim] {
|
||
return Err(anyhow!(
|
||
"Mamba2Block::backward_from_h_enriched_seq: d_h_enriched_seq \
|
||
shape {:?} != [{}, {}, {}]",
|
||
d_h_enriched_seq.shape(), n_batch, c.seq_len, c.hidden_dim
|
||
));
|
||
}
|
||
|
||
let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim;
|
||
let per_sample_n = n_batch * c.hidden_dim * c.state_dim;
|
||
let d_a_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_a_per_channel: {e}"))?;
|
||
let d_b_per_channel = self.stream
|
||
.alloc_zeros::<f32>(per_chan_n)
|
||
.map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?;
|
||
let d_w_c_per_sample = self.stream
|
||
.alloc_zeros::<f32>(per_sample_n)
|
||
.map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?;
|
||
let d_h_s2 = self.stream
|
||
.alloc_zeros::<f32>(n_batch * c.hidden_dim)
|
||
.map_err(|e| anyhow!("alloc d_h_s2: {e}"))?;
|
||
|
||
let block_threads: u32 = 32;
|
||
let grid_y_h: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(cache.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(cache.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(d_h_enriched_seq.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(d_h_s2.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_bwd_seq.cu_function(),
|
||
(n_batch as u32, grid_y_h, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_bwd_seq launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// Reductions are identical to the final-step variant — d_a/d_b
|
||
// tensors have the same [N, sh2, K, state_d] / [N, sh2, state_d]
|
||
// shapes regardless of how many steps contributed to them.
|
||
let red_grid_z: u32 =
|
||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let red_cfg = (
|
||
(n_batch as u32, c.seq_len as u32, red_grid_z),
|
||
(block_threads, 1, 1),
|
||
);
|
||
let d_a_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?;
|
||
let d_b_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_a_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_a_proj seq: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_b_proj seq: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
let dw_c: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc dw_c seq: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(dw_c.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_w_c.cu_function(),
|
||
(c.hidden_dim as u32, red_grid_z, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce dw_c seq: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_b_proj seq: {e}"))?;
|
||
let x_act = LinearActivations { input: cache.x.clone() };
|
||
let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
|
||
.w_b.inner
|
||
.backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_b backward seq: {e}"))?;
|
||
|
||
let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_a_proj seq: {e}"))?;
|
||
let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
|
||
.w_a.inner
|
||
.backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_a backward seq: {e}"))?;
|
||
|
||
let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
|
||
.map_err(|e| anyhow!("sum d_x branches seq: {e}"))?;
|
||
|
||
let input_act = LinearActivations { input: cache.input_2d.clone() };
|
||
let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
|
||
.w_in.inner
|
||
.backward_with_slices(&d_x, &input_act, &self.w_in.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_in backward seq: {e}"))?;
|
||
|
||
// W_out is unused by callers of forward_train_seq — zero grads
|
||
// freeze those parameters under AdamW.
|
||
let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero dw_out seq: {e}"))?;
|
||
let db_out = GpuTensor::zeros(&[1], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero db_out seq: {e}"))?;
|
||
|
||
Ok(Mamba2BackwardGrads {
|
||
dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c, dw_out, db_out,
|
||
})
|
||
}
|
||
|
||
/// Perf-tuned backward variant — same math as
|
||
/// [`backward_from_h_enriched_seq`] but uses caller-provided
|
||
/// pre-allocated scratch buffers instead of allocating 7 fresh
|
||
/// CudaSlices per call. Cuts ~10-20ms/step on the training hot
|
||
/// path by eliminating `cudaMalloc` churn.
|
||
///
|
||
/// Callers MUST pre-allocate [`Mamba2BackwardScratch`] with shapes
|
||
/// matching the (n_batch, seq_len, hidden_dim, state_dim) used
|
||
/// here. The scratch buffer contents are overwritten by every call
|
||
/// (no carry between calls; each backward pass is self-contained).
|
||
pub fn backward_from_h_enriched_seq_into(
|
||
&self,
|
||
cache: &Mamba2ForwardCacheSeq,
|
||
d_h_enriched_seq: &GpuTensor,
|
||
scratch: &mut Mamba2BackwardScratch,
|
||
) -> Result<Mamba2BackwardGrads> {
|
||
let c = &self.config;
|
||
let n_batch = cache.h_enriched_seq.shape()[0];
|
||
let n_rows = n_batch * c.seq_len;
|
||
|
||
// Validate scratch shapes match — wrong scratch silently
|
||
// produces wrong gradients (size mismatch → out-of-bounds
|
||
// writes). Assert eagerly.
|
||
anyhow::ensure!(
|
||
scratch.n_batch == n_batch
|
||
&& scratch.seq_len == c.seq_len
|
||
&& scratch.hidden_dim == c.hidden_dim
|
||
&& scratch.state_dim == c.state_dim,
|
||
"Mamba2BackwardScratch shape mismatch: expected ({},{},{},{}) got ({},{},{},{})",
|
||
n_batch, c.seq_len, c.hidden_dim, c.state_dim,
|
||
scratch.n_batch, scratch.seq_len, scratch.hidden_dim, scratch.state_dim
|
||
);
|
||
anyhow::ensure!(
|
||
d_h_enriched_seq.shape() == [n_batch, c.seq_len, c.hidden_dim],
|
||
"d_h_enriched_seq shape {:?} != [{}, {}, {}]",
|
||
d_h_enriched_seq.shape(), n_batch, c.seq_len, c.hidden_dim
|
||
);
|
||
|
||
let block_threads: u32 = 32;
|
||
let grid_y_h: u32 = ((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(cache.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(cache.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(d_h_enriched_seq.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(scratch.d_a_per_channel.raw_ptr());
|
||
args.push_ptr(scratch.d_b_per_channel.raw_ptr());
|
||
args.push_ptr(scratch.d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(scratch.d_h_s2.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_bwd_seq.cu_function(),
|
||
(n_batch as u32, grid_y_h, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_bwd_seq launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// Small per-call buffers (~few KB each, no measurable cost).
|
||
let red_grid_z: u32 = ((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let red_cfg = (
|
||
(n_batch as u32, c.seq_len as u32, red_grid_z),
|
||
(block_threads, 1, 1),
|
||
);
|
||
let d_a_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_a_proj_flat into: {e}"))?;
|
||
let d_b_proj_flat: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc d_b_proj_flat into: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.d_a_per_channel.raw_ptr());
|
||
args.push_ptr(d_a_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_a_proj seq into: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.d_b_per_channel.raw_ptr());
|
||
args.push_ptr(d_b_proj_flat.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_b_proj seq into: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
let dw_c: CudaSlice<f32> = self.stream
|
||
.alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
|
||
.map_err(|e| anyhow!("alloc dw_c into: {e}"))?;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(scratch.d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(dw_c.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_w_c.cu_function(),
|
||
(c.hidden_dim as u32, red_grid_z, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce dw_c seq into: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_b_proj seq into: {e}"))?;
|
||
let x_act = LinearActivations { input: cache.x.clone() };
|
||
let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
|
||
.w_b.inner
|
||
.backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_b backward into: {e}"))?;
|
||
|
||
let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
|
||
.map_err(|e| anyhow!("reshape d_a_proj seq into: {e}"))?;
|
||
let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
|
||
.w_a.inner
|
||
.backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_a backward into: {e}"))?;
|
||
|
||
let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
|
||
.map_err(|e| anyhow!("sum d_x branches into: {e}"))?;
|
||
|
||
let input_act = LinearActivations { input: cache.input_2d.clone() };
|
||
let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
|
||
.w_in.inner
|
||
.backward_with_slices(&d_x, &input_act, &self.w_in.weight,
|
||
&self.cublas, &self.stream)
|
||
.map_err(|e| anyhow!("w_in backward into: {e}"))?;
|
||
|
||
let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero dw_out seq into: {e}"))?;
|
||
let db_out = GpuTensor::zeros(&[1], &self.stream)
|
||
.map_err(|e| anyhow!("alloc zero db_out seq into: {e}"))?;
|
||
|
||
Ok(Mamba2BackwardGrads {
|
||
dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c, dw_out, db_out,
|
||
})
|
||
}
|
||
|
||
/// Fully-pre-allocated backward — companion to [`forward_train_seq_into`].
|
||
/// Uses [`Mamba2BackwardScratch`] for the big per-channel scan
|
||
/// buffers, [`Mamba2BackwardGradsBuffers`] for the cuBLAS linear
|
||
/// backward outputs + reduction-result tensors. Zero allocation in
|
||
/// the call. Caller reads the final grads from `grads_buffers`.
|
||
///
|
||
/// `input` is the same `[B, K, in_dim]` tensor passed to the
|
||
/// matching `forward_train_seq_into` call. Needed for the W_in
|
||
/// backward (which dots dY^T against the original X).
|
||
pub fn backward_from_h_enriched_seq_full_into(
|
||
&self,
|
||
input: &GpuTensor,
|
||
fwd_scratch: &Mamba2BlockForwardScratch,
|
||
d_h_enriched_seq: &GpuTensor,
|
||
bwd_scratch: &mut Mamba2BackwardScratch,
|
||
grads_buffers: &mut Mamba2BackwardGradsBuffers,
|
||
) -> Result<()> {
|
||
let c = &self.config;
|
||
let n_batch = fwd_scratch.n_batch;
|
||
anyhow::ensure!(
|
||
n_batch == bwd_scratch.n_batch
|
||
&& n_batch == grads_buffers.n_batch
|
||
&& c.seq_len == bwd_scratch.seq_len
|
||
&& c.seq_len == grads_buffers.seq_len
|
||
&& c.hidden_dim == bwd_scratch.hidden_dim
|
||
&& c.hidden_dim == grads_buffers.hidden_dim
|
||
&& c.state_dim == bwd_scratch.state_dim
|
||
&& c.state_dim == grads_buffers.state_dim,
|
||
"backward_full_into: scratch shape mismatch"
|
||
);
|
||
anyhow::ensure!(
|
||
d_h_enriched_seq.shape() == [n_batch, c.seq_len, c.hidden_dim],
|
||
"d_h_enriched_seq shape {:?} != [{}, {}, {}]",
|
||
d_h_enriched_seq.shape(), n_batch, c.seq_len, c.hidden_dim
|
||
);
|
||
|
||
// ── 5′. Scan backward kernel → bwd_scratch slots ──────────────
|
||
let block_threads: u32 = 32;
|
||
let grid_y_h: u32 =
|
||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n_batch as i32;
|
||
let k_i32 = c.seq_len as i32;
|
||
let sh2_i32 = c.hidden_dim as i32;
|
||
let st_i32 = c.state_dim as i32;
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(fwd_scratch.a_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(fwd_scratch.b_proj.cuda_data().raw_ptr());
|
||
args.push_ptr(d_h_enriched_seq.cuda_data().raw_ptr());
|
||
args.push_ptr(self.w_c.raw_ptr());
|
||
args.push_ptr(bwd_scratch.d_a_per_channel.raw_ptr());
|
||
args.push_ptr(bwd_scratch.d_b_per_channel.raw_ptr());
|
||
args.push_ptr(bwd_scratch.d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(bwd_scratch.d_h_s2.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_bwd_seq.cu_function(),
|
||
(n_batch as u32, grid_y_h, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("scan_bwd_seq_full_into launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Reductions: per-channel → flat. Writes directly into
|
||
// grads_buffers.{d_a_proj_2d, d_b_proj_2d, dw_c} since
|
||
// their flat layouts match [N, K, state_d] / [sh2, state_d].
|
||
let red_grid_z: u32 =
|
||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let red_cfg = (
|
||
(n_batch as u32, c.seq_len as u32, red_grid_z),
|
||
(block_threads, 1, 1),
|
||
);
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(bwd_scratch.d_a_per_channel.raw_ptr());
|
||
args.push_ptr(grads_buffers.d_a_proj_2d.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_a_proj full_into: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(bwd_scratch.d_b_per_channel.raw_ptr());
|
||
args.push_ptr(grads_buffers.d_b_proj_2d.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(k_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_proj.cu_function(),
|
||
red_cfg.0, red_cfg.1, 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce d_b_proj full_into: {e:?}"))?;
|
||
}
|
||
}
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(bwd_scratch.d_w_c_per_sample.raw_ptr());
|
||
args.push_ptr(grads_buffers.dw_c.data_mut().raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_i32(sh2_i32);
|
||
args.push_i32(st_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_reduce_d_w_c.cu_function(),
|
||
(c.hidden_dim as u32, red_grid_z, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("reduce dw_c full_into: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── Linear backward: w_b, w_a, w_in (all _into variants).
|
||
// Each takes the saved forward input as a raw slice + batch
|
||
// dim (NO GpuTensor wrapping, NO CudaSlice::clone — cuMemAlloc
|
||
// inside capture is forbidden). w_b/w_a use fwd_scratch.x as
|
||
// their input; w_in uses the original input tensor.
|
||
let n_rows = n_batch * c.seq_len;
|
||
self.w_b.inner.backward_with_slices_into(
|
||
&grads_buffers.d_b_proj_2d,
|
||
fwd_scratch.x.cuda_data(), n_rows,
|
||
&self.w_b.weight,
|
||
&self.cublas, &self.stream,
|
||
&mut grads_buffers.dw_b, &mut grads_buffers.db_b, &mut grads_buffers.d_x_from_b,
|
||
).map_err(|e| anyhow!("w_b bwd_into: {e}"))?;
|
||
|
||
self.w_a.inner.backward_with_slices_into(
|
||
&grads_buffers.d_a_proj_2d,
|
||
fwd_scratch.x.cuda_data(), n_rows,
|
||
&self.w_a.weight,
|
||
&self.cublas, &self.stream,
|
||
&mut grads_buffers.dw_a, &mut grads_buffers.db_a, &mut grads_buffers.d_x_from_a,
|
||
).map_err(|e| anyhow!("w_a bwd_into: {e}"))?;
|
||
|
||
// d_x = d_x_from_a + d_x_from_b (in-place add — zero alloc).
|
||
grads_buffers.d_x_from_a
|
||
.add_into(&grads_buffers.d_x_from_b, &grads_buffers.d_x, &self.stream)
|
||
.map_err(|e| anyhow!("d_x add_into: {e}"))?;
|
||
|
||
self.w_in.inner.backward_with_slices_into(
|
||
&grads_buffers.d_x,
|
||
input.cuda_data(), n_rows,
|
||
&self.w_in.weight,
|
||
&self.cublas, &self.stream,
|
||
&mut grads_buffers.dw_in, &mut grads_buffers.db_in, &mut grads_buffers.d_x_from_in,
|
||
).map_err(|e| anyhow!("w_in bwd_into: {e}"))?;
|
||
|
||
// dw_out / db_out stay zero — W_out is unused in the seq path.
|
||
// (The grads_buffers init already created them zero; no action.)
|
||
Ok(())
|
||
}
|
||
|
||
/// Total trainable parameter count (sum of all projections + W_c).
|
||
pub fn param_count(&self) -> usize {
|
||
let c = &self.config;
|
||
// W_in: in*hidden + hidden_bias
|
||
let n_w_in = c.in_dim * c.hidden_dim + c.hidden_dim;
|
||
// W_a / W_b: hidden*state + state_bias
|
||
let n_w_ab = (c.hidden_dim * c.state_dim + c.state_dim) * 2;
|
||
// W_c: hidden*state (no bias)
|
||
let n_w_c = c.hidden_dim * c.state_dim;
|
||
// W_out: hidden*1 + 1
|
||
let n_w_out = c.hidden_dim + 1;
|
||
n_w_in + n_w_ab + n_w_c + n_w_out
|
||
}
|
||
}
|
||
|
||
// =====================================================================
|
||
// AdamW optimizer for Mamba2Block
|
||
// =====================================================================
|
||
|
||
/// AdamW hyperparameters.
|
||
#[derive(Debug, Clone)]
|
||
pub struct Mamba2AdamWConfig {
|
||
pub lr: f32,
|
||
pub beta1: f32,
|
||
pub beta2: f32,
|
||
pub epsilon: f32,
|
||
pub weight_decay: f32,
|
||
pub grad_clip_max_norm: Option<f32>,
|
||
}
|
||
|
||
impl Default for Mamba2AdamWConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
lr: 1e-3,
|
||
beta1: 0.9,
|
||
beta2: 0.999,
|
||
epsilon: 1e-8,
|
||
weight_decay: 1e-2,
|
||
grad_clip_max_norm: Some(1.0),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Per-parameter Adam moment buffers (m, v) on GPU. Allocated once at
|
||
/// optimizer construction; reused across all training steps.
|
||
struct AdamState {
|
||
m: CudaSlice<f32>,
|
||
v: CudaSlice<f32>,
|
||
}
|
||
|
||
/// GPU-native AdamW optimizer specialized for Mamba2Block's nine parameter
|
||
/// tensors. Holds the (m, v) moment state for each parameter; `step` applies
|
||
/// the AdamW update kernel to every weight + bias in place.
|
||
pub struct Mamba2AdamW {
|
||
pub config: Mamba2AdamWConfig,
|
||
pub step_count: i32,
|
||
stream: Arc<CudaStream>,
|
||
raw_stream: CUstream,
|
||
kernel: CudaFunction,
|
||
|
||
// ─── GPU-side grad-clip support ─────────────────────────────────
|
||
// Eliminates the 9× per-step memcpy_dtoh that the host grad-norm
|
||
// path required. norm_sq_phase1 → block_partials → norm_sq_phase2
|
||
// accumulates into grad_norm_sq_d; clip_scale converts → grad_scale_d
|
||
// which the devscale AdamW kernel reads as a device pointer.
|
||
kernel_norm_sq_phase1: CudaFunction,
|
||
kernel_norm_sq_phase2: CudaFunction,
|
||
kernel_clip_scale: CudaFunction,
|
||
kernel_adamw_devscale: CudaFunction,
|
||
/// 1-thread kernel that advances `step_count_d` by 1. Launched
|
||
/// inside the captured graph so each replay increments the counter
|
||
/// before the AdamW kernels read it.
|
||
kernel_increment_step: CudaFunction,
|
||
/// Per-block partial sums for grad-norm reduction phase 1.
|
||
/// Sized for the largest parameter tensor's grid: max(n)/256 blocks.
|
||
block_partials_d: CudaSlice<f32>,
|
||
/// Accumulator across all 9 tensors — single scalar.
|
||
grad_norm_sq_d: CudaSlice<f32>,
|
||
/// Final grad-clip scale (= min(1, max_norm/sqrt(norm_sq))). Read
|
||
/// by mamba2_alpha_adamw_step_devscale as a device pointer arg.
|
||
grad_scale_d: CudaSlice<f32>,
|
||
/// Device-resident 1-indexed step counter for the gpu_clip path.
|
||
/// Advanced via `kernel_increment_step` inside the captured graph,
|
||
/// so each replay sees the correct bias-correction step. Legacy
|
||
/// non-captured paths (`step`, `step_from_buffers`) keep using
|
||
/// the host-side `step_count` field.
|
||
step_count_d: CudaSlice<i32>,
|
||
/// max(grid_blocks) across all 9 tensors — capacity check for
|
||
/// block_partials_d.
|
||
max_grid_blocks: usize,
|
||
|
||
// (m, v) state for each of the 9 Mamba2 parameter tensors.
|
||
s_w_in: AdamState,
|
||
s_b_in: AdamState,
|
||
s_w_a: AdamState,
|
||
s_b_a: AdamState,
|
||
s_w_b: AdamState,
|
||
s_b_b: AdamState,
|
||
s_w_c: AdamState,
|
||
s_w_out: AdamState,
|
||
s_b_out: AdamState,
|
||
}
|
||
|
||
impl Mamba2AdamW {
|
||
/// Allocate optimizer state matching the shapes of `block`'s parameters.
|
||
/// State buffers are zero-initialised; first AdamW step is unbiased.
|
||
pub fn new(block: &Mamba2Block, config: Mamba2AdamWConfig) -> Result<Self> {
|
||
let stream = Arc::clone(&block.stream);
|
||
let kernel = block.kernel_adamw.clone();
|
||
|
||
// Build all AdamStates first (immutable borrow of `stream` during
|
||
// allocation), then move `stream` into Self.
|
||
let alloc = |n: usize, name: &str| -> Result<AdamState> {
|
||
let m = stream.alloc_zeros::<f32>(n)
|
||
.map_err(|e| anyhow!("alloc Adam m for {name}: {e}"))?;
|
||
let v = stream.alloc_zeros::<f32>(n)
|
||
.map_err(|e| anyhow!("alloc Adam v for {name}: {e}"))?;
|
||
Ok(AdamState { m, v })
|
||
};
|
||
let s_w_in = alloc(block.w_in.weight.len(), "w_in.weight")?;
|
||
let s_b_in = alloc(block.w_in.bias.len(), "w_in.bias")?;
|
||
let s_w_a = alloc(block.w_a.weight.len(), "w_a.weight")?;
|
||
let s_b_a = alloc(block.w_a.bias.len(), "w_a.bias")?;
|
||
let s_w_b = alloc(block.w_b.weight.len(), "w_b.weight")?;
|
||
let s_b_b = alloc(block.w_b.bias.len(), "w_b.bias")?;
|
||
let s_w_c = alloc(block.w_c.len(), "w_c")?;
|
||
let s_w_out = alloc(block.w_out.weight.len(), "w_out.weight")?;
|
||
let s_b_out = alloc(block.w_out.bias.len(), "w_out.bias")?;
|
||
|
||
// Load grad-norm kernel handles (compiled by build.rs into a
|
||
// dedicated grad_norm.cubin) + the devscale AdamW variant
|
||
// (compiled into mamba2_alpha_kernel.cubin alongside the
|
||
// host-scalar version).
|
||
static GRAD_NORM_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/grad_norm.cubin"));
|
||
let grad_norm_module = stream.context()
|
||
.load_cubin(GRAD_NORM_CUBIN.to_vec())
|
||
.map_err(|e| anyhow!("Mamba2AdamW: grad_norm cubin load: {e}"))?;
|
||
let kernel_norm_sq_phase1 = grad_norm_module
|
||
.load_function("grad_norm_sq_phase1")
|
||
.map_err(|e| anyhow!("grad_norm_sq_phase1 symbol: {e}"))?;
|
||
let kernel_norm_sq_phase2 = grad_norm_module
|
||
.load_function("grad_norm_sq_phase2")
|
||
.map_err(|e| anyhow!("grad_norm_sq_phase2 symbol: {e}"))?;
|
||
let kernel_clip_scale = grad_norm_module
|
||
.load_function("grad_clip_scale")
|
||
.map_err(|e| anyhow!("grad_clip_scale symbol: {e}"))?;
|
||
let kernel_adamw_devscale = block
|
||
._module
|
||
.load_function("mamba2_alpha_adamw_step_devscale")
|
||
.map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale symbol: {e}"))?;
|
||
let kernel_increment_step = block
|
||
._module
|
||
.load_function("mamba2_alpha_increment_step_counter")
|
||
.map_err(|e| anyhow!("mamba2_alpha_increment_step_counter symbol: {e}"))?;
|
||
|
||
// Sizing: largest parameter tensor's grid dim. With block=256,
|
||
// the largest tensor is w_in.weight = hidden_dim * in_dim.
|
||
let max_n = [
|
||
block.w_in.weight.len(), block.w_in.bias.len(),
|
||
block.w_a.weight.len(), block.w_a.bias.len(),
|
||
block.w_b.weight.len(), block.w_b.bias.len(),
|
||
block.w_c.len(),
|
||
block.w_out.weight.len(), block.w_out.bias.len(),
|
||
].into_iter().max().unwrap_or(0);
|
||
let max_grid_blocks = (max_n + 255) / 256;
|
||
let block_partials_d = stream.alloc_zeros::<f32>(max_grid_blocks.max(1))
|
||
.map_err(|e| anyhow!("block_partials alloc: {e}"))?;
|
||
let grad_norm_sq_d = stream.alloc_zeros::<f32>(1)
|
||
.map_err(|e| anyhow!("grad_norm_sq alloc: {e}"))?;
|
||
let grad_scale_d = stream.alloc_zeros::<f32>(1)
|
||
.map_err(|e| anyhow!("grad_scale alloc: {e}"))?;
|
||
let step_count_d = stream.alloc_zeros::<i32>(1)
|
||
.map_err(|e| anyhow!("step_count_d alloc: {e}"))?;
|
||
|
||
let raw_stream = stream.cu_stream();
|
||
Ok(Self {
|
||
stream,
|
||
raw_stream,
|
||
kernel,
|
||
kernel_norm_sq_phase1,
|
||
kernel_norm_sq_phase2,
|
||
kernel_clip_scale,
|
||
kernel_adamw_devscale,
|
||
kernel_increment_step,
|
||
block_partials_d,
|
||
grad_norm_sq_d,
|
||
grad_scale_d,
|
||
step_count_d,
|
||
max_grid_blocks,
|
||
step_count: 0,
|
||
s_w_in, s_b_in, s_w_a, s_b_a, s_w_b, s_b_b, s_w_c, s_w_out, s_b_out,
|
||
config,
|
||
})
|
||
}
|
||
|
||
/// One AdamW step. Applies grad_clip_max_norm if configured (computed
|
||
/// from the L2 norm of the concatenated parameter gradient vector).
|
||
pub fn step(&mut self, block: &mut Mamba2Block, grads: &Mamba2BackwardGrads) -> Result<()> {
|
||
self.step_count += 1;
|
||
let t = self.step_count;
|
||
|
||
// Compute grad-clip scale on host. The L2 norm requires a sum-of-
|
||
// squares reduction across all 9 tensors; cheapest is a host-side
|
||
// reduction over the 9 individual norms (since each is small).
|
||
let grad_scale = if let Some(max_norm) = self.config.grad_clip_max_norm {
|
||
// Each tensor's L2 squared norm = sum of (g_i * g_i). For Phase
|
||
// 1d.1 sizes the total parameter count is ~10k, so the host
|
||
// dtoh + sum is ~tens of microseconds. Re-evaluate at scale if
|
||
// it shows up in the training profile.
|
||
let mut total_sq = 0.0_f32;
|
||
for slice in [
|
||
grads.dw_in.cuda_data(), grads.db_in.cuda_data(),
|
||
grads.dw_a.cuda_data(), grads.db_a.cuda_data(),
|
||
grads.dw_b.cuda_data(), grads.db_b.cuda_data(),
|
||
&grads.dw_c,
|
||
grads.dw_out.cuda_data(),grads.db_out.cuda_data(),
|
||
] {
|
||
let mut host = vec![0.0_f32; slice.len()];
|
||
self.stream.memcpy_dtoh(slice, &mut host)
|
||
.map_err(|e| anyhow!("grad-norm dtoh: {e}"))?;
|
||
for g in &host {
|
||
total_sq += g * g;
|
||
}
|
||
}
|
||
let norm = total_sq.sqrt();
|
||
if norm > max_norm { max_norm / norm } else { 1.0 }
|
||
} else {
|
||
1.0
|
||
};
|
||
|
||
// Borrow the immutable bits as locals so we can mutably borrow
|
||
// self.s_* fields independently.
|
||
let kernel = &self.kernel;
|
||
let cfg = &self.config;
|
||
let rs = self.raw_stream;
|
||
|
||
// Launch the AdamW kernel for each parameter tensor.
|
||
adamw_apply(rs, kernel, cfg, block.w_in.weight.len(),
|
||
&mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_in.bias.len(),
|
||
&mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_a.weight.len(),
|
||
&mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_a.bias.len(),
|
||
&mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_b.weight.len(),
|
||
&mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_b.bias.len(),
|
||
&mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_c.len(),
|
||
&mut block.w_c, &grads.dw_c, &mut self.s_w_c, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_out.weight.len(),
|
||
&mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_out.bias.len(),
|
||
&mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, t, grad_scale)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
pub fn set_learning_rate(&mut self, lr: f32) {
|
||
self.config.lr = lr;
|
||
}
|
||
|
||
/// FULLY GPU-RESIDENT AdamW step — no host roundtrips, no per-step
|
||
/// `memcpy_dtoh`. Grad-norm reduction + clip-scale computed on
|
||
/// device; the devscale AdamW kernel reads `grad_scale_d` as a
|
||
/// device pointer.
|
||
///
|
||
/// Replaces [`step_from_buffers`] (which did 9× `memcpy_dtoh` per
|
||
/// step — each a stream sync barrier costing ~ms on the hot path)
|
||
/// for the production training path. Same math, zero pipeline stalls.
|
||
///
|
||
/// Requires `config.grad_clip_max_norm` to be Some at every call —
|
||
/// the no-clip path would need a kernel-write to init grad_scale_d
|
||
/// to 1.0 which the trainer doesn't currently set up. (Callers
|
||
/// who want no clipping should set max_norm to a very large value
|
||
/// like f32::INFINITY; the clip becomes a no-op since grad_scale=1.)
|
||
pub fn step_from_buffers_gpu_clip(
|
||
&mut self,
|
||
block: &mut Mamba2Block,
|
||
grads: &Mamba2BackwardGradsBuffers,
|
||
) -> Result<()> {
|
||
// Device step counter is the canonical step number for this
|
||
// path. Increment kernel advances it BEFORE the adamw kernels
|
||
// read it, so the first call observes step=1.
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(self.step_count_d.raw_ptr());
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_increment_step.cu_function(),
|
||
(1, 1, 1), (1, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("inc_step launch: {e:?}"))?;
|
||
}
|
||
}
|
||
|
||
// ── 1. Compute grad_scale_d on device.
|
||
if let Some(max_norm) = self.config.grad_clip_max_norm {
|
||
// Zero the cross-tensor accumulator.
|
||
self.stream
|
||
.memset_zeros(&mut self.grad_norm_sq_d)
|
||
.map_err(|e| anyhow!("zero grad_norm_sq: {e}"))?;
|
||
|
||
// Reduce each gradient tensor's sum-of-squares INTO the accumulator.
|
||
let grad_slices: [&CudaSlice<f32>; 9] = [
|
||
grads.dw_in.cuda_data(), grads.db_in.cuda_data(),
|
||
grads.dw_a.cuda_data(), grads.db_a.cuda_data(),
|
||
grads.dw_b.cuda_data(), grads.db_b.cuda_data(),
|
||
grads.dw_c.cuda_data(),
|
||
grads.dw_out.cuda_data(), grads.db_out.cuda_data(),
|
||
];
|
||
for grad in grad_slices.iter() {
|
||
let n = grad.len();
|
||
if n == 0 { continue; }
|
||
let block_threads: u32 = 256;
|
||
let n_blocks: u32 = (n as u32).div_ceil(block_threads);
|
||
debug_assert!((n_blocks as usize) <= self.max_grid_blocks,
|
||
"block_partials capacity {} < n_blocks {}", self.max_grid_blocks, n_blocks);
|
||
let n_i32 = n as i32;
|
||
// Phase 1: per-block tree-reduce → block_partials_d.
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(grad.raw_ptr());
|
||
args.push_i32(n_i32);
|
||
args.push_ptr(self.block_partials_d.raw_ptr());
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_norm_sq_phase1.cu_function(),
|
||
(n_blocks, 1, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("grad_norm phase1: {e:?}"))?;
|
||
}
|
||
}
|
||
// Phase 2: single-block reduce of partials → ACCUMULATE into norm_sq_d.
|
||
{
|
||
let nb_i32 = n_blocks as i32;
|
||
let accumulate_i32: i32 = 1; // ALWAYS accumulate (zero was done before loop)
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(self.block_partials_d.raw_ptr());
|
||
args.push_i32(nb_i32);
|
||
args.push_ptr(self.grad_norm_sq_d.raw_ptr());
|
||
args.push_i32(accumulate_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_norm_sq_phase2.cu_function(),
|
||
(1, 1, 1), (block_threads, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("grad_norm phase2: {e:?}"))?;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Convert norm_sq → grad_scale_d = min(1, max_norm / sqrt(norm_sq + eps)).
|
||
{
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(self.grad_norm_sq_d.raw_ptr());
|
||
args.push_f32(max_norm);
|
||
args.push_ptr(self.grad_scale_d.raw_ptr());
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
self.kernel_clip_scale.cu_function(),
|
||
(1, 1, 1), (1, 1, 1), 0,
|
||
self.raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("grad_clip_scale: {e:?}"))?;
|
||
}
|
||
}
|
||
} else {
|
||
return Err(anyhow!(
|
||
"step_from_buffers_gpu_clip requires grad_clip_max_norm to be Some \
|
||
(set to f32::INFINITY to disable clipping in practice)"
|
||
));
|
||
}
|
||
|
||
// ── 2. AdamW per param via the devscale kernel (reads
|
||
// grad_scale_d AND step_count_d as device pointers).
|
||
let kernel = &self.kernel_adamw_devscale;
|
||
let cfg = &self.config;
|
||
let grad_scale = &self.grad_scale_d;
|
||
let step_d = &self.step_count_d;
|
||
let rs = self.raw_stream;
|
||
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_in.weight.len(),
|
||
&mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_in.bias.len(),
|
||
&mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_a.weight.len(),
|
||
&mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_a.bias.len(),
|
||
&mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_b.weight.len(),
|
||
&mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_b.bias.len(),
|
||
&mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_c.len(),
|
||
&mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_out.weight.len(),
|
||
&mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, step_d, grad_scale)?;
|
||
adamw_apply_devscale(rs, kernel, cfg, block.w_out.bias.len(),
|
||
&mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, step_d, grad_scale)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// AdamW step that reads gradients directly from a
|
||
/// [`Mamba2BackwardGradsBuffers`] (the pre-allocated buffer set
|
||
/// produced by [`Mamba2Block::backward_from_h_enriched_seq_full_into`]).
|
||
/// Avoids constructing a temporary [`Mamba2BackwardGrads`] wrapper
|
||
/// every step.
|
||
pub fn step_from_buffers(
|
||
&mut self,
|
||
block: &mut Mamba2Block,
|
||
grads: &Mamba2BackwardGradsBuffers,
|
||
) -> Result<()> {
|
||
self.step_count += 1;
|
||
let t = self.step_count;
|
||
|
||
let grad_scale = if let Some(max_norm) = self.config.grad_clip_max_norm {
|
||
let mut total_sq = 0.0_f32;
|
||
for slice in [
|
||
grads.dw_in.cuda_data(), grads.db_in.cuda_data(),
|
||
grads.dw_a.cuda_data(), grads.db_a.cuda_data(),
|
||
grads.dw_b.cuda_data(), grads.db_b.cuda_data(),
|
||
grads.dw_c.cuda_data(),
|
||
grads.dw_out.cuda_data(),grads.db_out.cuda_data(),
|
||
] {
|
||
let mut host = vec![0.0_f32; slice.len()];
|
||
self.stream.memcpy_dtoh(slice, &mut host)
|
||
.map_err(|e| anyhow!("grad-norm dtoh: {e}"))?;
|
||
for g in &host {
|
||
total_sq += g * g;
|
||
}
|
||
}
|
||
let norm = total_sq.sqrt();
|
||
if norm > max_norm { max_norm / norm } else { 1.0 }
|
||
} else {
|
||
1.0
|
||
};
|
||
|
||
let kernel = &self.kernel;
|
||
let cfg = &self.config;
|
||
let rs = self.raw_stream;
|
||
|
||
adamw_apply(rs, kernel, cfg, block.w_in.weight.len(),
|
||
&mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_in.bias.len(),
|
||
&mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_a.weight.len(),
|
||
&mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_a.bias.len(),
|
||
&mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_b.weight.len(),
|
||
&mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_b.bias.len(),
|
||
&mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_c.len(),
|
||
&mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_out.weight.len(),
|
||
&mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, t, grad_scale)?;
|
||
adamw_apply(rs, kernel, cfg, block.w_out.bias.len(),
|
||
&mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, t, grad_scale)?;
|
||
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// One AdamW kernel launch for a single parameter tensor. Free function
|
||
/// (rather than a method) so the caller can mutably borrow distinct fields
|
||
/// of `Mamba2AdamW` (the per-param Adam state) while sharing immutable
|
||
/// references to the stream + kernel + hyperparameters.
|
||
/// Companion to [`adamw_apply`] — launches `mamba2_alpha_adamw_step_devscale`
|
||
/// which reads `grad_scale` from a device pointer instead of a host
|
||
/// scalar. Used by [`Mamba2AdamW::step_from_buffers_gpu_clip`] to keep
|
||
/// the entire AdamW step host-roundtrip-free.
|
||
fn adamw_apply_devscale(
|
||
raw_stream: CUstream,
|
||
kernel: &CudaFunction,
|
||
cfg: &Mamba2AdamWConfig,
|
||
n: usize,
|
||
param: &mut CudaSlice<f32>,
|
||
grad: &CudaSlice<f32>,
|
||
state: &mut AdamState,
|
||
step_d: &CudaSlice<i32>,
|
||
grad_scale_d: &CudaSlice<f32>,
|
||
) -> Result<()> {
|
||
let block_threads: u32 = 256;
|
||
let blocks: u32 = (n as u32).div_ceil(block_threads);
|
||
let n_i32 = n as i32;
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(param.raw_ptr());
|
||
args.push_ptr(grad.raw_ptr());
|
||
args.push_ptr(state.m.raw_ptr());
|
||
args.push_ptr(state.v.raw_ptr());
|
||
args.push_f32(cfg.lr);
|
||
args.push_f32(cfg.beta1);
|
||
args.push_f32(cfg.beta2);
|
||
args.push_f32(cfg.epsilon);
|
||
args.push_f32(cfg.weight_decay);
|
||
args.push_ptr(grad_scale_d.raw_ptr()); // DEVICE pointer instead of host scalar
|
||
args.push_ptr(step_d.raw_ptr()); // DEVICE pointer instead of host scalar
|
||
args.push_i32(n_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
kernel.cu_function(),
|
||
(blocks, 1, 1), (block_threads, 1, 1), 0,
|
||
raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("mamba2_alpha_adamw_step_devscale launch (n={n}): {e:?}"))?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn adamw_apply(
|
||
raw_stream: CUstream,
|
||
kernel: &CudaFunction,
|
||
cfg: &Mamba2AdamWConfig,
|
||
n: usize,
|
||
param: &mut CudaSlice<f32>,
|
||
grad: &CudaSlice<f32>,
|
||
state: &mut AdamState,
|
||
t: i32,
|
||
grad_scale: f32,
|
||
) -> Result<()> {
|
||
let block_threads: u32 = 256;
|
||
let blocks: u32 = ((n + block_threads as usize - 1) / block_threads as usize) as u32;
|
||
let n_i32 = n as i32;
|
||
let mut args = RawArgs::new();
|
||
args.push_ptr(param.raw_ptr());
|
||
args.push_ptr(grad.raw_ptr());
|
||
args.push_ptr(state.m.raw_ptr());
|
||
args.push_ptr(state.v.raw_ptr());
|
||
args.push_f32(cfg.lr);
|
||
args.push_f32(cfg.beta1);
|
||
args.push_f32(cfg.beta2);
|
||
args.push_f32(cfg.epsilon);
|
||
args.push_f32(cfg.weight_decay);
|
||
args.push_f32(grad_scale);
|
||
args.push_i32(t);
|
||
args.push_i32(n_i32);
|
||
let mut ptrs = args.build_arg_ptrs();
|
||
unsafe {
|
||
raw_launch(
|
||
kernel.cu_function(),
|
||
(blocks, 1, 1), (block_threads, 1, 1), 0,
|
||
raw_stream, &mut ptrs[..args.len()],
|
||
).map_err(|e| anyhow!("mamba2_alpha_adamw_step launch (n={n}): {e:?}"))?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use cudarc::driver::CudaContext;
|
||
|
||
fn cuda_stream_or_skip() -> Option<Arc<CudaStream>> {
|
||
match CudaContext::new(0) {
|
||
Ok(ctx) => Some(ctx.default_stream()),
|
||
Err(_) => None,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_config_rejects_state_over_32() {
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 64, state_dim: 33, seq_len: 32,
|
||
};
|
||
assert!(cfg.validate().is_err(), "state_dim > 32 must be rejected");
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_config_rejects_seq_len_over_kernel_max() {
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 64, state_dim: 16, seq_len: MAMBA2_KERNEL_SEQ_MAX + 1,
|
||
};
|
||
assert!(
|
||
cfg.validate().is_err(),
|
||
"seq_len > {} must be rejected (backward x_hist limit)",
|
||
MAMBA2_KERNEL_SEQ_MAX
|
||
);
|
||
}
|
||
|
||
/// THE end-to-end correctness check: build a tiny Mamba2Block, run
|
||
/// 20 AdamW steps on a fixed input with a fixed binary label, and
|
||
/// assert mean BCE loss falls monotonically. If analytical backward
|
||
/// has a sign flip or wrong scale somewhere, this test fails by
|
||
/// observing the loss either flatten or rise.
|
||
#[test]
|
||
fn test_mamba2_block_training_loop_decreases_loss() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => return,
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 4, hidden_dim: 8, state_dim: 4, seq_len: 8,
|
||
};
|
||
let n_batch = 4;
|
||
let mut block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||
let mut opt = Mamba2AdamW::new(&block, Mamba2AdamWConfig {
|
||
lr: 1e-2, // bigger lr — small model, 20-step convergence
|
||
..Default::default()
|
||
}).expect("opt init");
|
||
|
||
// Fixed input + fixed binary labels (half +1, half 0).
|
||
let n = n_batch * cfg.seq_len * cfg.in_dim;
|
||
let h: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.137).sin()).collect();
|
||
let dev_in = stream.clone_htod(&h).expect("htod");
|
||
let input = GpuTensor::new(dev_in, vec![n_batch, cfg.seq_len, cfg.in_dim]).expect("input");
|
||
let labels: Vec<f32> = (0..n_batch).map(|i| if i < n_batch / 2 { 1.0 } else { 0.0 }).collect();
|
||
|
||
let bce_loss = |logits_host: &[f32], labels_host: &[f32]| -> f32 {
|
||
// Numerically-stable BCE-with-logits, mean over batch.
|
||
let mut s = 0.0_f32;
|
||
let eps = 1e-7_f32;
|
||
for (&z, &y) in logits_host.iter().zip(labels_host.iter()) {
|
||
let z = z.clamp(-50.0, 50.0);
|
||
let p = (1.0 / (1.0 + (-z).exp())).clamp(eps, 1.0 - eps);
|
||
s += -(y * p.ln() + (1.0 - y) * (1.0 - p).ln());
|
||
}
|
||
s / logits_host.len() as f32
|
||
};
|
||
|
||
let mut last_loss = f32::INFINITY;
|
||
let mut decrease_streak = 0;
|
||
let mut total_decreases = 0;
|
||
for step in 0..20 {
|
||
let (logit, cache) = block.forward_train(&input).expect("forward_train");
|
||
let logit_host = logit.to_host(&stream).expect("logit dtoh");
|
||
let loss = bce_loss(&logit_host, &labels);
|
||
|
||
// d_logit = (sigmoid(z) - y) / N for mean-BCE-with-logits.
|
||
let d_logit_host: Vec<f32> = logit_host.iter().zip(labels.iter())
|
||
.map(|(&z, &y)| {
|
||
let p = 1.0 / (1.0 + (-z.clamp(-50.0, 50.0)).exp());
|
||
(p - y) / (n_batch as f32)
|
||
})
|
||
.collect();
|
||
let d_logit_dev = stream.clone_htod(&d_logit_host).expect("htod d_logit");
|
||
let d_logit = GpuTensor::new(d_logit_dev, vec![n_batch, 1]).expect("d_logit");
|
||
|
||
let grads = block.backward(&cache, &d_logit).expect("backward");
|
||
opt.step(&mut block, &grads).expect("opt step");
|
||
|
||
// Track convergence behaviour.
|
||
if loss < last_loss {
|
||
decrease_streak += 1;
|
||
total_decreases += 1;
|
||
} else {
|
||
decrease_streak = 0;
|
||
}
|
||
eprintln!("step={step:2} loss={loss:.5} streak={decrease_streak}");
|
||
last_loss = loss;
|
||
}
|
||
|
||
// Strict invariants:
|
||
// - Loss must have decreased a majority of the steps (allow some
|
||
// non-monotonicity from the AdamW momentum + small batch noise).
|
||
assert!(
|
||
total_decreases >= 15,
|
||
"expected ≥15 decreasing steps out of 20, got {total_decreases} — backward likely broken"
|
||
);
|
||
// - Final loss must be meaningfully below the initial random-init loss
|
||
// (which starts near ln(2) ≈ 0.693 for balanced labels). Initial
|
||
// loss after random Xavier init varies; we just require we got
|
||
// below the chance baseline.
|
||
let (final_logit, _) = block.forward_train(&input).expect("final forward");
|
||
let final_logit_host = final_logit.to_host(&stream).expect("final dtoh");
|
||
let final_loss = bce_loss(&final_logit_host, &labels);
|
||
assert!(
|
||
final_loss < 0.65,
|
||
"final loss {final_loss} did not drop below 0.65 — training did not converge"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_backward_returns_finite_grads_with_correct_shapes() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => return,
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8,
|
||
};
|
||
let n_batch = 3;
|
||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||
|
||
// Forward with cache.
|
||
let n_in = n_batch * cfg.seq_len * cfg.in_dim;
|
||
let h: Vec<f32> = (0..n_in).map(|i| ((i as f32) * 0.137).sin()).collect();
|
||
let d = stream.clone_htod(&h).expect("htod");
|
||
let input = GpuTensor::new(d, vec![n_batch, cfg.seq_len, cfg.in_dim]).expect("input");
|
||
let (logit, cache) = block.forward_train(&input).expect("forward_train");
|
||
|
||
// Pretend d_logit is just sigmoid(logit) - 0.5 (a BCE-shaped scalar).
|
||
let logit_h = logit.to_host(&stream).expect("logit dtoh");
|
||
let d_logit_h: Vec<f32> = logit_h.iter()
|
||
.map(|z| 1.0 / (1.0 + (-z).exp()) - 0.5)
|
||
.collect();
|
||
let d_logit_dev = stream.clone_htod(&d_logit_h).expect("htod");
|
||
let d_logit = GpuTensor::new(d_logit_dev, vec![n_batch, 1]).expect("d_logit tensor");
|
||
|
||
let grads = block.backward(&cache, &d_logit).expect("backward");
|
||
|
||
// Shape assertions match each weight's forward shape.
|
||
assert_eq!(grads.dw_in.shape(), &[cfg.hidden_dim, cfg.in_dim], "dw_in shape");
|
||
assert_eq!(grads.db_in.shape(), &[cfg.hidden_dim], "db_in shape");
|
||
assert_eq!(grads.dw_a.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_a shape");
|
||
assert_eq!(grads.db_a.shape(), &[cfg.state_dim], "db_a shape");
|
||
assert_eq!(grads.dw_b.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_b shape");
|
||
assert_eq!(grads.db_b.shape(), &[cfg.state_dim], "db_b shape");
|
||
assert_eq!(grads.dw_c.len(), cfg.hidden_dim * cfg.state_dim, "dw_c numel");
|
||
assert_eq!(grads.dw_out.shape(), &[1, cfg.hidden_dim], "dw_out shape");
|
||
assert_eq!(grads.db_out.shape(), &[1], "db_out shape");
|
||
|
||
// Finiteness for every gradient tensor.
|
||
for (name, g) in [
|
||
("dw_in", &grads.dw_in), ("db_in", &grads.db_in),
|
||
("dw_a", &grads.dw_a), ("db_a", &grads.db_a),
|
||
("dw_b", &grads.dw_b), ("db_b", &grads.db_b),
|
||
("dw_out",&grads.dw_out),("db_out",&grads.db_out),
|
||
] {
|
||
let h = g.to_host(&stream).expect("dtoh");
|
||
for (i, &v) in h.iter().enumerate() {
|
||
assert!(v.is_finite(), "{} grad index {} non-finite: {}", name, i, v);
|
||
}
|
||
}
|
||
// dw_c is a raw CudaSlice<f32> (no GpuTensor wrapper); read directly.
|
||
let mut dw_c_host = vec![0.0_f32; grads.dw_c.len()];
|
||
stream.memcpy_dtoh(&grads.dw_c, &mut dw_c_host).expect("dw_c dtoh");
|
||
for (i, &v) in dw_c_host.iter().enumerate() {
|
||
assert!(v.is_finite(), "dw_c grad index {} non-finite: {}", i, v);
|
||
}
|
||
|
||
// At least one element of dw_in should be non-zero (gradient
|
||
// actually flowed through the chain — not silently zero).
|
||
let dw_in_h = grads.dw_in.to_host(&stream).expect("dw_in dtoh");
|
||
let any_nonzero = dw_in_h.iter().any(|&v| v.abs() > 1e-12);
|
||
assert!(any_nonzero, "dw_in entirely zero — gradient didn't propagate to W_in");
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_backward_rejects_wrong_d_logit_shape() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => return,
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8,
|
||
};
|
||
let n_batch = 3;
|
||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||
let input = GpuTensor::zeros(&[n_batch, cfg.seq_len, cfg.in_dim], &stream).expect("input");
|
||
let (_logit, cache) = block.forward_train(&input).expect("forward_train");
|
||
// Wrong d_logit shape: [n_batch, 2] instead of [n_batch, 1].
|
||
let bad = GpuTensor::zeros(&[n_batch, 2], &stream).expect("bad d_logit");
|
||
assert!(block.backward(&cache, &bad).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_forward_train_returns_cache() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => return,
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 32, state_dim: 8, seq_len: 16,
|
||
};
|
||
let n_batch = 4;
|
||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||
let n = n_batch * cfg.seq_len * cfg.in_dim;
|
||
let h: Vec<f32> = (0..n).map(|i| (i as f32 * 1e-3).cos()).collect();
|
||
let d = stream.clone_htod(&h).expect("htod");
|
||
let input = GpuTensor::new(d, vec![n_batch, cfg.seq_len, cfg.in_dim]).expect("input");
|
||
let (logit, cache) = block.forward_train(&input).expect("forward_train");
|
||
assert_eq!(logit.shape(), &[n_batch, 1]);
|
||
assert_eq!(cache.input_2d.shape(), &[n_batch * cfg.seq_len, cfg.in_dim]);
|
||
assert_eq!(cache.x.shape(), &[n_batch * cfg.seq_len, cfg.hidden_dim]);
|
||
assert_eq!(cache.a_proj.shape(), &[n_batch * cfg.seq_len, cfg.state_dim]);
|
||
assert_eq!(cache.b_proj.shape(), &[n_batch * cfg.seq_len, cfg.state_dim]);
|
||
assert_eq!(cache.h_enriched.shape(), &[n_batch, cfg.hidden_dim]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_config_rejects_zero_dims() {
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 0, hidden_dim: 64, state_dim: 16, seq_len: 32,
|
||
};
|
||
assert!(cfg.validate().is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_forward_shape_and_finite() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => {
|
||
eprintln!("CUDA unavailable; skipping forward smoke");
|
||
return;
|
||
}
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 32, state_dim: 8, seq_len: 16,
|
||
};
|
||
let n_batch = 4;
|
||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("Mamba2Block::new");
|
||
|
||
// Build deterministic synthetic input [N, K, in_dim] = [4, 16, 81] = 5184 f32.
|
||
let n_input = n_batch * cfg.seq_len * cfg.in_dim;
|
||
let host_input: Vec<f32> = (0..n_input).map(|i| ((i as f32) * 1e-3).sin()).collect();
|
||
let dev_input = stream.clone_htod(&host_input).expect("htod");
|
||
let input_tensor = GpuTensor::new(dev_input, vec![n_batch, cfg.seq_len, cfg.in_dim])
|
||
.expect("input tensor");
|
||
|
||
let logit = block.forward(&input_tensor).expect("forward");
|
||
// Output shape must be [N, 1].
|
||
assert_eq!(logit.shape(), &[n_batch, 1]);
|
||
// All logits must be finite (no NaN / Inf — sanity check on the
|
||
// selective scan + projections).
|
||
let host_logit = logit.to_host(&stream).expect("dtoh");
|
||
assert_eq!(host_logit.len(), n_batch);
|
||
for (i, &z) in host_logit.iter().enumerate() {
|
||
assert!(z.is_finite(), "logit[{i}] non-finite: {z}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_forward_rejects_wrong_shape() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => return,
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 32, state_dim: 8, seq_len: 16,
|
||
};
|
||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||
// Wrong in_dim.
|
||
let bad = GpuTensor::zeros(&[2, 16, 99], &stream).expect("bad alloc");
|
||
assert!(block.forward(&bad).is_err());
|
||
// Wrong seq_len.
|
||
let bad2 = GpuTensor::zeros(&[2, 8, 81], &stream).expect("bad alloc 2");
|
||
assert!(block.forward(&bad2).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_mamba2_block_constructs_and_loads_kernels() {
|
||
let stream = match cuda_stream_or_skip() {
|
||
Some(s) => s,
|
||
None => {
|
||
eprintln!("CUDA unavailable; skipping GPU construction test");
|
||
return;
|
||
}
|
||
};
|
||
let cfg = Mamba2BlockConfig {
|
||
in_dim: 81, hidden_dim: 64, state_dim: 16, seq_len: 32,
|
||
};
|
||
let block = Mamba2Block::new(cfg, stream).expect("Mamba2Block::new");
|
||
// Both kernel handles must be resolved.
|
||
let _ = &block.kernel_fwd;
|
||
let _ = &block.kernel_bwd;
|
||
// Parameter count sanity: 81*64 + 64 + (64*16 + 16)*2 + 64*16 + 64 + 1
|
||
// = 5184 + 64 + 2080 + 1024 + 65 = 8417
|
||
let expected = 81 * 64 + 64 + (64 * 16 + 16) * 2 + 64 * 16 + 64 + 1;
|
||
assert_eq!(block.param_count(), expected);
|
||
}
|
||
}
|