mamba-rs 0.6.8

//! AdamW optimizer in f32 master precision.
//!
//! Mirrors `torch.optim.AdamW` (Loshchilov & Hutter, "Decoupled Weight Decay
//! Regularization", ICLR 2019). Designed for AMP/mixed-precision training:
//!
//! - **Master weights** stay in f32 (per-tensor `GpuBuffer`s in
//!   `GpuMambaTrainWeights` / `GpuMamba3Weights`).
//! - **Gradients** are accumulated in f32 (flat `GpuBuffer` in
//!   `GpuMambaGrads.flat` / `GpuMamba3Grads.flat`), which the loss-scaler
//!   has already unscaled before this step runs.
//! - **Optimizer state** (`m`, `v`) lives in f32 alongside grads. Storing
//!   Adam moments in bf16 empirically diverges within ~1k SSM steps;
//!   PyTorch / DeepSpeed / Apex / CleanRL all keep them in f32.
//!
//! ## Usage
//! ```ignore
//! let mut adam = GpuAdamW::new(&ctx.stream, grads.flat.len())?
//!     .with_lr(3e-4)
//!     .with_weight_decay(1e-2);
//!
//! for _ in 0..n_steps {
//!     grads.zero(&ctx.stream)?;
//!     forward_backward(...)?;
//!     // (optional) loss-scaler unscale + grad clip here
//!     adam.step_m1(&ctx, &kernels, &mut weights.master, &grads)?;
//!     weights.sync_master_to_compute(&ctx)?;
//! }
//! ```
//!
//! ## Precision rationale
//! Default `f32` for everything inside the optimizer matches:
//! - PyTorch `torch.optim.AdamW` (always f32 even under AMP)
//! - NVIDIA Apex `FusedAdam`
//! - DeepSpeed ZeRO stage 0/1/2
//! - Hugging Face `accelerate.optimizer.AcceleratedOptimizer`
//!
//! Bias-correction factors `1/(1-β1ᵗ)` and `1/(1-β2ᵗ)` are computed CPU-side
//! (one `powf` per step) so the kernel sees plain scalar multiplies — same
//! as PyTorch's `_single_tensor_adamw` non-capturable path.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, PushKernelArg};

use crate::mamba_ssm::gpu::buffers::{GpuBuffer, GradSlice};
use crate::mamba_ssm::gpu::context::GpuCtx;
use crate::mamba_ssm::gpu::launch::grid_1d;

/// 2-element device buffer holding `[bias_c1, bias_c2]` for the
/// CUDA-Graph-capturable AdamW kernel. CPU writes the next-step values
/// here BEFORE each graph replay so the captured kernel reads fresh
/// bias-correction factors via a stable device pointer.
///
/// Mirrors the device-side state PyTorch's `AdamW(capturable=True)` keeps
/// per param group.
pub struct AdamWBiasFactors {
    pub buf: GpuBuffer,
}

impl AdamWBiasFactors {
    /// Allocate the 2-element device buffer. Initialises to `[1.0, 1.0]`
    /// — the neutral bias-correction value (`1 / (1 - β^1)` for step 1 is
    /// `~10` for β1=0.9 but the kernel tolerates any finite factor; 1.0
    /// produces an "Adam without bias correction" update if the buffer is
    /// ever read before `write()`). This guards against a silent-wrong
    /// update path where a zero-init buffer would make `m_hat = v_hat = 0`
    /// and the captured kernel would apply ONLY weight decay, no Adam
    /// step, on the first replay if `write()` was forgotten.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, String> {
        let buf = GpuBuffer::zeros(stream, 3)?;
        let mut this = Self { buf };
        this.write(stream, 1.0, 1.0, 1e-3)?;
        Ok(this)
    }

    /// Write `(bc1, bc2, lr)` for the upcoming step. Async H2D — the next
    /// graph replay will see these values via the device pointer. The lr
    /// rides the buffer so a warmup/cosine schedule works under a
    /// captured graph (the legacy per-tensor kernels baked lr by value
    /// at capture time).
    pub fn write(
        &mut self,
        stream: &Arc<CudaStream>,
        bc1: f32,
        bc2: f32,
        lr: f32,
    ) -> Result<(), String> {
        debug_assert!(
            bc1.is_finite() && bc2.is_finite() && bc1 > 0.0 && bc2 > 0.0,
            "AdamWBiasFactors::write got non-finite or non-positive values: bc1={bc1} bc2={bc2}"
        );
        debug_assert!(
            lr.is_finite() && lr > 0.0,
            "AdamWBiasFactors::write got invalid lr {lr}"
        );
        self.buf.upload(stream, &[bc1, bc2, lr])
    }

    pub fn ptr(&self) -> cudarc::driver::sys::CUdeviceptr {
        self.buf.cached_ptr()
    }
}

/// Device pointers for one AdamW tensor update: weight + grad + the
/// matching `m`/`v` slices (offset into the flat arenas by the same
/// amount as `grad` is into `grads.flat`).
#[derive(Clone, Copy, Debug)]
pub struct AdamWParamPtrs {
    pub weight: cudarc::driver::sys::CUdeviceptr,
    pub grad: cudarc::driver::sys::CUdeviceptr,
    pub m: cudarc::driver::sys::CUdeviceptr,
    pub v: cudarc::driver::sys::CUdeviceptr,
}

/// CPU-side snapshot of the full AdamW state for a bit-continuous
/// resume: without the moment buffers and step counter, a "resumed" run
/// re-warms Adam from zero and provably diverges from the unbroken run.
/// The learning rate is NOT part of the blob — it belongs to the
/// caller's schedule (see [`GpuAdamW::export_state`]).
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWStateBlob {
    pub m: Vec<f32>,
    pub v: Vec<f32>,
    pub step: u64,
    pub beta1: f32,
    pub beta2: f32,
    pub eps: f32,
    pub weight_decay: f32,
    pub reference_no_decay: bool,
}

/// f32 fused AdamW optimizer (matches `torch.optim.AdamW`).
pub struct GpuAdamW {
    /// First moment (m) in f32, layout matches the flat grad arena.
    pub m: GpuBuffer,
    /// Second moment (v) in f32, layout matches the flat grad arena.
    pub v: GpuBuffer,
    /// Step counter (1-indexed at first call to [`Self::step_one`]).
    pub step: u64,
    pub lr: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub eps: f32,
    pub weight_decay: f32,
    /// Reference-faithful no-decay parameter groups. When true, the
    /// per-tensor step functions ([`step_m1_capturable`] /
    /// [`step_m3_capturable`]) pass `weight_decay = 0` for the parameters
    /// the reference implementation marks `_no_weight_decay` — M1: `a_log`,
    /// `d_param`, `dt_proj_b` and the RMSNorm scales; M3: `dt_bias`,
    /// `d_param` and every norm scale. Decaying `a_log` pulls all decay
    /// rates toward A = -1 over long runs. Default `false` preserves the
    /// historical decay-everything behavior bit-for-bit; costs nothing
    /// either way (the decay coefficient is a per-launch scalar).
    pub reference_no_decay: bool,
}

impl GpuAdamW {
    /// Allocate zero-initialized `m`, `v` of length `n_params` (= total
    /// number of f32 master weights = `grads.flat.len()`).
    ///
    /// Defaults match `torch.optim.AdamW(params)`: lr=1e-3, β1=0.9, β2=0.999,
    /// eps=1e-8, weight_decay=1e-2.
    pub fn new(stream: &Arc<CudaStream>, n_params: usize) -> Result<Self, String> {
        Ok(Self {
            m: GpuBuffer::zeros(stream, n_params)?,
            v: GpuBuffer::zeros(stream, n_params)?,
            step: 0,
            lr: 1e-3,
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            weight_decay: 1e-2,
            reference_no_decay: false,
        })
    }

    /// Toggle the reference-faithful no-decay parameter groups (see the
    /// field docs on [`Self::reference_no_decay`]).
    #[must_use]
    pub fn with_reference_no_decay(mut self, on: bool) -> Self {
        self.reference_no_decay = on;
        self
    }

    #[must_use]
    pub fn with_lr(mut self, lr: f32) -> Self {
        assert!(lr.is_finite() && lr >= 0.0, "lr must be finite and >= 0");
        self.lr = lr;
        self
    }

    #[must_use]
    pub fn with_betas(mut self, beta1: f32, beta2: f32) -> Self {
        assert!(
            (0.0..1.0).contains(&beta1) && (0.0..1.0).contains(&beta2),
            "betas must be in [0, 1), got beta1={beta1} beta2={beta2}"
        );
        self.beta1 = beta1;
        self.beta2 = beta2;
        self
    }

    #[must_use]
    pub fn with_eps(mut self, eps: f32) -> Self {
        assert!(eps.is_finite() && eps > 0.0, "eps must be finite and > 0");
        self.eps = eps;
        self
    }

    #[must_use]
    pub fn with_weight_decay(mut self, wd: f32) -> Self {
        assert!(
            wd.is_finite() && wd >= 0.0,
            "weight_decay must be finite and >= 0"
        );
        self.weight_decay = wd;
        self
    }

    /// Reset `m`, `v`, and step counter — useful between training phases or
    /// after a checkpoint load that doesn't include optimizer state.
    pub fn zero_state(&mut self, stream: &Arc<CudaStream>) -> Result<(), String> {
        self.m.zero(stream)?;
        self.v.zero(stream)?;
        self.step = 0;
        Ok(())
    }

    /// Save (step, lr) — not the f32 buffers. Use `m`/`v` field access for
    /// full state-dict (download with `to_cpu()`).
    pub fn state(&self) -> (u64, f32) {
        (self.step, self.lr)
    }

    /// Download the full optimizer state for a bit-continuous resume:
    /// the moment buffers, the step counter, and the hyperparameters
    /// that change the update math. `lr` is deliberately NOT included —
    /// the learning rate belongs to the caller's schedule, which
    /// re-applies it after a load (the in-house trainers set it per
    /// accumulation window anyway).
    pub fn export_state(&self, stream: &Arc<CudaStream>) -> Result<AdamWStateBlob, String> {
        Ok(AdamWStateBlob {
            m: self.m.to_cpu(stream)?,
            v: self.v.to_cpu(stream)?,
            step: self.step,
            beta1: self.beta1,
            beta2: self.beta2,
            eps: self.eps,
            weight_decay: self.weight_decay,
            reference_no_decay: self.reference_no_decay,
        })
    }

    /// Upload a previously exported state, adopting its step counter and
    /// update hyperparameters. Errs on a moment-length mismatch — that
    /// means the blob belongs to a different parameterization and a
    /// "resume" from it would be silent corruption, not a resume.
    pub fn import_state(
        &mut self,
        stream: &Arc<CudaStream>,
        blob: &AdamWStateBlob,
    ) -> Result<(), String> {
        if blob.m.len() != self.m.len() || blob.v.len() != self.v.len() {
            return Err(format!(
                "AdamW state mismatch: blob m/v = {}/{} elements, optimizer = {}/{} — \
                 the blob belongs to a different parameterization",
                blob.m.len(),
                blob.v.len(),
                self.m.len(),
                self.v.len()
            ));
        }
        self.m.upload(stream, &blob.m)?;
        self.v.upload(stream, &blob.v)?;
        self.step = blob.step;
        self.beta1 = blob.beta1;
        self.beta2 = blob.beta2;
        self.eps = blob.eps;
        self.weight_decay = blob.weight_decay;
        self.reference_no_decay = blob.reference_no_decay;
        Ok(())
    }

    /// Run one fused AdamW update on a single tensor. Caller supplies the
    /// device pointers (weight + grad + matching `m`/`v` slices) bundled
    /// in [`AdamWParamPtrs`].
    ///
    /// You normally want [`step_m1`] / [`step_m3`]; this is the low-level
    /// building block they use.
    pub fn step_one(
        &self,
        ctx: &GpuCtx,
        adamw_kernel: &CudaFunction,
        ptrs: AdamWParamPtrs,
        len: usize,
        bias_c1: f32,
        bias_c2: f32,
    ) -> Result<(), String> {
        if len == 0 {
            return Ok(());
        }
        let n = len as i32;
        let cfg = grid_1d(len);
        let mut bld = ctx.stream.launch_builder(adamw_kernel);
        bld.arg(&ptrs.weight);
        bld.arg(&ptrs.grad);
        bld.arg(&ptrs.m);
        bld.arg(&ptrs.v);
        bld.arg(&self.lr);
        bld.arg(&self.beta1);
        bld.arg(&self.beta2);
        bld.arg(&self.eps);
        bld.arg(&self.weight_decay);
        bld.arg(&bias_c1);
        bld.arg(&bias_c2);
        bld.arg(&n);
        unsafe { bld.launch(cfg) }.map_err(|e| format!("adamw_step_f32: {e:?}"))?;
        Ok(())
    }

    /// CUDA-Graph-capturable variant of [`Self::step_one`]. Bias factors
    /// are read from a 2-element device buffer (see [`AdamWBiasFactors`]),
    /// which the CPU updates BEFORE each graph replay.
    pub fn step_one_capturable(
        &self,
        ctx: &GpuCtx,
        adamw_kernel: &CudaFunction,
        ptrs: AdamWParamPtrs,
        bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
        len: usize,
    ) -> Result<(), String> {
        self.step_one_capturable_wd(
            ctx,
            adamw_kernel,
            ptrs,
            bias_factors_ptr,
            len,
            self.weight_decay,
        )
    }

    /// [`Self::step_one_capturable`] with an explicit per-tensor weight
    /// decay — the mechanism behind the reference no-decay groups (the
    /// decay coefficient is a per-launch scalar, so per-tensor decay costs
    /// nothing).
    pub fn step_one_capturable_wd(
        &self,
        ctx: &GpuCtx,
        adamw_kernel: &CudaFunction,
        ptrs: AdamWParamPtrs,
        bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
        len: usize,
        weight_decay: f32,
    ) -> Result<(), String> {
        if len == 0 {
            return Ok(());
        }
        let n = len as i32;
        let cfg = grid_1d(len);
        let mut bld = ctx.stream.launch_builder(adamw_kernel);
        bld.arg(&ptrs.weight);
        bld.arg(&ptrs.grad);
        bld.arg(&ptrs.m);
        bld.arg(&ptrs.v);
        bld.arg(&self.lr);
        bld.arg(&self.beta1);
        bld.arg(&self.beta2);
        bld.arg(&self.eps);
        bld.arg(&weight_decay);
        bld.arg(&bias_factors_ptr);
        bld.arg(&n);
        unsafe { bld.launch(cfg) }.map_err(|e| format!("adamw_step_f32_capturable: {e:?}"))?;
        Ok(())
    }

    /// Pre-compute `(bias_c1, bias_c2)` for the *next* step (i.e. after
    /// incrementing `self.step`). Returns the new step number and the two
    /// bias-correction multipliers used inside the kernel.
    pub fn advance(&mut self) -> (u64, f32, f32) {
        self.step += 1;
        // `powi` takes i32 for the exponent. Clamp to a step count beyond
        // which `β^t` is already below f64 round-off (≈ 1e-300 at step
        // ~3000 for β=0.9, step ~700k for β=0.999). 2^30 is ≈ 1.07B, well
        // inside i32 range and well past any realistic training horizon.
        // This avoids the silent overflow that cast `u64 as i32` produced
        // at step ≥ 2^31 (negative exponent → garbage bias factors).
        let t = self.step.min(1 << 30) as i32;
        let denom1 = 1.0 - (self.beta1 as f64).powi(t);
        let denom2 = 1.0 - (self.beta2 as f64).powi(t);
        let bias_c1 = (1.0 / denom1.max(1e-30)) as f32;
        let bias_c2 = (1.0 / denom2.max(1e-30)) as f32;
        (self.step, bias_c1, bias_c2)
    }
}

/// Iterate (weight_buffer, grad_slice) pairs in arena order and launch one
/// `adamw_step_f32` per tensor. The flat `m`/`v` are sliced at the same
/// offset that `grad_slice` has into `grads.flat`.
///
/// Caller must ensure the `(weight, grad)` pairs are fed in the SAME order
/// as `GpuMambaGrads::new` wrote them, so the m/v offsets align.
pub fn run_pairs(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &GpuAdamW,
    bias_c1: f32,
    bias_c2: f32,
    flat_grad_base: cudarc::driver::sys::CUdeviceptr,
    pairs: &[(&GpuBuffer, &GradSlice)],
) -> Result<(), String> {
    let m_base = adam.m.cached_ptr();
    let v_base = adam.v.cached_ptr();
    for (w, g) in pairs {
        // Skip empty master tensors. `MambaWeights` clears `input_proj_w`
        // to zero-length for HF Mamba's identity input projection — the
        // grad arena still reserves a slot for layout symmetry, but
        // there's nothing to update. Skipping leaves m/v at zero, which
        // is the correct AdamW state for an absent param.
        if w.is_empty() {
            continue;
        }
        if g.len() != w.len() {
            return Err(format!(
                "adamw: weight/grad len mismatch: w={} g={}",
                w.len(),
                g.len()
            ));
        }
        let g_ptr = g.ptr();
        let off_bytes = g_ptr - flat_grad_base;
        // Element offset (f32 = 4 bytes).
        let off_elems = off_bytes / 4;
        let m_ptr = m_base + off_bytes;
        let v_ptr = v_base + off_bytes;
        debug_assert!(
            off_elems as usize + g.len() <= adam.m.len(),
            "adamw m/v slice OOB: off_elems={off_elems} len={} m.len={}",
            g.len(),
            adam.m.len()
        );
        adam.step_one(
            ctx,
            adamw_kernel,
            AdamWParamPtrs {
                weight: w.cached_ptr(),
                grad: g_ptr,
                m: m_ptr,
                v: v_ptr,
            },
            g.len(),
            bias_c1,
            bias_c2,
        )?;
    }
    Ok(())
}

/// CUDA-Graph-capturable variant of [`run_pairs`]. Reads bias factors
/// from `bias_factors_ptr` (a 2-element device buffer) instead of taking
/// scalars. Caller is responsible for writing fresh `(bc1, bc2)` into that
/// buffer BEFORE each graph replay.
pub fn run_pairs_capturable(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &GpuAdamW,
    bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
    flat_grad_base: cudarc::driver::sys::CUdeviceptr,
    pairs: &[(&GpuBuffer, &GradSlice)],
) -> Result<(), String> {
    for pair in pairs {
        run_one_capturable_wd(
            ctx,
            adamw_kernel,
            adam,
            bias_factors_ptr,
            flat_grad_base,
            *pair,
            adam.weight_decay,
        )?;
    }
    Ok(())
}

/// One capturable AdamW launch for a (weight, grad-slice) pair with an
/// explicit weight decay. Shared body of [`run_pairs_capturable`] and the
/// per-tensor no-decay paths in [`step_m1_capturable`] /
/// [`step_m3_capturable`].
fn run_one_capturable_wd(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &GpuAdamW,
    bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
    flat_grad_base: cudarc::driver::sys::CUdeviceptr,
    (w, g): (&GpuBuffer, &GradSlice),
    weight_decay: f32,
) -> Result<(), String> {
    // Skip empty master tensors (HF Mamba identity input_proj). See
    // `run_pairs` for the rationale.
    if w.is_empty() {
        return Ok(());
    }
    if g.len() != w.len() {
        return Err(format!(
            "adamw: weight/grad len mismatch: w={} g={}",
            w.len(),
            g.len()
        ));
    }
    let g_ptr = g.ptr();
    let off_bytes = g_ptr - flat_grad_base;
    let off_elems = off_bytes / 4;
    let m_ptr = adam.m.cached_ptr() + off_bytes;
    let v_ptr = adam.v.cached_ptr() + off_bytes;
    debug_assert!(
        off_elems as usize + g.len() <= adam.m.len(),
        "adamw m/v slice OOB"
    );
    adam.step_one_capturable_wd(
        ctx,
        adamw_kernel,
        AdamWParamPtrs {
            weight: w.cached_ptr(),
            grad: g_ptr,
            m: m_ptr,
            v: v_ptr,
        },
        bias_factors_ptr,
        g.len(),
        weight_decay,
    )
}

/// Mamba SSM backbone AdamW step. Iterates the per-tensor master weights in
/// the SAME order as `GpuMambaGrads::new` wrote them into the flat arena,
/// so the m/v offsets line up element-for-element with grads.flat.
pub fn step_m1(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &mut GpuAdamW,
    weights: &mut crate::mamba_ssm::gpu::weights::GpuMambaTrainWeights,
    grads: &crate::mamba_ssm::gpu::weights::GpuMambaGrads,
) -> Result<(), String> {
    let (_, bc1, bc2) = adam.advance();
    let flat_base = grads.flat.cached_ptr();

    // Build paired iterator in the EXACT layout of `GpuMambaGrads::new`:
    // input_proj_w, input_proj_b, [layers...], norm_f_weight.
    let mut pairs: Vec<(&GpuBuffer, &GradSlice)> =
        Vec::with_capacity(3 + 10 * weights.layers.len());
    pairs.push((&weights.input_proj_w, &grads.input_proj_w));
    pairs.push((&weights.input_proj_b, &grads.input_proj_b));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        pairs.push((&lw.norm_weight, &lg.norm_weight));
        pairs.push((&lw.in_proj_w, &lg.in_proj_w));
        pairs.push((&lw.conv1d_weight, &lg.conv1d_weight));
        pairs.push((&lw.conv1d_bias, &lg.conv1d_bias));
        pairs.push((&lw.x_proj_w, &lg.x_proj_w));
        pairs.push((&lw.dt_proj_w, &lg.dt_proj_w));
        pairs.push((&lw.dt_proj_b, &lg.dt_proj_b));
        pairs.push((&lw.a_log, &lg.a_log));
        pairs.push((&lw.d_param, &lg.d_param));
        pairs.push((&lw.out_proj_w, &lg.out_proj_w));
    }
    pairs.push((&weights.norm_f_weight, &grads.norm_f_weight));

    run_pairs(ctx, adamw_kernel, adam, bc1, bc2, flat_base, &pairs)
}

/// CUDA-Graph-capturable variant of [`step_m1`]. Bias factors come from
/// the 2-element device buffer `bias_factors_ptr`, which the CPU rewrites
/// before each graph replay.
pub fn step_m1_capturable(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &GpuAdamW,
    bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
    weights: &mut crate::mamba_ssm::gpu::weights::GpuMambaTrainWeights,
    grads: &crate::mamba_ssm::gpu::weights::GpuMambaGrads,
) -> Result<(), String> {
    let flat_base = grads.flat.cached_ptr();
    // (weight, grad, reference-no-decay?). The no-decay group is the set the
    // reference implementation marks `_no_weight_decay`: a_log, D, dt bias,
    // plus every RMSNorm scale — active only when `adam.reference_no_decay`.
    let mut pairs: Vec<(&GpuBuffer, &GradSlice, bool)> =
        Vec::with_capacity(3 + 10 * weights.layers.len());
    pairs.push((&weights.input_proj_w, &grads.input_proj_w, false));
    pairs.push((&weights.input_proj_b, &grads.input_proj_b, false));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        pairs.push((&lw.norm_weight, &lg.norm_weight, true));
        pairs.push((&lw.in_proj_w, &lg.in_proj_w, false));
        pairs.push((&lw.conv1d_weight, &lg.conv1d_weight, false));
        pairs.push((&lw.conv1d_bias, &lg.conv1d_bias, false));
        pairs.push((&lw.x_proj_w, &lg.x_proj_w, false));
        pairs.push((&lw.dt_proj_w, &lg.dt_proj_w, false));
        pairs.push((&lw.dt_proj_b, &lg.dt_proj_b, true));
        pairs.push((&lw.a_log, &lg.a_log, true));
        pairs.push((&lw.d_param, &lg.d_param, true));
        pairs.push((&lw.out_proj_w, &lg.out_proj_w, false));
    }
    pairs.push((&weights.norm_f_weight, &grads.norm_f_weight, true));
    for (w, g, no_decay) in pairs {
        let wd = if no_decay && adam.reference_no_decay {
            0.0
        } else {
            adam.weight_decay
        };
        run_one_capturable_wd(
            ctx,
            adamw_kernel,
            adam,
            bias_factors_ptr,
            flat_base,
            (w, g),
            wd,
        )?;
    }
    Ok(())
}

/// CUDA-Graph-capturable variant of [`step_m3`].
pub fn step_m3_capturable(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &GpuAdamW,
    bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
    weights: &mut crate::mamba3_siso::gpu::weights::GpuMamba3Weights,
    grads: &crate::mamba3_siso::gpu::weights::GpuMamba3Grads,
) -> Result<(), String> {
    let flat_base = grads.flat.cached_ptr();
    // M3 no-decay group (mirror of the M1 set's principle): dt bias, D,
    // every norm scale AND every bias — the reference parameter grouping
    // exempts all params named `*bias`, which covers input_proj_b and the
    // all-ones B/C biases (decay would walk them out of the positive
    // regime). M3 has no fixed a_log (A is input-dependent). Active only
    // when `adam.reference_no_decay`.
    let mut pairs: Vec<(&GpuBuffer, &GradSlice, bool)> =
        Vec::with_capacity(3 + 10 * weights.layers.len());
    pairs.push((&weights.input_proj_w, &grads.input_proj_w, false));
    pairs.push((&weights.input_proj_b, &grads.input_proj_b, true));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        pairs.push((&lw.norm_weight, &lg.norm_weight, true));
        pairs.push((&lw.in_proj_w, &lg.in_proj_w, false));
        pairs.push((&lw.dt_bias, &lg.dt_bias, true));
        pairs.push((&lw.b_norm_weight, &lg.b_norm_weight, true));
        pairs.push((&lw.c_norm_weight, &lg.c_norm_weight, true));
        pairs.push((&lw.b_bias, &lg.b_bias, true));
        pairs.push((&lw.c_bias, &lg.c_bias, true));
        pairs.push((&lw.d_param, &lg.d_param, true));
        pairs.push((&lw.norm_gate_weight, &lg.norm_gate_weight, true));
        pairs.push((&lw.out_proj_w, &lg.out_proj_w, false));
    }
    pairs.push((&weights.norm_f_weight, &grads.norm_f_weight, true));
    for (w, g, no_decay) in pairs {
        let wd = if no_decay && adam.reference_no_decay {
            0.0
        } else {
            adam.weight_decay
        };
        run_one_capturable_wd(
            ctx,
            adamw_kernel,
            adam,
            bias_factors_ptr,
            flat_base,
            (w, g),
            wd,
        )?;
    }
    Ok(())
}

/// Mamba-3 backbone AdamW step. Same idea as [`step_m1`] but for the M3
/// weight set (`GpuMamba3Weights` / `GpuMamba3Grads`).
pub fn step_m3(
    ctx: &GpuCtx,
    adamw_kernel: &CudaFunction,
    adam: &mut GpuAdamW,
    weights: &mut crate::mamba3_siso::gpu::weights::GpuMamba3Weights,
    grads: &crate::mamba3_siso::gpu::weights::GpuMamba3Grads,
) -> Result<(), String> {
    let (_, bc1, bc2) = adam.advance();
    let flat_base = grads.flat.cached_ptr();

    let mut pairs: Vec<(&GpuBuffer, &GradSlice)> =
        Vec::with_capacity(3 + 10 * weights.layers.len());
    pairs.push((&weights.input_proj_w, &grads.input_proj_w));
    pairs.push((&weights.input_proj_b, &grads.input_proj_b));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        pairs.push((&lw.norm_weight, &lg.norm_weight));
        pairs.push((&lw.in_proj_w, &lg.in_proj_w));
        pairs.push((&lw.dt_bias, &lg.dt_bias));
        pairs.push((&lw.b_norm_weight, &lg.b_norm_weight));
        pairs.push((&lw.c_norm_weight, &lg.c_norm_weight));
        pairs.push((&lw.b_bias, &lg.b_bias));
        pairs.push((&lw.c_bias, &lg.c_bias));
        pairs.push((&lw.d_param, &lg.d_param));
        pairs.push((&lw.norm_gate_weight, &lg.norm_gate_weight));
        pairs.push((&lw.out_proj_w, &lg.out_proj_w));
    }
    pairs.push((&weights.norm_f_weight, &grads.norm_f_weight));

    run_pairs(ctx, adamw_kernel, adam, bc1, bc2, flat_base, &pairs)
}

#[cfg(test)]
mod cpu_state_tests {
    #[test]
    fn bias_correction_step_one() {
        let beta1 = 0.9_f64;
        let beta2 = 0.999_f64;
        let bc1 = (1.0 / (1.0 - beta1.powi(1))) as f32;
        let bc2 = (1.0 / (1.0 - beta2.powi(1))) as f32;
        // At t=1: bc1 = 1/(1-0.9) = 10, bc2 = 1/(1-0.999) = 1000.
        assert!((bc1 - 10.0).abs() < 1e-3, "bc1={bc1}");
        assert!((bc2 - 1000.0).abs() < 1e-1, "bc2={bc2}");
    }

    #[test]
    fn bias_correction_step_large() {
        // At t=2000 (post-warmup), both bias factors → ~1.
        let beta1 = 0.9_f64;
        let beta2 = 0.999_f64;
        let bc1 = (1.0 / (1.0 - beta1.powi(2000))) as f32;
        let bc2 = (1.0 / (1.0 - beta2.powi(2000))) as f32;
        assert!(bc1 < 1.001, "bc1={bc1}");
        assert!(bc2 < 1.2, "bc2={bc2}");
    }
}

/// One tensor's coordinates for the fused multi-tensor AdamW step.
#[derive(Clone, Copy, Debug)]
pub struct AdamWTensorSpec {
    /// f32 master weight base pointer.
    pub weight: cudarc::driver::sys::CUdeviceptr,
    /// Gradient slice base pointer (inside the flat arena).
    pub grad: cudarc::driver::sys::CUdeviceptr,
    /// Optional typed shadow to store `FROM_F(new_p)` into (0 = none).
    pub out: cudarc::driver::sys::CUdeviceptr,
    /// Element size of the shadow in bytes (2 = bf16/f16; ignored when
    /// `out == 0`).
    pub out_elt_bytes: usize,
    /// The shadow slot is f32 (a "f32 stays f32" tensor). The fused
    /// kernel then writes it directly instead of the caller running a
    /// per-tensor device-to-device copy after the step.
    pub out_is_f32: bool,
    pub len: usize,
    /// Member of the reference no-decay group (a_log, D, dt bias, norm
    /// scales) — decays only when `reference_no_decay` is off.
    pub no_decay: bool,
}

/// Chunk table for `adamw_step_multi_*`: built ONCE at construction (the
/// flat-arena layout is static for the life of a trainer) and replayed
/// every step — 243 per-tensor launches become one kernel.
pub struct AdamWMultiPlan {
    table: crate::mamba_ssm::gpu::buffers::GpuByteBuffer,
    n_chunks: usize,
}

/// Elements per chunk. 8k x 4B = 32 KB of f32 per block: at 256 threads
/// that is 32 elements per thread instead of 256, and the grid grows from
/// the hundreds into the thousands, so the machine fills instead of
/// running a handful of long serial blocks alongside eight single-block
/// tails. Chunking is a pure work split - AdamW is elementwise with no
/// cross-element interaction, so the size never touches a single bit.
pub const ADAMW_MULTI_CHUNK: usize = 8_192;

/// Build the device chunk table. `flat_base` is `grads.flat` base — the
/// m/v slices mirror the grad offsets exactly (same arena layout).
pub fn build_multi_plan(
    stream: &Arc<CudaStream>,
    adam: &GpuAdamW,
    flat_base: cudarc::driver::sys::CUdeviceptr,
    specs: &[AdamWTensorSpec],
    reference_no_decay: bool,
    weight_decay: f32,
) -> Result<AdamWMultiPlan, String> {
    let mut bytes: Vec<u8> = Vec::new();
    let mut n_chunks = 0usize;
    let m_base = adam.m.cached_ptr();
    let v_base = adam.v.cached_ptr();
    for spec in specs {
        if spec.len == 0 {
            // Empty master tensors (HF identity input_proj) hold no slot
            // in the arena at all — no chunk, no phantom m/v stepping.
            continue;
        }
        let off_bytes = spec.grad - flat_base;
        debug_assert!(
            (off_bytes / 4) as usize + spec.len <= adam.m.len(),
            "adamw multi plan: m/v slice OOB"
        );
        let wd = if spec.no_decay && reference_no_decay {
            0.0f32
        } else {
            weight_decay
        };
        let mut start = 0usize;
        while start < spec.len {
            let n = (spec.len - start).min(ADAMW_MULTI_CHUNK);
            let b4 = (start * 4) as u64;
            let entry: [u64; 6] = [
                spec.weight + b4,
                spec.grad + b4,
                m_base + off_bytes + b4,
                v_base + off_bytes + b4,
                if spec.out == 0 {
                    0
                } else {
                    spec.out + (start * spec.out_elt_bytes) as u64
                },
                // Bit 31 tags an f32 shadow; the chunk length needs 17.
                (n as u64) | (u64::from(spec.out_is_f32) << 31) | (u64::from(wd.to_bits()) << 32),
            ];
            for w in entry {
                bytes.extend_from_slice(&w.to_le_bytes());
            }
            n_chunks += 1;
            start += n;
        }
    }
    if n_chunks == 0 {
        return Err("adamw multi plan: no non-empty tensors".to_string());
    }
    let mut table = crate::mamba_ssm::gpu::buffers::GpuByteBuffer::zeros(stream, bytes.len())?;
    table.upload_bytes(stream, &bytes)?;
    Ok(AdamWMultiPlan { table, n_chunks })
}

/// Launch the fused step: one kernel over every chunk. `bias_factors_ptr`
/// is the 3-element `[bc1, bc2, lr]` device buffer.
pub fn step_multi(
    ctx: &GpuCtx,
    kernel: &CudaFunction,
    plan: &AdamWMultiPlan,
    adam: &GpuAdamW,
    bias_factors_ptr: cudarc::driver::sys::CUdeviceptr,
) -> Result<(), String> {
    let table_ptr = plan.table.cached_ptr();
    let cfg = cudarc::driver::LaunchConfig {
        grid_dim: (plan.n_chunks as u32, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut b = ctx.stream.launch_builder(kernel);
    b.arg(&table_ptr);
    b.arg(&adam.beta1);
    b.arg(&adam.beta2);
    b.arg(&adam.eps);
    b.arg(&bias_factors_ptr);
    unsafe { b.launch(cfg) }
        .map(|_| ())
        .map_err(|e| format!("adamw_step_multi: {e:?}"))
}

/// The M1 tensor walk in the EXACT `GpuMambaGrads::new` layout, f32 lane
/// (no typed shadow — `out = 0`).
pub fn m1_specs(
    weights: &crate::mamba_ssm::gpu::weights::GpuMambaTrainWeights,
    grads: &crate::mamba_ssm::gpu::weights::GpuMambaGrads,
) -> Vec<AdamWTensorSpec> {
    let f32_spec = |w: &GpuBuffer, g: &GradSlice, no_decay: bool| AdamWTensorSpec {
        weight: w.cached_ptr(),
        grad: g.ptr(),
        out: 0,
        out_elt_bytes: 0,
        out_is_f32: false,
        len: w.len(),
        no_decay,
    };
    let mut specs = Vec::with_capacity(3 + 10 * weights.layers.len());
    specs.push(f32_spec(&weights.input_proj_w, &grads.input_proj_w, false));
    specs.push(f32_spec(&weights.input_proj_b, &grads.input_proj_b, false));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        specs.push(f32_spec(&lw.norm_weight, &lg.norm_weight, true));
        specs.push(f32_spec(&lw.in_proj_w, &lg.in_proj_w, false));
        specs.push(f32_spec(&lw.conv1d_weight, &lg.conv1d_weight, false));
        specs.push(f32_spec(&lw.conv1d_bias, &lg.conv1d_bias, false));
        specs.push(f32_spec(&lw.x_proj_w, &lg.x_proj_w, false));
        specs.push(f32_spec(&lw.dt_proj_w, &lg.dt_proj_w, false));
        specs.push(f32_spec(&lw.dt_proj_b, &lg.dt_proj_b, true));
        specs.push(f32_spec(&lw.a_log, &lg.a_log, true));
        specs.push(f32_spec(&lw.d_param, &lg.d_param, true));
        specs.push(f32_spec(&lw.out_proj_w, &lg.out_proj_w, false));
    }
    specs.push(f32_spec(&weights.norm_f_weight, &grads.norm_f_weight, true));
    specs
}

/// Mixed-lane walk: identical order, and the four BULK weights carry
/// their typed compute slot as the fused shadow — the optimizer stores
/// `FROM_F(new_p)` in the same kernel, killing the per-step cast pass.
pub fn m1_specs_mixed(
    weights: &crate::mamba_ssm::gpu::weights_mixed_train::GpuMambaTrainMixedWeights,
    grads: &crate::mamba_ssm::gpu::weights::GpuMambaGrads,
) -> Vec<AdamWTensorSpec> {
    let out_elt = match weights.dtype {
        crate::mamba_ssm::gpu::dtype::WeightDtype::F32 => 4usize,
        _ => 2usize,
    };
    let mut specs = m1_specs(&weights.master, grads);
    // Positions in the walk: 0 input_proj_w (bulk), then per layer at
    // base+1 in_proj_w, base+4 x_proj_w, base+5 dt_proj_w, base+9
    // out_proj_w. Wire the shadows by walking the compute side in the
    // same order instead of indexing arithmetic.
    let mut wire = |idx: usize, slot: &crate::mamba_ssm::gpu::buffers::WeightSliceDyn| {
        debug_assert_eq!(specs[idx].len, slot.len_elems());
        specs[idx].out = slot.ptr();
        specs[idx].out_elt_bytes = out_elt;
    };
    wire(0, &weights.compute.input_proj_w);
    for (li, cw) in weights.compute.layers.iter().enumerate() {
        let base = 2 + li * 10;
        wire(base + 1, &cw.in_proj_w);
        wire(base + 4, &cw.x_proj_w);
        wire(base + 5, &cw.dt_proj_w);
        wire(base + 9, &cw.out_proj_w);
    }
    // The f32-stays-f32 shadows ride the same kernel: their compute copy
    // is the new master value verbatim, so the fused step writes it and
    // the per-tensor copy walk disappears. `a_log` is deliberately absent
    // - its compute copy has no training-path reader (the forward and
    // backward ride a_neg_all, recomputed from the master every step).
    let last = specs.len() - 1;
    let mut wire_f32 = |idx: usize, slot: &crate::mamba_ssm::gpu::buffers::WeightSliceDyn| {
        debug_assert_eq!(specs[idx].len, slot.len_elems());
        specs[idx].out = slot.ptr();
        specs[idx].out_elt_bytes = 4;
        specs[idx].out_is_f32 = true;
    };
    wire_f32(1, &weights.compute.input_proj_b);
    for (li, cw) in weights.compute.layers.iter().enumerate() {
        let base = 2 + li * 10;
        wire_f32(base, &cw.norm_weight);
        wire_f32(base + 2, &cw.conv1d_weight);
        wire_f32(base + 3, &cw.conv1d_bias);
        wire_f32(base + 6, &cw.dt_proj_b);
        wire_f32(base + 8, &cw.d_param);
    }
    wire_f32(last, &weights.compute.norm_f_weight);
    specs
}

/// The M3 tensor walk in the `Mamba3Grads` layout, f32 lane.
pub fn m3_specs(
    weights: &crate::mamba3_siso::gpu::weights::GpuMamba3Weights,
    grads: &crate::mamba3_siso::gpu::weights::GpuMamba3Grads,
) -> Vec<AdamWTensorSpec> {
    let f32_spec = |w: &GpuBuffer, g: &GradSlice, no_decay: bool| AdamWTensorSpec {
        weight: w.cached_ptr(),
        grad: g.ptr(),
        out: 0,
        out_elt_bytes: 0,
        out_is_f32: false,
        len: w.len(),
        no_decay,
    };
    let mut specs = Vec::with_capacity(3 + 10 * weights.layers.len());
    specs.push(f32_spec(&weights.input_proj_w, &grads.input_proj_w, false));
    // Every bias is no-decay per the reference parameter grouping (the
    // name rule in the official param_grouping catches input_proj_b and
    // the all-ones B/C biases; decaying an all-ones bias walks it out of
    // the positive regime the Mamba-3 ablation requires).
    specs.push(f32_spec(&weights.input_proj_b, &grads.input_proj_b, true));
    for (lw, lg) in weights.layers.iter().zip(&grads.layers) {
        specs.push(f32_spec(&lw.norm_weight, &lg.norm_weight, true));
        specs.push(f32_spec(&lw.in_proj_w, &lg.in_proj_w, false));
        specs.push(f32_spec(&lw.dt_bias, &lg.dt_bias, true));
        specs.push(f32_spec(&lw.b_norm_weight, &lg.b_norm_weight, true));
        specs.push(f32_spec(&lw.c_norm_weight, &lg.c_norm_weight, true));
        specs.push(f32_spec(&lw.b_bias, &lg.b_bias, true));
        specs.push(f32_spec(&lw.c_bias, &lg.c_bias, true));
        specs.push(f32_spec(&lw.d_param, &lg.d_param, true));
        specs.push(f32_spec(&lw.norm_gate_weight, &lg.norm_gate_weight, true));
        specs.push(f32_spec(&lw.out_proj_w, &lg.out_proj_w, false));
    }
    specs.push(f32_spec(&weights.norm_f_weight, &grads.norm_f_weight, true));
    specs
}

/// M3 mixed-lane walk: in_proj_w / out_proj_w carry their typed compute
/// shadows (the M3 bulk set), everything else stays f32-synced.
pub fn m3_specs_mixed(
    weights: &crate::mamba3_siso::gpu::weights_mixed_train::GpuMamba3TrainMixedWeights,
    grads: &crate::mamba3_siso::gpu::weights::GpuMamba3Grads,
) -> Vec<AdamWTensorSpec> {
    let out_elt = match weights.dtype {
        crate::mamba_ssm::gpu::dtype::WeightDtype::F32 => 4usize,
        _ => 2usize,
    };
    let mut specs = m3_specs(&weights.master, grads);
    let mut wire = |idx: usize, slot: &crate::mamba_ssm::gpu::buffers::WeightSliceDyn| {
        debug_assert_eq!(specs[idx].len, slot.len_elems());
        specs[idx].out = slot.ptr();
        specs[idx].out_elt_bytes = out_elt;
    };
    wire(0, &weights.compute.input_proj_w);
    for (li, cw) in weights.compute.layers.iter().enumerate() {
        let base = 2 + li * 10;
        wire(base + 1, &cw.in_proj_w);
        wire(base + 9, &cw.out_proj_w);
    }
    // f32-stays-f32 shadows ride the same kernel (see the M1 twin).
    let last = specs.len() - 1;
    let mut wire_f32 = |idx: usize, slot: &crate::mamba_ssm::gpu::buffers::WeightSliceDyn| {
        debug_assert_eq!(specs[idx].len, slot.len_elems());
        specs[idx].out = slot.ptr();
        specs[idx].out_elt_bytes = 4;
        specs[idx].out_is_f32 = true;
    };
    wire_f32(1, &weights.compute.input_proj_b);
    for (li, cw) in weights.compute.layers.iter().enumerate() {
        let base = 2 + li * 10;
        wire_f32(base, &cw.norm_weight);
        wire_f32(base + 2, &cw.dt_bias);
        wire_f32(base + 3, &cw.b_norm_weight);
        wire_f32(base + 4, &cw.c_norm_weight);
        wire_f32(base + 5, &cw.b_bias);
        wire_f32(base + 6, &cw.c_bias);
        wire_f32(base + 7, &cw.d_param);
        wire_f32(base + 8, &cw.norm_gate_weight);
    }
    wire_f32(last, &weights.compute.norm_f_weight);
    specs
}