mamba-rs 0.6.5

//! Mamba-3 SISO GPU **prefill**: process a whole prompt window through the
//! chunked SSD pipeline in one pass, leaving the persistent recurrent state
//! (SSM + K + V + RoPE angle) positioned for token-by-token decode.
//!
//! This is the inference twin of the training forward's chunked path with
//! the backward tape removed: no per-layer activation saves, one reusable
//! scratch across layers. Two window semantics, chosen per call:
//!
//! - **Stateless window** (`carry_state = false`): all four persistent
//!   state buffers are zeroed first — the classify-serve contract, where
//!   every page is its own world.
//! - **Continued window** (`carry_state = true`): the entering SSM state is
//!   seeded from the persistent buffers WITH the trapezoidal boundary fold
//!   (`v_state (x) k_state * dt0 * (1 - trap0)`): the discretization's beta
//!   term reaches one step back across the window seam, so carrying the
//!   SSM state alone would be wrong at exactly the first position of the
//!   new window. The RoPE angle accumulator continues natively.
//!
//! Launch configurations mirror the training forward launch-for-launch —
//! that is what makes prefill-vs-training-forward parity testable at the
//! bit level on the same device.

use super::kernels::Mamba3Kernels;
use super::state::{CUptr, GpuMamba3Dims, GpuMamba3StateBufs, GpuMamba3TargetScratch};
use super::weights::GpuMamba3WeightsInf;
use crate::mamba_ssm::gpu::buffers::GpuBuffer;
use crate::mamba_ssm::gpu::context::GpuCtx;
use crate::mamba_ssm::gpu::launch::{grid_1d, grid_norm};
use cudarc::driver::PushKernelArg;
use std::sync::Arc;

/// Chunk-pipeline intermediates the prefill needs beyond the no-save layer
/// scratch: per-window buffers consumed inside F6 and released to the next
/// layer immediately (nothing is kept for a backward).
pub struct Mamba3PrefillChunkScratch {
    /// a_val * dt per (b,t,h) — the log-decay increments.
    d_alpha: GpuBuffer, // [B*T*nh]
    /// K prescaled by the per-step scale (chunk preprocess output).
    k_scaled: GpuBuffer, // [B*T*nh*ds]
    /// Per-step qk_dot lane (chunk preprocess output).
    qk_dot: GpuBuffer, // [B*T*nh]
    /// Per-step scale lane (chunk preprocess output).
    scale: GpuBuffer, // [B*T*nh]
    /// Per-step gamma lane (chunk preprocess output).
    gamma_pre: GpuBuffer, // [B*T*nh]
    /// Intra-chunk cumulative decay.
    da_cumsum: GpuBuffer, // [B*n_chunks*nh*chunk_size]
    /// Chunk contributions, then (in-place) entering states per chunk.
    chunk_states: GpuBuffer, // [B*n_chunks*nh*hd*ds]
    /// State exiting the window per (b,h).
    final_states: GpuBuffer, // [B*nh*hd*ds]
    /// Entering SSM state incl. the trapezoidal boundary fold.
    init_state: GpuBuffer,
    /// fp64 staging for the chunk-parallel angle accumulation.
    angle_chunk_sums: crate::mamba_ssm::gpu::buffers::GpuByteBuffer,
    angle_chunk_carries: crate::mamba_ssm::gpu::buffers::GpuByteBuffer,
}

impl Mamba3PrefillChunkScratch {
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        dims: &GpuMamba3Dims,
    ) -> Result<Self, String> {
        let bt = dims.bt();
        let nh = dims.nheads;
        let hd = dims.headdim;
        let ds = dims.d_state;
        let nc = dims.n_chunks();
        let cs = dims.chunk_size();
        Ok(Self {
            d_alpha: GpuBuffer::zeros(stream, bt * nh)?,
            k_scaled: GpuBuffer::zeros(stream, bt * nh * ds)?,
            qk_dot: GpuBuffer::zeros(stream, bt * nh)?,
            scale: GpuBuffer::zeros(stream, bt * nh)?,
            gamma_pre: GpuBuffer::zeros(stream, bt * nh)?,
            da_cumsum: GpuBuffer::zeros(stream, dims.batch * nc * nh * cs)?,
            chunk_states: GpuBuffer::zeros(stream, dims.batch * nc * nh * hd * ds)?,
            final_states: GpuBuffer::zeros(stream, dims.batch * nh * hd * ds)?,
            init_state: GpuBuffer::zeros(stream, dims.batch * nh * hd * ds)?,
            angle_chunk_sums: crate::mamba_ssm::gpu::buffers::GpuByteBuffer::zeros(
                stream,
                dims.batch * nc * dims.nheads * dims.n_angles.max(1) * std::mem::size_of::<f64>(),
            )?,
            angle_chunk_carries: crate::mamba_ssm::gpu::buffers::GpuByteBuffer::zeros(
                stream,
                dims.batch * nc * dims.nheads * dims.n_angles.max(1) * std::mem::size_of::<f64>(),
            )?,
        })
    }
}

/// One prefill executor for a fixed `(batch, seq_len)` shape: the no-save
/// layer scratch plus the chunk intermediates, allocated once and reused
/// across calls (zero allocation on the hot path, CUDA-Graph friendly).
pub struct Mamba3Prefill {
    pub tgt: GpuMamba3TargetScratch,
    chunk: Mamba3PrefillChunkScratch,
    /// The shape every buffer above was sized for; `run` refuses any
    /// other dims — a longer window would silently index past the
    /// scratch allocations.
    sized_for: GpuMamba3Dims,
}

/// Output request for [`Mamba3Prefill::run_full`] - the m3 serve surface.
///
/// `full_temporal`, when present, receives the post-`norm_f` temporal for
/// ALL T positions (`[B * T * d_model]`) - the same buffer content the m3
/// trainer forward emits, bit for bit. `pooled_sum`, when present,
/// receives the COLUMN SUM over T (`[d_model]`, batch = 1 only): the
/// on-device colsum walks ascending t accumulating pure f32 adds from
/// 0.0, so `download / T` reproduces a CPU mean pool bit for bit while
/// the transfer drops from the full temporal to 1.5 KB.
pub struct Mamba3PrefillOutputs<'a> {
    /// `[B * d_model]` - final-timestep post-norm hidden state.
    pub last_hidden: &'a mut GpuBuffer,
    /// `[B * T * d_model]` - optional post-norm_f output for every position.
    pub full_temporal: Option<&'a mut GpuBuffer>,
    /// `[d_model]` - optional column sum of the post-norm_f temporal.
    pub pooled_sum: Option<&'a mut GpuBuffer>,
}

/// Everything a prefill launch needs besides the executor itself.
pub struct Mamba3PrefillRun<'a> {
    pub ctx: &'a GpuCtx,
    pub kernels: &'a Mamba3Kernels,
    pub dims: &'a GpuMamba3Dims,
    pub weights: &'a GpuMamba3WeightsInf,
    /// `[B*T*mamba_input_dim]` (or `[B*T*d_model]` with an identity proj).
    pub mamba_input: &'a GpuBuffer,
    /// Empty `input_proj_w` on the checkpoint = identity input: skip the
    /// projection GEMM and feed the input straight to the layers.
    pub identity_proj: bool,
    /// `false` = stateless window (zero all four states first);
    /// `true` = continue from the persistent state across the seam.
    pub carry_state: bool,
}

/// Async D2D region copy between raw device pointers (stream-ordered, no
/// sync — the same primitive `GpuBuffer::copy_from_raw` uses, needed here
/// for per-layer slices of the persistent state buffers).
fn dtod_region(
    dst: CUptr,
    src: CUptr,
    floats: usize,
    stream: &Arc<cudarc::driver::CudaStream>,
) -> Result<(), String> {
    let result = unsafe {
        cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
            dst,
            src,
            floats * std::mem::size_of::<f32>(),
            stream.cu_stream(),
        )
    };
    if result != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
        return Err(format!(
            "prefill D2D region copy ({floats} floats): {result:?}"
        ));
    }
    Ok(())
}

impl Mamba3Prefill {
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        dims: &GpuMamba3Dims,
    ) -> Result<Self, String> {
        dims.validate_index_budget()?;
        Ok(Self {
            tgt: GpuMamba3TargetScratch::new(stream, dims)?,
            chunk: Mamba3PrefillChunkScratch::new(stream, dims)?,
            sized_for: *dims,
        })
    }

    /// Run the prompt window through every layer. On return
    /// `last_hidden` holds the FINAL-timestep post-norm hidden state
    /// `[B*d_model]` (the decode loop's next input to the lm head) and the
    /// persistent state buffers are positioned after the window's last
    /// token, ready for `step()`.
    pub fn run(
        &mut self,
        run: &Mamba3PrefillRun<'_>,
        states: GpuMamba3StateBufs<'_>,
        last_hidden: &mut GpuBuffer,
    ) -> Result<(), String> {
        self.run_full(
            run,
            states,
            Mamba3PrefillOutputs {
                last_hidden,
                full_temporal: None,
                pooled_sum: None,
            },
        )
    }

    /// [`Self::run`] with the all-T serve surface: identical kernel chain
    /// plus the optional full post-norm_f temporal and the on-device
    /// pooled column sum (see [`Mamba3PrefillOutputs`]).
    pub fn run_full(
        &mut self,
        run: &Mamba3PrefillRun<'_>,
        states: GpuMamba3StateBufs<'_>,
        outputs: Mamba3PrefillOutputs<'_>,
    ) -> Result<(), String> {
        let Mamba3PrefillOutputs {
            last_hidden,
            full_temporal,
            pooled_sum,
        } = outputs;
        let Mamba3PrefillRun {
            ctx,
            kernels: m3k,
            dims,
            weights,
            mamba_input,
            identity_proj,
            carry_state,
        } = *run;
        if *dims != self.sized_for {
            return Err(format!(
                "prefill executor was sized for {:?} but run was asked for {:?} —                  allocate a prefill for the shape you run",
                self.sized_for, dims
            ));
        }
        let tgt = &mut self.tgt;
        let ck = &mut self.chunk;
        let bt = dims.bt();
        let dm = dims.d_model;
        let di = dims.d_inner;
        let ds = dims.d_state;
        let nh = dims.nheads;
        let hd = dims.headdim;
        let ng = dims.ngroups;
        let ip = dims.in_proj_dim;
        let na = dims.n_angles;
        let na_alloc = na.max(1);
        let nc = dims.n_chunks();
        let cs = dims.chunk_size() as i32;
        let b_i = dims.batch as i32;
        let t_i = dims.seq_len as i32;
        let nh_i = nh as i32;
        let hd_i = hd as i32;
        let ds_i = ds as i32;
        let ng_i = ng as i32;
        let f32_sz = std::mem::size_of::<f32>() as u64;

        if !carry_state {
            states.ssm.zero(&ctx.stream)?;
            states.k.zero(&ctx.stream)?;
            states.v.zero(&ctx.stream)?;
            states.angle.zero(&ctx.stream)?;
        }

        // Input projection (or identity feed) into the working temporal.
        if identity_proj {
            tgt.temporal_work.copy_from_raw(mamba_input, &ctx.stream)?;
        } else {
            crate::mamba_ssm::gpu::blas::gpu_sgemm_forward_raw(
                ctx,
                &mut tgt.temporal_work,
                mamba_input,
                weights.input_proj_w.ptr(),
                Some(weights.input_proj_b.ptr()),
                (bt, dims.mamba_input_dim, dm),
            )?;
        }

        let ssm_base = states.ssm.raw_ptr(&ctx.stream);
        let k_base = states.k.raw_ptr(&ctx.stream);
        let v_base = states.v.raw_ptr(&ctx.stream);
        let a_base = states.angle.raw_ptr(&ctx.stream);

        for l in 0..dims.n_layers {
            let lw = &weights.layers[l];
            let ssm_ptr = ssm_base + (dims.batch * l * nh * hd * ds) as u64 * f32_sz;
            let k_ptr = k_base + (dims.batch * l * nh * ds) as u64 * f32_sz;
            let v_ptr = v_base + (dims.batch * l * nh * hd) as u64 * f32_sz;
            let a_ptr = a_base + (dims.batch * l * nh * na_alloc) as u64 * f32_sz;

            // Layer norm over the residual stream.
            tgt.residual
                .copy_from_raw(&tgt.temporal_work, &ctx.stream)?;
            {
                let bt_i = bt as i32;
                let dm_i = dm as i32;
                let eps: f32 = dims.rms_norm_eps;
                let nw = lw.norm_weight.ptr();
                let mut b = ctx.stream.launch_builder(&m3k.rmsnorm_fwd);
                b.arg(tgt.out_flat.inner_mut());
                b.arg(tgt.rms_discard.inner_mut());
                b.arg(tgt.residual.inner());
                b.arg(&nw);
                b.arg(&bt_i);
                b.arg(&dm_i);
                b.arg(&eps);
                unsafe { b.launch(grid_norm(bt, dm)) }
                    .map_err(|e| format!("prefill rmsnorm L{l}: {e:?}"))?;
            }
            crate::mamba_ssm::gpu::blas::gpu_sgemm_forward_raw(
                ctx,
                &mut tgt.proj_flat,
                &tgt.out_flat,
                lw.in_proj_w.ptr(),
                None,
                (bt, dm, ip),
            )?;
            // 8-way split + fused activations.
            {
                let n_i = bt as i32;
                let di_i = di as i32;
                let na_i = na as i32;
                let db = lw.dt_bias.ptr();
                let mut b = ctx.stream.launch_builder(&m3k.m3_split);
                b.arg(tgt.z.inner_mut());
                b.arg(tgt.x.inner_mut());
                b.arg(tgt.b_raw.inner_mut());
                b.arg(tgt.c_raw.inner_mut());
                b.arg(tgt.dt.inner_mut());
                b.arg(tgt.a_val.inner_mut());
                b.arg(tgt.trap.inner_mut());
                b.arg(tgt.angles_raw.inner_mut());
                b.arg(tgt.dd_dt_raw.inner_mut());
                b.arg(tgt.dd_a_raw.inner_mut());
                b.arg(tgt.trap_raw.inner_mut());
                b.arg(tgt.proj_flat.inner());
                b.arg(&db);
                b.arg(&dims.a_floor);
                b.arg(&n_i);
                b.arg(&di_i);
                b.arg(&ng_i);
                b.arg(&ds_i);
                b.arg(&nh_i);
                b.arg(&na_i);
                unsafe { b.launch(grid_1d(bt * ip)) }
                    .map_err(|e| format!("prefill m3_split L{l}: {e:?}"))?;
            }
            // BCNorm + bias for B and C.
            for (out, rms, raw, w, tag) in [
                (
                    &mut tgt.b_normed,
                    &mut tgt.b_rms,
                    &tgt.b_raw,
                    lw.b_norm_weight.ptr(),
                    "B",
                ),
                (
                    &mut tgt.c_normed,
                    &mut tgt.c_rms,
                    &tgt.c_raw,
                    lw.c_norm_weight.ptr(),
                    "C",
                ),
            ] {
                let n_i = bt as i32;
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: ((bt * ng) as u32, 1, 1),
                    block_dim: (ds as u32, 1, 1),
                    shared_mem_bytes: ds as u32 * 4,
                };
                let eps: f32 = dims.rms_norm_eps;
                let mut b = ctx.stream.launch_builder(&m3k.bcnorm_fwd);
                b.arg(out.inner_mut());
                b.arg(rms.inner_mut());
                b.arg(raw.inner());
                b.arg(&w);
                b.arg(&n_i);
                b.arg(&ng_i);
                b.arg(&ds_i);
                b.arg(&eps);
                unsafe { b.launch(cfg) }
                    .map_err(|e| format!("prefill bcnorm {tag} L{l}: {e:?}"))?;
            }
            // RoPE angle accumulation continues from the persistent
            // accumulator (zeroed above for a stateless window).
            if na > 0 {
                crate::mamba3_siso::gpu::forward::gpu_angle_chunked_fwd(
                    ctx,
                    m3k,
                    &mut tgt.angle_cumsum,
                    a_ptr,
                    &tgt.angles_raw,
                    &tgt.dt,
                    &ck.angle_chunk_sums,
                    &ck.angle_chunk_carries,
                    dims.batch,
                    dims.seq_len,
                    nh,
                    na,
                )?;
            }
            // Fused bias add (B + C) + RoPE: one launch, biased tensors
            // still materialize; n_angles == 0 passes through (replacing
            // the old copy branch).
            {
                let n_i = bt as i32;
                let na_i = na as i32;
                let bb_p = lw.b_bias.ptr();
                let cb_p = lw.c_bias.ptr();
                let mut b = ctx.stream.launch_builder(&m3k.m3_bias_rope_fwd);
                b.arg(tgt.b_biased.inner_mut());
                b.arg(tgt.c_biased.inner_mut());
                b.arg(tgt.k.inner_mut());
                b.arg(tgt.q.inner_mut());
                b.arg(tgt.b_normed.inner());
                b.arg(tgt.c_normed.inner());
                b.arg(&bb_p);
                b.arg(&cb_p);
                b.arg(tgt.angle_cumsum.inner());
                b.arg(&n_i);
                b.arg(&nh_i);
                b.arg(&ng_i);
                b.arg(&ds_i);
                b.arg(&na_i);
                unsafe { b.launch(grid_1d(bt * nh * ds)) }
                    .map_err(|e| format!("prefill bias_rope L{l}: {e:?}"))?;
            }
            // Trapezoidal coefficients.
            {
                let n_total = (bt * nh) as i32;
                let mut b = ctx.stream.launch_builder(&m3k.m3_compute_abg);
                b.arg(tgt.alpha.inner_mut());
                b.arg(tgt.beta.inner_mut());
                b.arg(tgt.gamma.inner_mut());
                b.arg(tgt.dt.inner());
                b.arg(tgt.a_val.inner());
                b.arg(tgt.trap.inner());
                b.arg(&n_total);
                unsafe { b.launch(grid_1d(bt * nh)) }
                    .map_err(|e| format!("prefill abg L{l}: {e:?}"))?;
            }
            // Chunked SSD pipeline (the training forward's F6, tape-free).
            {
                let n_total = (bt * nh) as i32;
                let mut b = ctx.stream.launch_builder(&m3k.elementwise_mul);
                b.arg(ck.d_alpha.inner_mut());
                b.arg(tgt.a_val.inner());
                b.arg(tgt.dt.inner());
                b.arg(&n_total);
                unsafe { b.launch(grid_1d(bt * nh)) }
                    .map_err(|e| format!("prefill adt L{l}: {e:?}"))?;
            }
            {
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: ((dims.batch * nc) as u32, nh as u32, 1),
                    block_dim: (dims.chunk_size() as u32, 1, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_preprocess_chunks);
                b.arg(ck.k_scaled.inner_mut());
                b.arg(ck.qk_dot.inner_mut());
                b.arg(ck.scale.inner_mut());
                b.arg(ck.gamma_pre.inner_mut());
                b.arg(tgt.k.inner());
                b.arg(tgt.q.inner());
                b.arg(tgt.dt.inner());
                b.arg(tgt.trap.inner());
                b.arg(&b_i);
                b.arg(&t_i);
                b.arg(&nh_i);
                b.arg(&ds_i);
                b.arg(&cs);
                unsafe { b.launch(cfg) }.map_err(|e| format!("prefill preprocess L{l}: {e:?}"))?;
            }
            {
                let block_x = nh.min(256) as u32;
                let grid_z = nh.div_ceil(block_x as usize) as u32;
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: (dims.batch as u32, nc as u32, grid_z),
                    block_dim: (block_x, 1, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_da_cumsum);
                b.arg(ck.da_cumsum.inner_mut());
                b.arg(ck.d_alpha.inner());
                b.arg(&b_i);
                b.arg(&t_i);
                b.arg(&nh_i);
                b.arg(&cs);
                unsafe { b.launch(cfg) }.map_err(|e| format!("prefill da_cumsum L{l}: {e:?}"))?;
            }
            {
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: ((dims.batch * nc) as u32, nh.div_ceil(2) as u32, 1),
                    block_dim: (hd as u32, 2, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_chunk_state_fwd);
                b.arg(ck.chunk_states.inner_mut());
                b.arg(tgt.x.inner());
                b.arg(ck.k_scaled.inner());
                b.arg(ck.da_cumsum.inner());
                b.arg(&b_i);
                b.arg(&t_i);
                b.arg(&nh_i);
                b.arg(&hd_i);
                b.arg(&ds_i);
                b.arg(&cs);
                unsafe { b.launch(cfg) }.map_err(|e| format!("prefill chunk_state L{l}: {e:?}"))?;
            }
            // Entering state: continued windows seed the inter-chunk scan
            // from the persistent SSM state plus the trapezoidal boundary
            // fold; stateless windows scan from zero (null pointer).
            let init_ptr: CUptr = if carry_state {
                dtod_region(
                    ck.init_state.cached_ptr(),
                    ssm_ptr,
                    dims.batch * nh * hd * ds,
                    &ctx.stream,
                )?;
                let dim = hd * ds;
                let block_x = dim.min(256) as u32;
                let grid_z = dim.div_ceil(block_x as usize) as u32;
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: (dims.batch as u32, nh as u32, grid_z),
                    block_dim: (block_x, 1, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_chunk_entering_state);
                b.arg(ck.init_state.inner_mut());
                b.arg(&k_ptr);
                b.arg(&v_ptr);
                b.arg(tgt.dt.inner());
                b.arg(tgt.trap.inner());
                b.arg(&b_i);
                b.arg(&nh_i);
                b.arg(&hd_i);
                b.arg(&ds_i);
                b.arg(&t_i);
                unsafe { b.launch(cfg) }
                    .map_err(|e| format!("prefill entering_state L{l}: {e:?}"))?;
                ck.init_state.cached_ptr()
            } else {
                0
            };
            {
                let dim = hd * ds;
                let block_x = dim.min(256) as u32;
                let grid_z = dim.div_ceil(block_x as usize) as u32;
                let nc_i = nc as i32;
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: (dims.batch as u32, nh as u32, grid_z),
                    block_dim: (block_x, 1, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_state_passing_fwd);
                b.arg(ck.chunk_states.inner_mut());
                b.arg(ck.final_states.inner_mut());
                b.arg(ck.da_cumsum.inner());
                b.arg(&init_ptr);
                b.arg(&b_i);
                b.arg(&nc_i);
                b.arg(&nh_i);
                b.arg(&hd_i);
                b.arg(&ds_i);
                b.arg(&cs);
                b.arg(&t_i);
                unsafe { b.launch(cfg) }
                    .map_err(|e| format!("prefill state_passing L{l}: {e:?}"))?;
            }
            {
                let dp = lw.d_param.ptr();
                let (coop, cfg) =
                    super::kernels::chunk_scan_cfg(dims.batch, nc, nh, hd, ds, dims.chunk_size());
                let kern = if coop {
                    &m3k.m3_chunk_scan_fwd_coop
                } else {
                    &m3k.m3_chunk_scan_fwd
                };
                let mut b = ctx.stream.launch_builder(kern);
                b.arg(tgt.y.inner_mut());
                b.arg(tgt.x.inner());
                b.arg(tgt.q.inner());
                b.arg(ck.k_scaled.inner());
                b.arg(ck.qk_dot.inner());
                b.arg(ck.da_cumsum.inner());
                b.arg(ck.chunk_states.inner());
                b.arg(&dp);
                b.arg(&b_i);
                b.arg(&t_i);
                b.arg(&nh_i);
                b.arg(&hd_i);
                b.arg(&ds_i);
                b.arg(&cs);
                unsafe { b.launch(cfg) }.map_err(|e| format!("prefill chunk_scan L{l}: {e:?}"))?;
            }
            // Exit states -> persistent buffers (SSM from the scan, final K
            // post-RoPE, final V = x at the window's last timestep).
            {
                let block_x = hd.max(ds) as u32;
                let cfg = cudarc::driver::LaunchConfig {
                    grid_dim: (dims.batch as u32, nh as u32, 1),
                    block_dim: (block_x, 1, 1),
                    shared_mem_bytes: 0,
                };
                let mut b = ctx.stream.launch_builder(&m3k.m3_writeback_parallel_states);
                b.arg(&ssm_ptr);
                b.arg(&k_ptr);
                b.arg(&v_ptr);
                b.arg(ck.final_states.inner());
                b.arg(tgt.k.inner());
                b.arg(tgt.x.inner());
                b.arg(&b_i);
                b.arg(&t_i);
                b.arg(&nh_i);
                b.arg(&hd_i);
                b.arg(&ds_i);
                unsafe { b.launch(cfg) }.map_err(|e| format!("prefill writeback L{l}: {e:?}"))?;
            }
            // Output gate + out_proj + residual.
            if dims.is_outproj_norm {
                assert!(
                    di <= 1024,
                    "d_inner ({di}) exceeds rmsnorm_gated shared memory limit"
                );
                let nw = lw.norm_gate_weight.ptr();
                let bt_i = bt as i32;
                let di_i = di as i32;
                let grid = cudarc::driver::LaunchConfig {
                    grid_dim: (bt as u32, 1, 1),
                    block_dim: (di as u32, 1, 1),
                    shared_mem_bytes: (di * std::mem::size_of::<f32>()) as u32,
                };
                let eps: f32 = dims.rms_norm_eps;
                let mut b = ctx.stream.launch_builder(&m3k.rmsnorm_gated_fwd);
                b.arg(tgt.gated.inner_mut());
                b.arg(tgt.rms_discard.inner_mut());
                b.arg(tgt.y.inner());
                b.arg(tgt.z.inner());
                b.arg(&nw);
                b.arg(&bt_i);
                b.arg(&di_i);
                b.arg(&hd_i);
                b.arg(&eps);
                unsafe { b.launch(grid) }.map_err(|e| format!("prefill gated L{l}: {e:?}"))?;
            } else {
                let n = (bt * di) as i32;
                let mut b = ctx.stream.launch_builder(&m3k.silu_gate_fwd);
                b.arg(tgt.gated.inner_mut());
                b.arg(tgt.y.inner());
                b.arg(tgt.z.inner());
                b.arg(&n);
                unsafe { b.launch(grid_1d(bt * di)) }
                    .map_err(|e| format!("prefill silu L{l}: {e:?}"))?;
            }
            crate::mamba_ssm::gpu::blas::gpu_sgemm_forward_raw(
                ctx,
                &mut tgt.out_flat,
                &tgt.gated,
                lw.out_proj_w.ptr(),
                None,
                (bt, di, dm),
            )?;
            {
                let ne = (bt * dm) as i32;
                let mut b = ctx.stream.launch_builder(&m3k.residual_add);
                b.arg(tgt.temporal_work.inner_mut());
                b.arg(tgt.out_flat.inner());
                b.arg(tgt.residual.inner());
                b.arg(&ne);
                unsafe { b.launch(grid_1d(bt * dm)) }
                    .map_err(|e| format!("prefill residual L{l}: {e:?}"))?;
            }
        }

        // Final norm over the whole window, then keep only the last
        // timestep's hidden state for the decode loop / lm head.
        tgt.residual
            .copy_from_raw(&tgt.temporal_work, &ctx.stream)?;
        {
            let nf = weights.norm_f_weight.ptr();
            let bt_i = bt as i32;
            let dm_i = dm as i32;
            let eps: f32 = dims.rms_norm_eps;
            let mut b = ctx.stream.launch_builder(&m3k.rmsnorm_fwd);
            b.arg(tgt.temporal_work.inner_mut());
            b.arg(tgt.rms_discard.inner_mut());
            b.arg(tgt.residual.inner());
            b.arg(&nf);
            b.arg(&bt_i);
            b.arg(&dm_i);
            b.arg(&eps);
            unsafe { b.launch(grid_norm(bt, dm)) }.map_err(|e| format!("prefill norm_f: {e:?}"))?;
        }
        {
            let dm_i = dm as i32;
            let mut b = ctx.stream.launch_builder(&m3k.gather_last_timestep);
            b.arg(last_hidden.inner_mut());
            b.arg(tgt.temporal_work.inner());
            b.arg(&b_i);
            b.arg(&t_i);
            b.arg(&dm_i);
            let grid = grid_1d(dims.batch * dm);
            unsafe { b.launch(grid) }.map_err(|e| format!("prefill gather_last: {e:?}"))?;
        }
        if let Some(full) = full_temporal {
            if full.len() != bt * dm {
                return Err(format!(
                    "m3 prefill outputs: full_temporal len {} != batch*seq_len*d_model = {}",
                    full.len(),
                    bt * dm
                ));
            }
            full.copy_from_raw(&tgt.temporal_work, &ctx.stream)?;
        }
        if let Some(pooled) = pooled_sum {
            if dims.batch != 1 {
                return Err(format!(
                    "m3 prefill pooled_sum: batch {} unsupported - the column sum \
                     would mix samples; pool per sample with batch = 1",
                    dims.batch
                ));
            }
            if pooled.len() != dm {
                return Err(format!(
                    "m3 prefill outputs: pooled_sum len {} != d_model = {dm}",
                    pooled.len()
                ));
            }
            // colsum_accumulate does db[j] += sum - zero first so the
            // result is exactly the sum over T.
            pooled.zero(&ctx.stream)?;
            let rows = bt as i32;
            let cols = dm as i32;
            let mut b = ctx.stream.launch_builder(&m3k.colsum_accumulate);
            b.arg(pooled.inner_mut());
            b.arg(tgt.temporal_work.inner());
            b.arg(&rows);
            b.arg(&cols);
            unsafe { b.launch(grid_1d(dm)) }
                .map_err(|e| format!("m3 prefill pooled colsum: {e:?}"))?;
        }
        Ok(())
    }
}

/// A captured CUDA graph of one prefill window over FIXED buffers: replay
/// re-runs the whole window at graph-launch cost. The GEMM-tier flags are
/// snapshotted at capture and asserted at replay — flipping the
/// batch-invariant / tensor-core routing after capture would silently
/// replay the OLD kernels, which is a numeric-route swap the crate treats
/// as a hard error.
pub struct Mamba3PrefillGraph {
    graph: cudarc::driver::CudaGraph,
    flags_at_capture: (bool, bool, bool),
    input_ptr: CUptr,
    ssm_ptr: CUptr,
    k_ptr: CUptr,
    v_ptr: CUptr,
    angle_ptr: CUptr,
    last_hidden_ptr: CUptr,
}

impl Mamba3PrefillGraph {
    /// Capture the window over the given fixed input/state/output buffers.
    /// Upload fresh bytes into the SAME input buffer before each replay.
    pub fn capture(
        prefill: &mut Mamba3Prefill,
        run: &Mamba3PrefillRun<'_>,
        mut states: GpuMamba3StateBufs<'_>,
        last_hidden: &mut GpuBuffer,
    ) -> Result<Self, String> {
        run.ctx.presize_bi_scratch()?;
        let flags_at_capture = run.ctx.gemm_flags();
        let input_ptr = run.mamba_input.cached_ptr();
        let ssm_ptr = states.ssm.cached_ptr();
        let k_ptr = states.k.cached_ptr();
        let v_ptr = states.v.cached_ptr();
        let angle_ptr = states.angle.cached_ptr();
        let last_hidden_ptr = last_hidden.cached_ptr();
        let graph =
            crate::mamba_ssm::gpu::graph_capture::capture_into_graph(&run.ctx.stream, || {
                prefill.run(run, states.reborrow(), last_hidden)
            })?;
        graph
            .upload()
            .map_err(|e| format!("prefill graph upload: {e:?}"))?;
        run.ctx.note_graph_capture();
        Ok(Self {
            graph,
            flags_at_capture,
            input_ptr,
            ssm_ptr,
            k_ptr,
            v_ptr,
            angle_ptr,
            last_hidden_ptr,
        })
    }

    /// Replay the captured window. `mamba_input` and ALL FOUR state
    /// buffers must be the SAME allocations the capture saw (the graph
    /// baked every device pointer in; a different buffer would leave
    /// the graph writing the old allocation while the caller reads the
    /// new one).
    pub fn replay(
        &self,
        ctx: &GpuCtx,
        mamba_input: &GpuBuffer,
        states: &GpuMamba3StateBufs<'_>,
        last_hidden: &GpuBuffer,
    ) -> Result<(), String> {
        if ctx.gemm_flags() != self.flags_at_capture {
            return Err(format!(
                "prefill graph replay refused: GEMM-tier flags changed since capture \
                 (captured {:?}, now {:?}) — a replay would silently run the old \
                 numeric route",
                self.flags_at_capture,
                ctx.gemm_flags()
            ));
        }
        if mamba_input.cached_ptr() != self.input_ptr
            || states.ssm.cached_ptr() != self.ssm_ptr
            || states.k.cached_ptr() != self.k_ptr
            || states.v.cached_ptr() != self.v_ptr
            || states.angle.cached_ptr() != self.angle_ptr
            || last_hidden.cached_ptr() != self.last_hidden_ptr
        {
            return Err(
                "prefill graph replay refused: input/state buffers differ from the \
                 captured allocations"
                    .to_string(),
            );
        }
        self.graph
            .launch()
            .map_err(|e| format!("prefill graph launch: {e:?}"))
    }
}

/// A captured CUDA graph of one POOLED prefill window over fixed buffers -
/// the m3 classify serve shape: state reset (carry_state = false zeroes
/// inside the capture, so every replay starts a fresh page) + the full
/// layer chain + norm_f + the on-device pooled column sum. Mirrors the M1
/// `PrefillPooledGraph` contract: upload the next page into the SAME
/// input buffer, replay, download the 1.5 KB pooled sum, divide by T on
/// the host.
pub struct Mamba3PrefillPooledGraph {
    graph: cudarc::driver::CudaGraph,
    flags_at_capture: (bool, bool, bool),
    input_ptr: CUptr,
    ssm_ptr: CUptr,
    k_ptr: CUptr,
    v_ptr: CUptr,
    angle_ptr: CUptr,
    pooled_ptr: CUptr,
}

impl Mamba3PrefillPooledGraph {
    /// Capture the pooled window. `run.carry_state` must be `false` -
    /// the state reset has to live INSIDE the graph for replays to score
    /// independent pages.
    pub fn capture(
        prefill: &mut Mamba3Prefill,
        run: &Mamba3PrefillRun<'_>,
        mut states: GpuMamba3StateBufs<'_>,
        last_hidden: &mut GpuBuffer,
        pooled_sum: &mut GpuBuffer,
    ) -> Result<Self, String> {
        if run.carry_state {
            return Err(
                "m3 pooled prefill graph: carry_state must be false - the state \
                 reset must be captured so each replay scores a fresh page"
                    .to_string(),
            );
        }
        run.ctx.presize_bi_scratch()?;
        let flags_at_capture = run.ctx.gemm_flags();
        let input_ptr = run.mamba_input.cached_ptr();
        let ssm_ptr = states.ssm.cached_ptr();
        let k_ptr = states.k.cached_ptr();
        let v_ptr = states.v.cached_ptr();
        let angle_ptr = states.angle.cached_ptr();
        let pooled_ptr = pooled_sum.cached_ptr();
        let graph =
            crate::mamba_ssm::gpu::graph_capture::capture_into_graph(&run.ctx.stream, || {
                prefill.run_full(
                    run,
                    states.reborrow(),
                    Mamba3PrefillOutputs {
                        last_hidden,
                        full_temporal: None,
                        pooled_sum: Some(pooled_sum),
                    },
                )
            })?;
        graph
            .upload()
            .map_err(|e| format!("m3 pooled prefill graph upload: {e:?}"))?;
        run.ctx.note_graph_capture();
        Ok(Self {
            graph,
            flags_at_capture,
            input_ptr,
            ssm_ptr,
            k_ptr,
            v_ptr,
            angle_ptr,
            pooled_ptr,
        })
    }

    /// Replay one page: same fixed buffers as capture, checked.
    pub fn replay(
        &self,
        ctx: &GpuCtx,
        mamba_input: &GpuBuffer,
        states: &GpuMamba3StateBufs<'_>,
        pooled_sum: &GpuBuffer,
    ) -> Result<(), String> {
        if ctx.gemm_flags() != self.flags_at_capture {
            return Err(format!(
                "m3 pooled prefill graph replay refused: GEMM-tier flags changed \
                 since capture (captured {:?}, now {:?})",
                self.flags_at_capture,
                ctx.gemm_flags()
            ));
        }
        if mamba_input.cached_ptr() != self.input_ptr
            || states.ssm.cached_ptr() != self.ssm_ptr
            || states.k.cached_ptr() != self.k_ptr
            || states.v.cached_ptr() != self.v_ptr
            || states.angle.cached_ptr() != self.angle_ptr
            || pooled_sum.cached_ptr() != self.pooled_ptr
        {
            return Err(
                "m3 pooled prefill graph replay refused: input/state/pooled buffers \
                 differ from the captured allocations"
                    .to_string(),
            );
        }
        self.graph
            .launch()
            .map_err(|e| format!("m3 pooled prefill graph launch: {e:?}"))
    }
}