//! 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);
    }
}
