mamba-rs 0.7.3

//! GPU Mamba inference engine (T=1 step with optional CUDA Graphs).
//!
//! Provides zero-copy persistent state on GPU, pre-allocated scratch buffers,
//! and optional CUDA Graph capture for minimal kernel launch overhead.
//!
//! All 12 existing CUDA kernels are reused — zero new kernel code needed.

use super::blas::{TypedPtr, gpu_gemm_forward_dispatch, gpu_gemm_typed_forward_raw};
use super::buffers::{DtypedBuf, GpuBuffer};
use super::context::{GemmMode, GemmRole, GpuCtx};
use super::device::GpuDevice;
use super::dtype::WeightDtype;
use super::forward::GpuMambaDims;
use super::gemm_bi_inference::prepare_inference_arch_rung;
use super::graph_capture::{
    capture_into_graph_with_gemm_plan, require_deterministic_gemm_graph_plan,
    with_validated_gemm_graph_launch,
};
use super::kernel_identity::{CapturedGemmGraphPlan, PreparedGemmCaptureManifest};
use super::launch::{grid_1d, grid_norm};
use super::weights::{
    GpuMambaMixedWeights, GpuMambaWeights, MambaLayerWeightsView, MambaWeightsView,
};
use crate::config::MambaConfig;
use crate::weights::MambaWeights;
use cudarc::driver::PushKernelArg;
use std::cell::Cell;
use std::sync::Arc;

#[cfg(test)]
mod model_gemm_manifest_tests {
    use super::super::blas::vendor_gemm_test::Guard;
    use super::super::context::BiGemmFamily;
    use super::super::graph_capture::model_gemm_guard_tests::{
        assert_inventory, assert_plan_mutations, configure,
    };
    use super::*;
    use crate::config::ScanMode;

    fn config() -> MambaConfig {
        MambaConfig {
            d_model: 32,
            n_layers: 2,
            d_state: 8,
            d_conv: 4,
            expand: 2,
            scan_mode: ScanMode::Sequential,
            rms_norm_eps: 1e-5,
        }
    }

    fn weights(cfg: &MambaConfig, input_dim: usize, identity: bool) -> MambaWeights {
        let mut weights = MambaWeights::init(cfg, input_dim, 0x9051);
        if identity {
            weights.input_proj_w.clear();
            weights.input_proj_b.clear();
        }
        weights
    }

    fn projections(
        cfg: &MambaConfig,
        batch: usize,
        input: Option<usize>,
    ) -> Vec<(usize, usize, usize)> {
        let mut expected = Vec::new();
        if let Some(input) = input {
            expected.push((batch, input, cfg.d_model));
        }
        let layer = [
            (batch, cfg.d_model, 2 * cfg.d_inner()),
            (batch, cfg.d_inner(), cfg.xdbl_dim()),
            (batch, cfg.dt_rank(), cfg.d_inner()),
            (batch, cfg.d_inner(), cfg.d_model),
        ];
        for _ in 0..cfg.n_layers {
            expected.extend(layer);
        }
        expected
    }

    fn bits(output: &[f32]) -> Vec<u32> {
        assert!(output.iter().all(|x| x.is_finite()));
        output.iter().map(|x| x.to_bits()).collect()
    }

    #[test]
    #[ignore = "needs a CUDA device"]
    fn m1_failed_steps_clear_permits_and_mixed_graphs_reject_wrong_path() {
        let device = GpuDevice::new(0).unwrap();
        let cfg = config();
        let weights = weights(&cfg, cfg.d_model, true);
        let input = vec![0.01; cfg.d_model];
        let mut output = vec![0.0; cfg.d_model];
        let mut f32 = GpuMambaInference::new(&device, &weights, cfg, cfg.d_model, 1).unwrap();
        configure(&f32.ctx, BiGemmFamily::Inference, true);
        let mut state = f32.alloc_state().unwrap();
        let mut scratch = f32.alloc_scratch().unwrap();
        assert!(
            unsafe { f32.capture_graph(&mut state, &mut scratch) }
                .unwrap_err()
                .contains("eager")
        );
        for gpu_only in [false, true] {
            f32.step(&input, &mut output, &mut state, &mut scratch)
                .unwrap();
            assert!(f32.eager_gemm_manifest.get().is_some());
            let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                if gpu_only {
                    f32.step_gpu_only(&[], &mut state, &mut scratch)
                } else {
                    f32.step(&[], &mut output, &mut state, &mut scratch)
                }
            }));
            assert!(failed.is_err());
            assert!(f32.eager_gemm_manifest.get().is_none());
            assert!(
                unsafe { f32.capture_graph(&mut state, &mut scratch) }
                    .unwrap_err()
                    .contains("eager")
            );
        }
        // A manifest prepared against different live buffers cannot install a
        // graph. The capture error must also consume that old permit.
        f32.step(&input, &mut output, &mut state, &mut scratch)
            .unwrap();
        let mut cold_scratch = f32.alloc_scratch().unwrap();
        assert!(unsafe { f32.capture_graph(&mut state, &mut cold_scratch) }.is_err());
        assert!(!f32.has_graph());
        assert!(f32.eager_gemm_manifest.get().is_none());
        drop(f32);

        for path in [MixedGraphPath::Legacy, MixedGraphPath::Native] {
            for gpu_only in [false, true] {
                for invalid_upload in [false, true] {
                    eprintln!(
                        "installed-opposite graph: attempted={path:?} gpu_only={gpu_only} invalid_upload={invalid_upload}"
                    );
                    let mut engine = GpuMambaInferenceMixed::new(
                        &device,
                        &weights,
                        cfg,
                        cfg.d_model,
                        1,
                        WeightDtype::Bf16,
                    )
                    .unwrap();
                    configure(&engine.engine.ctx, BiGemmFamily::Inference, true);
                    let mut state = engine.alloc_state().unwrap();
                    let mut legacy = engine.alloc_scratch().unwrap();
                    let mut native = engine.alloc_mixed_scratch().unwrap();
                    engine
                        .step(&input, &mut output, &mut state, &mut legacy)
                        .unwrap();
                    engine
                        .step_mixed_native(&input, &mut output, &mut state, &mut native)
                        .unwrap();
                    let installed_path = if path == MixedGraphPath::Legacy {
                        unsafe { engine.capture_graph_mixed_native(&mut state, &mut native) }
                            .unwrap();
                        assert!(engine.eager_legacy_gemm_manifest.get().is_some());
                        MixedGraphPath::Native
                    } else {
                        unsafe { engine.capture_graph(&mut state, &mut legacy) }.unwrap();
                        assert!(engine.eager_mixed_native_gemm_manifest.get().is_some());
                        MixedGraphPath::Legacy
                    };
                    let installed_routes = engine
                        .captured_gemm_plan
                        .as_ref()
                        .unwrap()
                        .routes()
                        .to_vec();
                    let installed_scratch = (
                        engine.captured_state_ptr,
                        engine.captured_scratch_ptr,
                        engine.captured_half_staging_ptr,
                        engine.captured_bi_upcast_ptrs,
                    );
                    let attempted_input = if invalid_upload {
                        &[][..]
                    } else {
                        input.as_slice()
                    };
                    let failed =
                        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                            match (path, gpu_only) {
                                (MixedGraphPath::Legacy, false) => engine.step(
                                    attempted_input,
                                    &mut output,
                                    &mut state,
                                    &mut legacy,
                                ),
                                (MixedGraphPath::Legacy, true) => {
                                    engine.step_gpu_only(attempted_input, &mut state, &mut legacy)
                                }
                                (MixedGraphPath::Native, false) => engine.step_mixed_native(
                                    attempted_input,
                                    &mut output,
                                    &mut state,
                                    &mut native,
                                ),
                                (MixedGraphPath::Native, true) => engine
                                    .step_gpu_only_mixed_native(
                                        attempted_input,
                                        &mut state,
                                        &mut native,
                                    ),
                            }
                        }));
                    if invalid_upload {
                        assert!(failed.is_err(), "invalid upload must fail before replay");
                    } else {
                        assert!(failed.unwrap().unwrap_err().contains("captured path"));
                    }
                    let recapture = if path == MixedGraphPath::Legacy {
                        assert!(
                            engine.eager_legacy_gemm_manifest.get().is_none(),
                            "failed legacy attempt retained its eager permit with native graph installed"
                        );
                        unsafe { engine.capture_graph(&mut state, &mut legacy) }
                    } else {
                        assert!(
                            engine.eager_mixed_native_gemm_manifest.get().is_none(),
                            "failed native attempt retained its eager permit with legacy graph installed"
                        );
                        unsafe { engine.capture_graph_mixed_native(&mut state, &mut native) }
                    };
                    assert!(recapture.unwrap_err().contains("eager"));
                    assert!(engine.has_graph());
                    assert_eq!(engine.captured_path, Some(installed_path));
                    assert_eq!(
                        engine.captured_gemm_plan.as_ref().unwrap().routes(),
                        installed_routes
                    );
                    assert_eq!(
                        (
                            engine.captured_state_ptr,
                            engine.captured_scratch_ptr,
                            engine.captured_half_staging_ptr,
                            engine.captured_bi_upcast_ptrs
                        ),
                        installed_scratch
                    );
                    if installed_path == MixedGraphPath::Native {
                        engine
                            .step_gpu_only_mixed_native(&input, &mut state, &mut native)
                            .unwrap();
                    } else {
                        engine
                            .step_gpu_only(&input, &mut state, &mut legacy)
                            .unwrap();
                    }
                    engine.engine.ctx.stream.synchronize().unwrap();
                }
            }
        }

        let mut engine =
            GpuMambaInferenceMixed::new(&device, &weights, cfg, cfg.d_model, 1, WeightDtype::Bf16)
                .unwrap();
        configure(&engine.engine.ctx, BiGemmFamily::Inference, true);
        let mut state = engine.alloc_state().unwrap();
        let mut legacy = engine.alloc_scratch().unwrap();
        let mut native = engine.alloc_mixed_scratch().unwrap();
        for path in [MixedGraphPath::Legacy, MixedGraphPath::Native] {
            for gpu_only in [false, true] {
                engine
                    .step(&input, &mut output, &mut state, &mut legacy)
                    .unwrap();
                engine
                    .step_mixed_native(&input, &mut output, &mut state, &mut native)
                    .unwrap();
                let failed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    match (path, gpu_only) {
                        (MixedGraphPath::Legacy, false) => {
                            engine.step(&[], &mut output, &mut state, &mut legacy)
                        }
                        (MixedGraphPath::Legacy, true) => {
                            engine.step_gpu_only(&[], &mut state, &mut legacy)
                        }
                        (MixedGraphPath::Native, false) => {
                            engine.step_mixed_native(&[], &mut output, &mut state, &mut native)
                        }
                        (MixedGraphPath::Native, true) => {
                            engine.step_gpu_only_mixed_native(&[], &mut state, &mut native)
                        }
                    }
                }));
                assert!(failed.is_err());
                if path == MixedGraphPath::Legacy {
                    assert!(engine.eager_legacy_gemm_manifest.get().is_none());
                    assert!(engine.eager_mixed_native_gemm_manifest.get().is_some());
                    assert!(
                        unsafe { engine.capture_graph(&mut state, &mut legacy) }
                            .unwrap_err()
                            .contains("eager")
                    );
                } else {
                    assert!(engine.eager_mixed_native_gemm_manifest.get().is_none());
                    assert!(engine.eager_legacy_gemm_manifest.get().is_some());
                    assert!(
                        unsafe { engine.capture_graph_mixed_native(&mut state, &mut native) }
                            .unwrap_err()
                            .contains("eager")
                    );
                }
            }
        }
        engine
            .step(&input, &mut output, &mut state, &mut legacy)
            .unwrap();
        engine
            .step_mixed_native(&input, &mut output, &mut state, &mut native)
            .unwrap();
        // Positive control: both paths were prepared successfully and no step
        // attempt intervenes between capturing one and recapturing the other.
        unsafe { engine.capture_graph_mixed_native(&mut state, &mut native) }.unwrap();
        // Recapture the other prepared path into the same slot; the association
        // and plan must change together, not retain the native plan.
        unsafe { engine.capture_graph(&mut state, &mut legacy) }.unwrap();
        assert_eq!(engine.captured_path, Some(MixedGraphPath::Legacy));
        assert!(
            engine
                .step_mixed_native(&input, &mut output, &mut state, &mut native)
                .unwrap_err()
                .contains("captured path")
        );
        assert!(
            engine
                .step_gpu_only_mixed_native(&input, &mut state, &mut native)
                .unwrap_err()
                .contains("captured path")
        );
        let ctx = &engine.engine.ctx;
        assert!(
            ctx.ensure_half_staging(usize::MAX)
                .unwrap_err()
                .contains("cannot grow")
        );
        let original = engine.captured_half_staging_ptr;
        engine.captured_half_staging_ptr ^= 16;
        assert!(
            engine
                .step_gpu_only(&input, &mut state, &mut legacy)
                .unwrap_err()
                .contains("staging scratch changed")
        );
        engine.captured_half_staging_ptr = original;
        engine
            .step_gpu_only(&input, &mut state, &mut legacy)
            .unwrap();
        let calls = Cell::new(0);
        engine.engine.ctx.poison_gemm_for_test();
        assert!(
            with_validated_gemm_graph_launch(
                &engine.engine.ctx,
                true,
                engine.captured_gemm_plan.as_ref(),
                "poison",
                || {
                    calls.set(calls.get() + 1);
                    Ok(())
                }
            )
            .unwrap_err()
            .contains("unusable")
        );
        assert_eq!(calls.get(), 0);
        drop(engine);
    }

    #[test]
    #[ignore = "needs a CUDA device"]
    fn m1_model_manifests_replay_all_paths_without_vendor_gemm() {
        let device = GpuDevice::new(0).unwrap();
        let cfg = config();
        let deny = Guard::new(true).unwrap();
        for (family, tc) in [
            (BiGemmFamily::Inference, true),
            (BiGemmFamily::Triad, false),
            (BiGemmFamily::Triad, true),
        ] {
            for batch in [1, 3] {
                eprintln!("M1 F32 {family:?} tc={tc} B{batch} nonidentity");
                let input = vec![0.01; batch * 24];
                let mut output = vec![0.0; batch * cfg.d_model];
                let mut engine =
                    GpuMambaInference::new(&device, &weights(&cfg, 24, false), cfg, 24, batch)
                        .unwrap();
                configure(&engine.ctx, family, tc);
                let mut state = engine.alloc_state().unwrap();
                let mut scratch = engine.alloc_scratch().unwrap();
                // First public step also exercises cold architecture preparation.
                engine
                    .step(&input, &mut output, &mut state, &mut scratch)
                    .unwrap();
                let trace = engine
                    .ctx
                    .record_eager_gemm_trace(|| engine.step_kernels(&mut state, &mut scratch))
                    .unwrap();
                state.reset(&engine.ctx.stream).unwrap();
                engine
                    .step_gpu_only(&input, &mut state, &mut scratch)
                    .unwrap();
                scratch
                    .temporal
                    .download(&engine.ctx.stream, &mut output)
                    .unwrap();
                let expected_bits = bits(&output);
                let manifest = engine.eager_gemm_manifest.get().unwrap();
                unsafe { engine.capture_graph(&mut state, &mut scratch) }.unwrap();
                assert!(engine.eager_gemm_manifest.get().is_none());
                assert_inventory(
                    &engine.ctx,
                    &trace,
                    manifest,
                    engine.captured_gemm_plan.as_ref().unwrap(),
                    &projections(&cfg, batch, Some(24)),
                );
                if family == BiGemmFamily::Inference && batch == 1 {
                    assert_plan_mutations(&engine.ctx, engine.captured_gemm_plan.as_ref().unwrap());
                }
                for gpu_only in [false, true] {
                    state.reset(&engine.ctx.stream).unwrap();
                    if gpu_only {
                        engine
                            .step_gpu_only(&input, &mut state, &mut scratch)
                            .unwrap();
                        scratch
                            .temporal
                            .download(&engine.ctx.stream, &mut output)
                            .unwrap();
                    } else {
                        engine
                            .step(&input, &mut output, &mut state, &mut scratch)
                            .unwrap();
                    }
                    assert_eq!(bits(&output), expected_bits);
                }
                drop(engine);

                for dtype in [WeightDtype::Bf16, WeightDtype::F16] {
                    for path in [MixedGraphPath::Legacy, MixedGraphPath::Native] {
                        eprintln!("M1 {dtype:?} {path:?} {family:?} tc={tc} B{batch}");
                        let native = path == MixedGraphPath::Native;
                        let input_dim = if native { cfg.d_model } else { 24 };
                        let input = vec![0.01; batch * input_dim];
                        let mut engine = GpuMambaInferenceMixed::new(
                            &device,
                            &weights(&cfg, input_dim, native),
                            cfg,
                            input_dim,
                            batch,
                            dtype,
                        )
                        .unwrap();
                        configure(&engine.engine.ctx, family, tc);
                        let mut state = engine.alloc_state().unwrap();
                        let mut legacy_scratch = engine.alloc_scratch().unwrap();
                        let mut native_scratch = engine.alloc_mixed_scratch().unwrap();
                        let ctx = &engine.engine.ctx;
                        let trace;
                        let manifest = if native {
                            engine
                                .step_mixed_native(
                                    &input,
                                    &mut output,
                                    &mut state,
                                    &mut native_scratch,
                                )
                                .unwrap();
                            trace = ctx
                                .record_eager_gemm_trace(|| {
                                    engine
                                        .step_kernels_mixed_native(&mut state, &mut native_scratch)
                                })
                                .unwrap();
                            state.reset(&ctx.stream).unwrap();
                            engine
                                .step_gpu_only_mixed_native(&input, &mut state, &mut native_scratch)
                                .unwrap();
                            native_scratch
                                .temporal
                                .download_f32(&ctx.stream, &mut output)
                                .unwrap();
                            engine.eager_mixed_native_gemm_manifest.get().unwrap()
                        } else {
                            engine
                                .step(&input, &mut output, &mut state, &mut legacy_scratch)
                                .unwrap();
                            trace = ctx
                                .record_eager_gemm_trace(|| {
                                    engine.step_kernels_mixed(&mut state, &mut legacy_scratch)
                                })
                                .unwrap();
                            state.reset(&ctx.stream).unwrap();
                            engine
                                .step_gpu_only(&input, &mut state, &mut legacy_scratch)
                                .unwrap();
                            legacy_scratch
                                .temporal
                                .download(&ctx.stream, &mut output)
                                .unwrap();
                            engine.eager_legacy_gemm_manifest.get().unwrap()
                        };
                        let expected_bits = bits(&output);
                        if native {
                            unsafe {
                                engine.capture_graph_mixed_native(&mut state, &mut native_scratch)
                            }
                            .unwrap();
                        } else {
                            unsafe { engine.capture_graph(&mut state, &mut legacy_scratch) }
                                .unwrap();
                        }
                        let ctx = &engine.engine.ctx;
                        assert_eq!(engine.captured_path, Some(path));
                        assert_inventory(
                            ctx,
                            &trace,
                            manifest,
                            engine.captured_gemm_plan.as_ref().unwrap(),
                            &projections(&cfg, batch, if native { None } else { Some(input_dim) }),
                        );
                        if family == BiGemmFamily::Triad && !tc {
                            assert!(trace.routes().iter().all(|r| r.symbol.contains("matvec")));
                        }
                        for gpu_only in [false, true] {
                            state.reset(&ctx.stream).unwrap();
                            match (native, gpu_only) {
                                (true, true) => {
                                    engine
                                        .step_gpu_only_mixed_native(
                                            &input,
                                            &mut state,
                                            &mut native_scratch,
                                        )
                                        .unwrap();
                                    native_scratch
                                        .temporal
                                        .download_f32(&ctx.stream, &mut output)
                                        .unwrap();
                                }
                                (true, false) => engine
                                    .step_mixed_native(
                                        &input,
                                        &mut output,
                                        &mut state,
                                        &mut native_scratch,
                                    )
                                    .unwrap(),
                                (false, true) => {
                                    engine
                                        .step_gpu_only(&input, &mut state, &mut legacy_scratch)
                                        .unwrap();
                                    legacy_scratch
                                        .temporal
                                        .download(&ctx.stream, &mut output)
                                        .unwrap();
                                }
                                (false, false) => engine
                                    .step(&input, &mut output, &mut state, &mut legacy_scratch)
                                    .unwrap(),
                            }
                            assert_eq!(bits(&output), expected_bits);
                        }
                        drop(engine);
                    }
                }
            }
        }
        assert_eq!(deny.calls(), 0);
    }
}

// ---------------------------------------------------------------------------
// GPU Inference State
// ---------------------------------------------------------------------------

/// Persistent GPU Mamba state for T=1 inference.
///
/// Lives in VRAM across steps. Reset on episode/sequence boundaries.
///
/// Layout:
/// - conv: `[n_layers][batch * d_inner * d_conv]` (layer-major, matches GPU training)
/// - ssm:  `[n_layers][batch * d_inner * d_state]`
///
/// Note: GPU conv uses `d_conv` slots (full window), while CPU inference uses
/// `d_conv - 1` (history only). The conv1d step kernels handle the
/// shift-register semantics with the full `d_conv` layout.
pub struct GpuInferenceState {
    pub conv: GpuBuffer,
    pub ssm: GpuBuffer,
    batch: usize,
    d_inner: usize,
    d_conv: usize,
    d_state: usize,
}

impl GpuInferenceState {
    /// Allocate zeroed inference state.
    pub fn zeros(
        stream: &Arc<cudarc::driver::CudaStream>,
        batch: usize,
        cfg: &MambaConfig,
    ) -> Result<Self, String> {
        let di = cfg.d_inner();
        let conv_len = cfg.n_layers * batch * di * cfg.d_conv;
        let ssm_len = cfg.n_layers * batch * di * cfg.d_state;
        Ok(Self {
            conv: GpuBuffer::zeros(stream, conv_len)?,
            ssm: GpuBuffer::zeros(stream, ssm_len)?,
            batch,
            d_inner: di,
            d_conv: cfg.d_conv,
            d_state: cfg.d_state,
        })
    }

    /// Reset all state to zero (episode boundary).
    pub fn reset(&mut self, stream: &Arc<cudarc::driver::CudaStream>) -> Result<(), String> {
        self.conv.zero(stream)?;
        self.ssm.zero(stream)
    }

    /// Per-layer conv state offset in floats.
    pub fn conv_offset(&self, layer: usize) -> usize {
        layer * self.batch * self.d_inner * self.d_conv
    }

    /// Per-layer SSM state offset in floats.
    pub fn ssm_offset(&self, layer: usize) -> usize {
        layer * self.batch * self.d_inner * self.d_state
    }

    /// Number of batch samples.
    pub fn batch(&self) -> usize {
        self.batch
    }
}

// ---------------------------------------------------------------------------
// GPU Inference Scratch
// ---------------------------------------------------------------------------

/// Pre-allocated scratch buffers for GPU inference (reused every step).
///
/// All GPU buffers are sized for a fixed batch size. Host staging buffers
/// handle H2D/D2H transfers outside CUDA Graph capture.
pub struct GpuInferenceScratch {
    // GPU scratch (reused per step)
    pub gpu_input: GpuBuffer,
    pub temporal: GpuBuffer,
    pub residual: GpuBuffer,
    pub proj: GpuBuffer,
    pub u: GpuBuffer,
    pub xdbl: GpuBuffer,
    pub dt_gather: GpuBuffer,
    pub delta: GpuBuffer,
    pub y: GpuBuffer,
    pub rms_buf: GpuBuffer,
}

impl GpuInferenceScratch {
    /// Allocate scratch buffers for the given batch size and config.
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        batch: usize,
        cfg: &MambaConfig,
        input_dim: usize,
    ) -> Result<Self, String> {
        let dm = cfg.d_model;
        let di = cfg.d_inner();
        let dt_rank = cfg.dt_rank();
        let xdbl_dim = cfg.xdbl_dim();
        Ok(Self {
            gpu_input: GpuBuffer::zeros(stream, batch * input_dim)?,
            temporal: GpuBuffer::zeros(stream, batch * dm)?,
            residual: GpuBuffer::zeros(stream, batch * dm)?,
            proj: GpuBuffer::zeros(stream, batch * 2 * di)?,
            u: GpuBuffer::zeros(stream, batch * di)?,
            xdbl: GpuBuffer::zeros(stream, batch * xdbl_dim)?,
            dt_gather: GpuBuffer::zeros(stream, batch * dt_rank)?,
            delta: GpuBuffer::zeros(stream, batch * di)?,
            y: GpuBuffer::zeros(stream, batch * di)?,
            rms_buf: GpuBuffer::zeros(stream, batch)?,
        })
    }
}

// ---------------------------------------------------------------------------
// GPU Inference Mixed Scratch (end-to-end bf16/f16 activations)
// ---------------------------------------------------------------------------

/// Scratch for mixed-precision inference with bf16/f16 activations.
///
/// Kept separate from `GpuInferenceScratch` (f32) so the RL training path
/// touches nothing. Dtype policy per-tensor:
///
/// - **Half dtype** (bf16/f16, matches weight dtype) — all linear-layer
///   I/O and activations: `temporal`, `proj`, `u`, `xdbl`, `dt_gather`,
///   `delta`, `y`. Storage
///   mantissa (bf16: 7-bit) is sufficient; compute happens in f32
///   (CUBLAS_COMPUTE_32F for GEMMs, upcast-inside-kernel for activation
///   kernels). Matches the reference state-spaces/mamba bf16 path.
/// - **f32** — `gpu_input` (CPU upload staging), `residual` (cross-layer
///   accumulator, HF default `residual_in_fp32=True`), `rms_buf`
///   (per-batch statistic). The residual stream is the only
///   non-negotiable f32: over 24+ layers, bf16 residual drifts
///   measurably (llama.cpp #9590).
pub struct GpuInferenceMixedScratch {
    pub gpu_input: GpuBuffer,
    pub temporal: DtypedBuf,
    pub residual: GpuBuffer,
    pub proj: DtypedBuf,
    pub u: DtypedBuf,
    pub xdbl: DtypedBuf,
    pub dt_gather: DtypedBuf,
    pub delta: DtypedBuf,
    pub y: DtypedBuf,
    pub rms_buf: GpuBuffer,
    pub dtype: WeightDtype,
}

impl GpuInferenceMixedScratch {
    /// Allocate mixed-precision scratch. `dtype` must match the weight dtype
    /// of the `GpuMambaInferenceMixed` engine that will use this scratch.
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        batch: usize,
        cfg: &MambaConfig,
        input_dim: usize,
        dtype: WeightDtype,
    ) -> Result<Self, String> {
        if matches!(dtype, WeightDtype::F32) {
            return Err("GpuInferenceMixedScratch requires bf16 or f16 dtype (use \
                 GpuInferenceScratch for f32)"
                .to_string());
        }
        let dm = cfg.d_model;
        let di = cfg.d_inner();
        let dt_rank = cfg.dt_rank();
        let xdbl_dim = cfg.xdbl_dim();
        Ok(Self {
            gpu_input: GpuBuffer::zeros(stream, batch * input_dim)?,
            temporal: DtypedBuf::zeros(stream, batch * dm, dtype)?,
            residual: GpuBuffer::zeros(stream, batch * dm)?,
            proj: DtypedBuf::zeros(stream, batch * 2 * di, dtype)?,
            u: DtypedBuf::zeros(stream, batch * di, dtype)?,
            xdbl: DtypedBuf::zeros(stream, batch * xdbl_dim, dtype)?,
            dt_gather: DtypedBuf::zeros(stream, batch * dt_rank, dtype)?,
            delta: DtypedBuf::zeros(stream, batch * di, dtype)?,
            y: DtypedBuf::zeros(stream, batch * di, dtype)?,
            rms_buf: GpuBuffer::zeros(stream, batch)?,
            dtype,
        })
    }
}

// ---------------------------------------------------------------------------
// GPU Inference Engine
// ---------------------------------------------------------------------------

/// GPU Mamba inference engine — owns kernels, weights, cuBLAS.
///
/// Lifecycle:
/// 1. `GpuMambaInference::new()` — compile kernels, upload weights
/// 2. Allocate state + scratch via `alloc_state()` / `alloc_scratch()`
/// 3. Call `step()` each timestep
/// 4. Optionally call `capture_graph()` for ~2-5x speedup
/// 5. Call `state.reset()` on episode boundaries
pub struct GpuMambaInference {
    pub(super) ctx: GpuCtx,
    pub(super) weights: GpuMambaWeights,
    pub(super) a_neg_all: GpuBuffer,
    pub(super) cfg: MambaConfig,
    pub(super) input_dim: usize,
    pub(super) batch: usize,
    /// When true (HF Mamba with no input_proj), skip input projection and copy
    /// `gpu_input` → `temporal` directly. Mirrors CPU `mamba_step_no_proj`.
    pub(super) identity_proj: bool,
    graph: Option<cudarc::driver::CudaGraph>,
    captured_gemm_route: Option<crate::mamba_ssm::gpu::context::GemmRoute>,
    captured_gemm_plan: Option<CapturedGemmGraphPlan>,
    eager_gemm_manifest: Cell<Option<PreparedGemmCaptureManifest>>,
    /// Raw pointers captured during graph capture for runtime validation.
    captured_state_ptr: u64,
    captured_scratch_ptr: u64,
}

impl Drop for GpuMambaInference {
    fn drop(&mut self) {
        let _ = self.ctx.stream.synchronize();
        drop(self.graph.take());
    }
}

impl GpuMambaInference {
    fn dims(&self, seq_len: usize) -> GpuMambaDims {
        GpuMambaDims {
            batch: self.batch,
            seq_len,
            n_layers: self.cfg.n_layers,
            d_model: self.cfg.d_model,
            d_inner: self.cfg.d_inner(),
            d_state: self.cfg.d_state,
            d_conv: self.cfg.d_conv,
            dt_rank: self.cfg.dt_rank(),
            xdbl_dim: self.cfg.xdbl_dim(),
            mamba_input_dim: self.input_dim,
            scan_mode: self.cfg.scan_mode,
            rms_norm_eps: self.cfg.rms_norm_eps,
        }
    }

    /// Create an f32 inference engine in the GEMM mode `MAMBA_RS_GEMM_MODE`
    /// names (`deterministic` when unset). A model context serves the Inference
    /// kernels in the deterministic mode; the two cuBLAS modes are explicit
    /// alternatives. Use [`Self::new_with_mode`] to pick the mode in code and
    /// [`Self::ctx`] to inspect the route that graph capture binds.
    ///
    /// # Errors
    ///
    /// Returns configuration, `MAMBA_RS_GEMM_MODE`, CUDA, upload, or allocation
    /// failures.
    pub fn new(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
    ) -> Result<Self, String> {
        Self::new_inner(
            device,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            None,
            WeightDtype::F32,
        )
    }

    /// Create an inference engine with an explicit GEMM execution mode.
    ///
    /// The weights remain f32; `mode` selects the deterministic kernels or a
    /// cuBLAS mode, and `MAMBA_RS_GEMM_MODE` is ignored. Invalid model
    /// configuration, kernel compilation, allocation, upload, or vendor setup
    /// returns an error.
    pub fn new_with_mode(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        mode: GemmMode,
    ) -> Result<Self, String> {
        Self::new_inner(
            device,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            Some(mode),
            WeightDtype::F32,
        )
    }

    /// `dtype` is the storage precision the caller asked for, `F32` or
    /// `Tf32`; both store f32 and differ only in the f32 numeric contract
    /// the context's GEMMs follow.
    pub(crate) fn new_inner(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        mode: Option<GemmMode>,
        dtype: WeightDtype,
    ) -> Result<Self, String> {
        cfg.validate()?;
        let state_cap = crate::mamba_ssm::gpu::kernels::state_capacity(cfg.d_state)?;
        let role = GemmRole::inference(dtype);
        let ctx = match mode {
            Some(mode) => GpuCtx::new_with_state_cap_mode_and_role(device, state_cap, mode, role)?,
            None => GpuCtx::new_from_env_with_state_cap_and_role(device, state_cap, role)?,
        };

        let weights = GpuMambaWeights::from_cpu(&ctx.stream, cpu_weights, &cfg)?;

        // Precompute a_neg = -exp(a_log) for all layers
        let di = cfg.d_inner();
        let ds = cfg.d_state;
        let total_aneg = cfg.n_layers * di * ds;
        let a_neg_all = GpuBuffer::zeros(&ctx.stream, total_aneg)?;
        for (layer_idx, lw) in weights.layers.iter().enumerate() {
            let offset = layer_idx * di * ds;
            let dst_ptr = a_neg_all.raw_ptr_at(&ctx.stream, offset);
            let src_ptr = lw.a_log.ptr();
            let n_i = (di * ds) as i32;
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.exp_negate);
            builder.arg(&dst_ptr);
            builder.arg(&src_ptr);
            builder.arg(&n_i);
            unsafe { builder.launch(grid_1d(di * ds)) }
                .map_err(|e| format!("exp_negate layer {layer_idx}: {e:?}"))?;
        }

        let identity_proj = cpu_weights.input_proj_w.is_empty();

        Ok(Self {
            ctx,
            weights,
            a_neg_all,
            cfg,
            input_dim,
            batch,
            identity_proj,
            graph: None,
            captured_gemm_route: None,
            captured_gemm_plan: None,
            eager_gemm_manifest: Cell::new(None),
            captured_state_ptr: 0,
            captured_scratch_ptr: 0,
        })
    }

    /// Capture CUDA Graph for the inference step.
    ///
    /// After capture, both step entries replay the fixed-buffer graph.
    ///
    /// Requires and consumes a successful eager step on these buffers. GEMM
    /// mode/family changes require new eager preparation and capture. A missing
    /// or mismatching GEMM manifest is an error. The inventory covers GEMMs;
    /// H2D/D2H transfers remain outside the recorded body and graph.
    ///
    /// # Safety
    ///
    /// `state`, `scratch`, and their views must remain unchanged until the
    /// graph is cleared and all replays complete. The engine context, stream,
    /// cuBLAS workspace, modules, functions, and weights must also stay fixed.
    pub unsafe fn capture_graph(
        &mut self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        let manifest = self.eager_gemm_manifest.take().ok_or_else(|| {
            "M1 f32 inference graph capture requires a successful eager step".to_string()
        })?;
        self.ctx.presize_bi_scratch()?;
        let snap_state = state.conv.cached_ptr();
        let snap_scratch = scratch.gpu_input.cached_ptr();
        let snap_gemm_route = self.ctx.gemm_route();
        let (graph, captured_gemm_plan) = unsafe {
            capture_into_graph_with_gemm_plan(&self.ctx, manifest.route_capacity, &manifest, || {
                self.step_kernels(state, scratch)
            })
        }?;
        require_deterministic_gemm_graph_plan(
            &self.ctx,
            self.has_gemm_work(),
            captured_gemm_plan.as_ref(),
            "M1 f32 inference graph capture",
        )?;
        self.graph = Some(graph);
        self.captured_gemm_route = Some(snap_gemm_route);
        self.captured_gemm_plan = captured_gemm_plan;
        self.captured_state_ptr = snap_state;
        self.captured_scratch_ptr = snap_scratch;
        self.ctx.note_graph_capture();
        Ok(())
    }

    fn launch_captured_graph(&self) -> Result<(), String> {
        let graph = self
            .graph
            .as_ref()
            .ok_or_else(|| "M1 f32 inference graph is not captured".to_string())?;
        with_validated_gemm_graph_launch(
            &self.ctx,
            self.has_gemm_work(),
            self.captured_gemm_plan.as_ref(),
            "M1 f32 inference graph replay",
            || {
                graph
                    .launch()
                    .map_err(|error| format!("graph launch: {error:?}"))
            },
        )
    }

    fn has_gemm_work(&self) -> bool {
        self.batch != 0 && (!self.identity_proj || self.cfg.n_layers != 0)
    }

    /// Whether a CUDA Graph has been captured.
    pub fn has_graph(&self) -> bool {
        self.graph.is_some()
    }

    /// Allocate zeroed inference state for this engine's batch size.
    pub fn alloc_state(&self) -> Result<GpuInferenceState, String> {
        GpuInferenceState::zeros(&self.ctx.stream, self.batch, &self.cfg)
    }

    /// Allocate scratch buffers for this engine.
    pub fn alloc_scratch(&self) -> Result<GpuInferenceScratch, String> {
        GpuInferenceScratch::new(&self.ctx.stream, self.batch, &self.cfg, self.input_dim)
    }

    /// Run one inference step: input → output.
    ///
    /// `input`: `[batch * input_dim]` on CPU.
    /// `output`: `[batch * d_model]` on CPU.
    /// State is mutated in-place (conv + SSM updated).
    ///
    /// # CUDA Graph mode
    /// When a graph is captured, all GPU pointers are baked in at capture time.
    /// The `state` and `scratch` parameters MUST be the same objects used during
    /// capture — passing different buffers silently operates on the original ones.
    /// Use [`GpuMambaBackbone`] which owns state/scratch and guarantees this.
    pub fn step(
        &self,
        input: &[f32],
        output: &mut [f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        // H2D: upload raw input (outside graph)
        self.eager_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.ctx.stream, input)?;

        // Run GPU kernel pipeline (graph replay or individual launches)
        if self.graph.is_some() {
            if self.captured_gemm_route != Some(self.ctx.gemm_route()) {
                return Err("inference graph replay: GEMM route changed since capture".into());
            }
            assert_eq!(
                state.conv.cached_ptr(),
                self.captured_state_ptr,
                "CUDA Graph replay requires the same state buffers used during capture"
            );
            assert_eq!(
                scratch.gpu_input.cached_ptr(),
                self.captured_scratch_ptr,
                "CUDA Graph replay requires the same scratch buffers used during capture"
            );
            self.launch_captured_graph()?;
        } else {
            if self.has_gemm_work() {
                prepare_inference_arch_rung(&self.ctx)?;
            }
            let manifest = self
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels(state, scratch))?;
            self.eager_gemm_manifest.set(Some(manifest));
        }

        // Sync: ensure all GPU work completes before D2H download.
        self.ctx
            .stream
            .synchronize()
            .map_err(|e| format!("sync: {e:?}"))?;
        scratch.temporal.download(&self.ctx.stream, output)?;
        Ok(())
    }

    /// Run backbone step without D2H download. Returns GPU temporal pointer.
    /// Use for chaining with lm_head GEMM on GPU without round-trip.
    pub fn step_gpu_only(
        &self,
        input: &[f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        self.eager_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.ctx.stream, input)?;
        if self.graph.is_some() {
            if self.captured_gemm_route != Some(self.ctx.gemm_route()) {
                return Err("inference graph replay: GEMM route changed since capture".into());
            }
            assert_eq!(state.conv.cached_ptr(), self.captured_state_ptr);
            assert_eq!(scratch.gpu_input.cached_ptr(), self.captured_scratch_ptr);
            self.launch_captured_graph()?;
        } else {
            if self.has_gemm_work() {
                prepare_inference_arch_rung(&self.ctx)?;
            }
            let manifest = self
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels(state, scratch))?;
            self.eager_gemm_manifest.set(Some(manifest));
        }
        Ok(())
    }

    /// Get the temporal output buffer (on GPU). Valid after `step_gpu_only`.
    pub fn temporal_buffer<'a>(&self, scratch: &'a GpuInferenceScratch) -> &'a GpuBuffer {
        &scratch.temporal
    }

    /// Launch the full T=1 forward pipeline on GPU.
    ///
    /// Pipeline per layer:
    /// ```text
    /// RmsNorm(+previous residual add) → in_proj GEMM → conv1d_step+silu →
    /// x_proj GEMM → gather_cols(dt) → dt_proj GEMM →
    /// ssm_step(+softplus+gather+gate) → out_proj GEMM
    /// ```
    fn step_kernels(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        // F32 path uses the typed-kernel dispatch via generic. This should be
        // identical to the legacy f32-only kernels because TypedKernel.get(F32)
        // returns the _f32 suffix variant which is the same PTX as legacy.
        self.step_kernels_generic(&self.weights, state, scratch)
    }

    /// Debug-only: f32 step stopping after `stop_after_layer` so the caller
    /// can download `scratch.temporal` for layer-by-layer parity comparison.
    #[doc(hidden)]
    pub fn step_kernels_f32_debug(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
        stop_after_layer: usize,
    ) -> Result<(), String> {
        self.step_kernels_generic_impl(&self.weights, state, scratch, Some(stop_after_layer))
    }

    /// Generic step pipeline — works with any weight view (f32 or bf16/f16).
    /// Bulk weights dispatch to SGEMM (f32) or cublasGemmEx (bf16/f16).
    /// Always-f32 weights read directly from their pointers.
    pub(super) fn step_kernels_generic<W: MambaWeightsView>(
        &self,
        weights: &W,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        self.step_kernels_generic_impl(weights, state, scratch, None)
    }

    fn step_kernels_generic_impl<W: MambaWeightsView>(
        &self,
        weights: &W,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
        stop_after_layer: Option<usize>,
    ) -> Result<(), String> {
        let b = self.batch;
        let cfg = &self.cfg;
        let dm = cfg.d_model;
        let di = cfg.d_inner();
        let ds = cfg.d_state;
        let dt_rank = cfg.dt_rank();
        let xdbl_dim = cfg.xdbl_dim();
        let d_conv = cfg.d_conv;
        let k = &self.ctx.kernels;

        // Input projection: [B, input_dim] → [B, d_model]
        // HF Mamba models have no input_proj (embedding is already d_model).
        // Identity branch mirrors CPU `mamba_step_no_proj`.
        if self.identity_proj {
            debug_assert_eq!(
                self.input_dim, dm,
                "identity_proj requires input_dim == d_model"
            );
            scratch
                .temporal
                .copy_from_raw(&scratch.gpu_input, &self.ctx.stream)?;
        } else {
            let (ipw_ptr, ipw_dtype) = weights.input_proj_w();
            gpu_gemm_forward_dispatch(
                &self.ctx,
                &mut scratch.temporal,
                &scratch.gpu_input,
                ipw_ptr,
                ipw_dtype,
                Some(weights.input_proj_b()),
                (b, self.input_dim, dm),
            )?;
        }

        let f32_sz = std::mem::size_of::<f32>() as u64;

        let layer_limit = stop_after_layer
            .map(|n| n.min(weights.n_layers()))
            .unwrap_or(weights.n_layers());

        for layer_idx in 0..layer_limit {
            let lw = weights.layer(layer_idx);
            let conv_ptr = state.conv.cached_ptr() + (state.conv_offset(layer_idx) as u64) * f32_sz;
            let ssm_ptr = state.ssm.cached_ptr() + (state.ssm_offset(layer_idx) as u64) * f32_sz;
            let aneg_ptr = self.a_neg_all.cached_ptr() + (layer_idx * di * ds) as u64 * f32_sz;

            // F1: RmsNorm. Layer 0 saves the input projection as the residual
            // and normalizes it. From the second layer on, the previous
            // layer's residual add rides in the same launch: the norm adds
            // the branch output (temporal) into the residual in place and
            // normalizes the sum, so neither the add kernel nor the residual
            // copy runs.
            {
                let b_i = b as i32;
                let dm_i = dm as i32;
                let eps: f32 = cfg.rms_norm_eps;
                let t_ptr = scratch.temporal.cached_ptr();
                let rms_ptr = scratch.rms_buf.cached_ptr();
                let res_ptr = scratch.residual.cached_ptr();
                let nw = lw.norm_weight();
                let mut bld = if layer_idx == 0 {
                    scratch
                        .residual
                        .copy_from_raw(&scratch.temporal, &self.ctx.stream)?;
                    let mut bld = self.ctx.stream.launch_builder(&k.rmsnorm_fwd);
                    bld.arg(&t_ptr); // output overwrites temporal
                    bld.arg(&rms_ptr);
                    bld.arg(&res_ptr); // input = saved residual
                    bld
                } else {
                    let mut bld = self
                        .ctx
                        .stream
                        .launch_builder(&k.rmsnorm_fwd_resadd_typed.f32);
                    bld.arg(&t_ptr); // output overwrites temporal
                    bld.arg(&rms_ptr);
                    bld.arg(&res_ptr); // residual += branch, then the norm input
                    bld.arg(&t_ptr); // branch = the previous layer's out_proj
                    bld
                };
                bld.arg(&nw);
                bld.arg(&b_i);
                bld.arg(&dm_i);
                bld.arg(&eps);
                unsafe { bld.launch(grid_norm(b, dm)) }
                    .map_err(|e| format!("rmsnorm_fwd L{layer_idx}: {e:?}"))?;
            }

            // F2: in_proj SGEMM [B, d_model] → [B, 2*d_inner]
            let (ipw, ipw_dt) = lw.in_proj_w();
            gpu_gemm_forward_dispatch(
                &self.ctx,
                &mut scratch.proj,
                &scratch.temporal,
                ipw,
                ipw_dt,
                None,
                (b, dm, 2 * di),
            )?;

            // F3+F4: conv1d_step with the SiLU fused into its store → u. It
            // reads the x half of the in_proj output straight from proj (row
            // stride 2 * d_inner), so the former split kernel is gone; the
            // gate half is read by the step kernel below.
            {
                let b_i = b as i32;
                let di_i = di as i32;
                let dc_i = d_conv as i32;
                let x_stride_i = (2 * di) as i32;
                let mut bld = self
                    .ctx
                    .stream
                    .launch_builder(&k.conv1d_step_fwd_silu_typed.f32);
                let u_ptr = scratch.u.cached_ptr();
                let proj_ptr = scratch.proj.cached_ptr();
                bld.arg(&u_ptr);
                bld.arg(&conv_ptr); // state mutated in-place
                bld.arg(&proj_ptr);
                bld.arg(&x_stride_i);
                let cw = lw.conv1d_weight();
                let cb = lw.conv1d_bias();
                bld.arg(&cw);
                bld.arg(&cb);
                bld.arg(&b_i);
                bld.arg(&di_i);
                bld.arg(&dc_i);
                unsafe { bld.launch(grid_1d(b * di)) }
                    .map_err(|e| format!("conv1d_step+silu L{layer_idx}: {e:?}"))?;
            }

            // F5: x_proj GEMM [B, d_inner] → [B, xdbl_dim]
            let (xpw, xpw_dt) = lw.x_proj_w();
            gpu_gemm_forward_dispatch(
                &self.ctx,
                &mut scratch.xdbl,
                &scratch.u,
                xpw,
                xpw_dt,
                None,
                (b, di, xdbl_dim),
            )?;

            // F6: gather dt from xdbl
            {
                let b_i = b as i32;
                let xdbl_i = xdbl_dim as i32;
                let dt_i = dt_rank as i32;
                let offset: i32 = 0;
                let mut bld = self.ctx.stream.launch_builder(&k.gather_cols);
                let dtg_ptr = scratch.dt_gather.cached_ptr();
                let xdbl_ptr = scratch.xdbl.cached_ptr();
                bld.arg(&dtg_ptr);
                bld.arg(&xdbl_ptr);
                bld.arg(&b_i);
                bld.arg(&xdbl_i);
                bld.arg(&dt_i);
                bld.arg(&offset);
                unsafe { bld.launch(grid_1d(b * dt_rank)) }
                    .map_err(|e| format!("gather_cols dt L{layer_idx}: {e:?}"))?;
            }

            // F7: dt_proj GEMM; the step kernel applies the softplus.
            let (dpw, dpw_dt) = lw.dt_proj_w();
            gpu_gemm_forward_dispatch(
                &self.ctx,
                &mut scratch.delta,
                &scratch.dt_gather,
                dpw,
                dpw_dt,
                Some(lw.dt_proj_b()),
                (b, dt_rank, di),
            )?;
            // F7-F10: the step kernel applies softplus to the dt_proj output,
            // reads B and C straight from xdbl, runs the recurrence, and
            // multiplies the gate (read from proj, SiLU recomputed) into y
            // before the store. Every folded value is spelled as the kernel
            // it replaces spelled it, so y keeps its bits.
            {
                let b_i = b as i32;
                let di_i = di as i32;
                let ds_i = ds as i32;
                let xdbl_stride_i = xdbl_dim as i32;
                let gate_stride_i = (2 * di) as i32;
                let b_off = dt_rank as i32;
                let c_off = (dt_rank + ds) as i32;
                // The step kernel keeps the state in registers sized at
                // compile time and silently returns without writing y
                // beyond that cap; T=1 decode has no parallel alternative
                // to route to, so refuse loudly.
                assert!(
                    ds <= k.state_cap,
                    "ssm_step_fwd_fused: d_state {ds} exceeds the compiled state capacity {}",
                    k.state_cap
                );
                let dp = lw.d_param();
                let mut bld = self
                    .ctx
                    .stream
                    .launch_builder(&k.ssm_step_fwd_fused_typed.f32);
                let y_ssm_ptr = scratch.y.cached_ptr();
                let delta_ssm_ptr = scratch.delta.cached_ptr();
                let u_ssm_ptr = scratch.u.cached_ptr();
                let xdbl_ssm_ptr = scratch.xdbl.cached_ptr();
                let proj_ptr = scratch.proj.cached_ptr();
                bld.arg(&ssm_ptr);
                bld.arg(&y_ssm_ptr);
                bld.arg(&delta_ssm_ptr);
                bld.arg(&u_ssm_ptr);
                bld.arg(&xdbl_ssm_ptr);
                bld.arg(&proj_ptr);
                bld.arg(&gate_stride_i);
                bld.arg(&aneg_ptr);
                bld.arg(&dp);
                bld.arg(&b_i);
                bld.arg(&di_i);
                bld.arg(&ds_i);
                bld.arg(&xdbl_stride_i);
                bld.arg(&b_off);
                bld.arg(&c_off);
                unsafe { bld.launch(grid_1d(b * di)) }
                    .map_err(|e| format!("ssm_step fused L{layer_idx}: {e:?}"))?;
            }

            // F11: out_proj GEMM [B, d_inner] → [B, d_model]
            let (opw, opw_dt) = lw.out_proj_w();
            gpu_gemm_forward_dispatch(
                &self.ctx,
                &mut scratch.temporal,
                &scratch.y,
                opw,
                opw_dt,
                None,
                (b, di, dm),
            )?;
        }

        // The last layer's residual add is still pending. A full step folds
        // it into norm_f: residual += temporal, then temporal <- norm(residual)
        // for the lm_head. A debug run that stopped early applies the add on
        // its own so the caller sees temporal as it stands after the
        // requested layer, without norm_f.
        if layer_limit > 0 && stop_after_layer.is_some() {
            let n = (b * dm) as i32;
            let mut bld = self.ctx.stream.launch_builder(&k.residual_add);
            let t_ptr = scratch.temporal.cached_ptr();
            let r_ptr = scratch.residual.cached_ptr();
            bld.arg(&t_ptr);
            bld.arg(&r_ptr);
            bld.arg(&t_ptr); // temporal = residual + temporal
            bld.arg(&n);
            unsafe { bld.launch(grid_1d(b * dm)) }
                .map_err(|e| format!("residual debug tail: {e:?}"))?;
        }
        if stop_after_layer.is_none() {
            let b_i = b as i32;
            let dm_i = dm as i32;
            let eps: f32 = cfg.rms_norm_eps;
            let mut bld = self
                .ctx
                .stream
                .launch_builder(&k.rmsnorm_fwd_resadd_typed.f32);
            let t_ptr = scratch.temporal.cached_ptr();
            let rms_ptr = scratch.rms_buf.cached_ptr();
            let res_ptr = scratch.residual.cached_ptr();
            bld.arg(&t_ptr);
            bld.arg(&rms_ptr);
            bld.arg(&res_ptr);
            bld.arg(&t_ptr); // branch = the last layer's out_proj
            let nfw = weights.norm_f_weight();
            bld.arg(&nfw);
            bld.arg(&b_i);
            bld.arg(&dm_i);
            bld.arg(&eps);
            unsafe { bld.launch(grid_norm(b, dm)) }.map_err(|e| format!("norm_f: {e:?}"))?;
        }

        Ok(())
    }

    /// Config reference.
    pub fn config(&self) -> &MambaConfig {
        &self.cfg
    }

    /// Batch size.
    pub fn batch(&self) -> usize {
        self.batch
    }

    /// Access the context that owns this engine's GEMM mode and family.
    ///
    /// Storage remains f32; inspect execution with [`GpuCtx::gemm_mode`] and
    /// [`GpuCtx::bi_gemm_family`] before graph capture.
    pub fn ctx(&self) -> &GpuCtx {
        &self.ctx
    }
}

// ---------------------------------------------------------------------------
// Mixed-precision inference engine (bf16/f16 weight storage, f32 compute).
// ---------------------------------------------------------------------------

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum MixedGraphPath {
    Legacy,
    Native,
}

/// GPU Mamba inference with mixed-precision weights (bf16 or f16).
///
/// Bulk linear weights use half storage; norms, biases and recurrence
/// parameters stay F32. The context selects deterministic or vendor GEMMs.
/// Legacy F32-activation and native half-activation paths have separate eager
/// permits; their shared graph slot is bound to exactly one captured path.
pub struct GpuMambaInferenceMixed {
    engine: GpuMambaInference, // owns ctx + (possibly unused) f32 weights
    mixed_weights: GpuMambaMixedWeights,
    a_neg_all: GpuBuffer,
    graph: Option<cudarc::driver::CudaGraph>,
    captured_gemm_route: Option<crate::mamba_ssm::gpu::context::GemmRoute>,
    captured_gemm_plan: Option<CapturedGemmGraphPlan>,
    captured_path: Option<MixedGraphPath>,
    eager_legacy_gemm_manifest: Cell<Option<PreparedGemmCaptureManifest>>,
    eager_mixed_native_gemm_manifest: Cell<Option<PreparedGemmCaptureManifest>>,
    captured_state_ptr: u64,
    captured_scratch_ptr: u64,
    captured_half_staging_ptr: u64,
    captured_bi_upcast_ptrs: [u64; 3],
}

impl Drop for GpuMambaInferenceMixed {
    fn drop(&mut self) {
        let _ = self.engine.ctx.stream.synchronize();
        drop(self.graph.take());
    }
}

impl GpuMambaInferenceMixed {
    fn ensure_graph_scratch(&self) -> Result<(), String> {
        self.engine.ctx.ensure_graph_scratch_ptrs(
            self.captured_half_staging_ptr,
            self.captured_bi_upcast_ptrs,
            "mixed inference graph replay",
        )
    }

    fn has_gemm_work(&self, path: MixedGraphPath) -> bool {
        self.engine.batch != 0
            && (self.engine.cfg.n_layers != 0
                || (path == MixedGraphPath::Legacy && !self.engine.identity_proj))
    }

    fn ensure_graph_path(&self, path: MixedGraphPath) -> Result<(), String> {
        if self.captured_path != Some(path) {
            return Err(
                "M1 mixed inference graph replay: captured path does not match entry".into(),
            );
        }
        Ok(())
    }

    fn launch_captured_graph(&self, path: MixedGraphPath) -> Result<(), String> {
        self.ensure_graph_path(path)?;
        let graph = self
            .graph
            .as_ref()
            .ok_or_else(|| "M1 mixed-native inference graph is not captured".to_string())?;
        with_validated_gemm_graph_launch(
            &self.engine.ctx,
            self.has_gemm_work(path),
            self.captured_gemm_plan.as_ref(),
            "M1 mixed inference graph replay",
            || {
                graph
                    .launch()
                    .map_err(|error| format!("graph launch mixed_native: {error:?}"))
            },
        )
    }

    /// Create a bf16/f16-storage engine in the GEMM mode `MAMBA_RS_GEMM_MODE`
    /// names (`deterministic` when unset).
    ///
    /// `bulk_dtype` controls storage, not the GEMM mode. Use [`Self::new_with_mode`]
    /// for an explicit mode and [`Self::ctx`] to inspect the route that graph
    /// capture binds. Configuration, dtype, upload, and allocation failures are
    /// returned.
    pub fn new(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        bulk_dtype: WeightDtype,
    ) -> Result<Self, String> {
        Self::new_inner(device, cpu_weights, cfg, input_dim, batch, bulk_dtype, None)
    }

    /// Create a mixed-storage inference engine with an explicit GEMM mode.
    ///
    /// `bulk_dtype` chooses bf16/f16 storage; `mode` separately controls GEMM
    /// execution, and `MAMBA_RS_GEMM_MODE` is ignored. Construction preserves the
    /// same validation, upload, and allocation errors as [`Self::new`].
    pub fn new_with_mode(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        bulk_dtype: WeightDtype,
        mode: GemmMode,
    ) -> Result<Self, String> {
        Self::new_inner(
            device,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            bulk_dtype,
            Some(mode),
        )
    }

    fn new_inner(
        device: &GpuDevice,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        bulk_dtype: WeightDtype,
        mode: Option<GemmMode>,
    ) -> Result<Self, String> {
        cfg.validate()?;
        // Create f32 engine first (builds ctx, kernels, a_neg_all via CPU upload path).
        // We'll then discard its `weights` flat buffer and replace with mixed arena.
        let engine = GpuMambaInference::new_inner(
            device,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            mode,
            WeightDtype::F32,
        )?;
        let mixed_weights =
            GpuMambaMixedWeights::from_cpu(&engine.ctx.stream, cpu_weights, &cfg, bulk_dtype)?;

        // Pre-size half_staging for the worst-case step-time GEMM operand.
        // Eliminates lazy-grow during graph capture (a re-allocation would
        // free the buffer the captured graph baked a pointer into → silent
        // dangling on replay). Native step path already bypasses staging,
        // but the legacy step_kernels_mixed path and compute_logits fallback
        // still go through it.
        engine
            .ctx
            .presize_half_staging_for_step(&cfg, batch, bulk_dtype)?;

        // Precompute a_neg into a separate arena (same as engine but from mixed weights'
        // f32 a_log — they match since a_log is f32 in both storages).
        let di = cfg.d_inner();
        let ds = cfg.d_state;
        let total_aneg = cfg.n_layers * di * ds;
        let a_neg_all = GpuBuffer::zeros(&engine.ctx.stream, total_aneg)?;
        for (layer_idx, lw) in mixed_weights.layers.iter().enumerate() {
            let offset = layer_idx * di * ds;
            let dst_ptr = a_neg_all.raw_ptr_at(&engine.ctx.stream, offset);
            let src_ptr = lw.a_log.ptr();
            let n_i = (di * ds) as i32;
            let mut builder = engine
                .ctx
                .stream
                .launch_builder(&engine.ctx.kernels.exp_negate);
            builder.arg(&dst_ptr);
            builder.arg(&src_ptr);
            builder.arg(&n_i);
            unsafe { builder.launch(grid_1d(di * ds)) }
                .map_err(|e| format!("exp_negate mixed L{layer_idx}: {e:?}"))?;
        }

        Ok(Self {
            engine,
            mixed_weights,
            a_neg_all,
            graph: None,
            captured_gemm_route: None,
            captured_gemm_plan: None,
            captured_path: None,
            eager_legacy_gemm_manifest: Cell::new(None),
            eager_mixed_native_gemm_manifest: Cell::new(None),
            captured_state_ptr: 0,
            captured_scratch_ptr: 0,
            captured_half_staging_ptr: 0,
            captured_bi_upcast_ptrs: [0; 3],
        })
    }

    pub fn step(
        &self,
        input: &[f32],
        output: &mut [f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        self.eager_legacy_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.engine.ctx.stream, input)?;
        if self.graph.is_some() {
            self.ensure_graph_path(MixedGraphPath::Legacy)?;
            if self.captured_gemm_route != Some(self.engine.ctx.gemm_route()) {
                return Err(
                    "mixed inference graph replay: GEMM route changed since capture".into(),
                );
            }
            self.ensure_graph_scratch()?;
            assert_eq!(state.conv.cached_ptr(), self.captured_state_ptr);
            assert_eq!(scratch.gpu_input.cached_ptr(), self.captured_scratch_ptr);
            self.launch_captured_graph(MixedGraphPath::Legacy)?;
        } else {
            if self.has_gemm_work(MixedGraphPath::Legacy) {
                prepare_inference_arch_rung(&self.engine.ctx)?;
            }
            let manifest = self
                .engine
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels_mixed(state, scratch))?;
            self.eager_legacy_gemm_manifest.set(Some(manifest));
        }
        self.engine
            .ctx
            .stream
            .synchronize()
            .map_err(|e| format!("sync: {e:?}"))?;
        scratch.temporal.download(&self.engine.ctx.stream, output)?;
        Ok(())
    }

    pub fn step_gpu_only(
        &self,
        input: &[f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        self.eager_legacy_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.engine.ctx.stream, input)?;
        if self.graph.is_some() {
            self.ensure_graph_path(MixedGraphPath::Legacy)?;
            if self.captured_gemm_route != Some(self.engine.ctx.gemm_route()) {
                return Err(
                    "mixed inference graph replay: GEMM route changed since capture".into(),
                );
            }
            self.ensure_graph_scratch()?;
            assert_eq!(state.conv.cached_ptr(), self.captured_state_ptr);
            assert_eq!(scratch.gpu_input.cached_ptr(), self.captured_scratch_ptr);
            self.launch_captured_graph(MixedGraphPath::Legacy)?;
        } else {
            if self.has_gemm_work(MixedGraphPath::Legacy) {
                prepare_inference_arch_rung(&self.engine.ctx)?;
            }
            let manifest = self
                .engine
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels_mixed(state, scratch))?;
            self.eager_legacy_gemm_manifest.set(Some(manifest));
        }
        Ok(())
    }

    fn step_kernels_mixed(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        // Legacy f32-activation path: upcasts activations to f32 around
        // every GEMM via the cast-staging buffer inside gpu_gemm_forward_dispatch.
        // Kept for backward compatibility; new callers should use the mixed
        // scratch + `step_kernels_mixed_native` for the end-to-end bf16 path.
        self.engine
            .step_kernels_generic(&self.mixed_weights, state, scratch)
    }

    /// End-to-end bf16/f16 activation pipeline.
    ///
    /// Unlike `step_kernels_mixed` (which cast-staged around each GEMM),
    /// this path keeps activations in the weight dtype throughout — GEMMs
    /// write half output directly, the rmsnorm/silu/softplus kernels are
    /// the typed variants (upcast compute, half storage), and only the
    /// residual stream stays f32 (HF `residual_in_fp32=True` semantics).
    ///
    /// Requires `identity_proj=true` (HF LLM use case). The RL path keeps
    /// using `step_kernels_generic` against `GpuInferenceScratch`.
    /// Debug-only: run the mixed-native step stopping after `stop_after_layer`
    /// has been processed (so the caller can download `scratch.residual` /
    /// `scratch.temporal` and compare layer-by-layer against the f32 path).
    /// `stop_after_layer == w.n_layers()` runs the full step (identical math
    /// to `step_kernels_mixed_native`).
    #[doc(hidden)]
    pub fn step_kernels_mixed_native_debug(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
        stop_after_layer: usize,
    ) -> Result<(), String> {
        self.step_kernels_mixed_native_impl(state, scratch, Some(stop_after_layer))
    }

    pub(super) fn step_kernels_mixed_native(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
    ) -> Result<(), String> {
        self.step_kernels_mixed_native_impl(state, scratch, None)
    }

    fn step_kernels_mixed_native_impl(
        &self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
        stop_after_layer: Option<usize>,
    ) -> Result<(), String> {
        let engine = &self.engine;
        assert!(
            engine.identity_proj,
            "step_kernels_mixed_native requires identity_proj=true (LLM path)"
        );
        assert_eq!(
            scratch.dtype, self.mixed_weights.bulk_dtype,
            "mixed scratch dtype must match mixed weights bulk_dtype"
        );
        let dt = scratch.dtype;
        let b = engine.batch;
        let cfg = &engine.cfg;
        let dm = cfg.d_model;
        let di = cfg.d_inner();
        let ds = cfg.d_state;
        let dt_rank = cfg.dt_rank();
        let xdbl_dim = cfg.xdbl_dim();
        let d_conv = cfg.d_conv;
        let k = &engine.ctx.kernels;
        let w = &self.mixed_weights;

        // Entry: residual_f32 <- gpu_input (identity_proj). Use copy_from_raw
        // (cuMemcpyDtoDAsync on raw ptrs) — CUDA Graph safe; copy_from creates
        // SyncOnDrop guards that invalidate the capture.
        scratch
            .residual
            .copy_from_raw(&scratch.gpu_input, &engine.ctx.stream)?;

        let f32_sz = std::mem::size_of::<f32>() as u64;

        let layer_limit = stop_after_layer
            .map(|n| n.min(w.n_layers()))
            .unwrap_or(w.n_layers());

        for layer_idx in 0..layer_limit {
            let lw = w.layer(layer_idx);
            let conv_ptr = state.conv.cached_ptr() + (state.conv_offset(layer_idx) as u64) * f32_sz;
            let ssm_ptr = state.ssm.cached_ptr() + (state.ssm_offset(layer_idx) as u64) * f32_sz;
            let aneg_ptr = self.a_neg_all.cached_ptr() + (layer_idx * di * ds) as u64 * f32_sz;

            // F1: rmsnorm f32_in → half_out (temporal_bf16 <- residual_f32 * norm_w).
            // From the second layer on, the previous layer's residual add
            // rides in the same launch: the norm adds the branch output
            // (temporal) into the f32 residual and normalizes the sum.
            {
                let b_i = b as i32;
                let dm_i = dm as i32;
                let eps: f32 = cfg.rms_norm_eps;
                let t_ptr = scratch.temporal.cached_ptr();
                let rms_ptr = scratch.rms_buf.cached_ptr();
                let res_ptr = scratch.residual.cached_ptr();
                let nw = lw.norm_weight();
                let mut bld = if layer_idx == 0 {
                    let mut bld = engine
                        .ctx
                        .stream
                        .launch_builder(k.rmsnorm_fwd_f32in_typed.get(dt));
                    bld.arg(&t_ptr);
                    bld.arg(&rms_ptr);
                    bld.arg(&res_ptr);
                    bld
                } else {
                    let mut bld = engine
                        .ctx
                        .stream
                        .launch_builder(k.rmsnorm_fwd_resadd_typed.get(dt));
                    bld.arg(&t_ptr);
                    bld.arg(&rms_ptr);
                    bld.arg(&res_ptr);
                    bld.arg(&t_ptr); // branch = the previous layer's out_proj
                    bld
                };
                bld.arg(&nw);
                bld.arg(&b_i);
                bld.arg(&dm_i);
                bld.arg(&eps);
                unsafe { bld.launch(grid_norm(b, dm)) }
                    .map_err(|e| format!("rmsnorm_f32in L{layer_idx}: {e:?}"))?;
            }

            // F2: in_proj GEMM — bf16 input, bf16 weights, bf16 output.
            let (ipw, ipw_dt) = lw.in_proj_w();
            gpu_gemm_typed_forward_raw(
                &engine.ctx,
                TypedPtr {
                    ptr: scratch.proj.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: scratch.temporal.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: ipw,
                    dtype: ipw_dt,
                },
                None,
                (b, dm, 2 * di),
            )?;

            // F3+F4: conv1d_step with the SiLU fused into its store. It reads
            // the x half of the in_proj output straight from proj (row stride
            // 2 * d_inner), so the former split kernel is gone; the gate half
            // is read by the step kernel below.
            {
                let b_i = b as i32;
                let di_i = di as i32;
                let dc_i = d_conv as i32;
                let x_stride_i = (2 * di) as i32;
                let mut bld = engine
                    .ctx
                    .stream
                    .launch_builder(k.conv1d_step_fwd_silu_typed.get(dt));
                let u_ptr = scratch.u.cached_ptr();
                let proj_ptr = scratch.proj.cached_ptr();
                bld.arg(&u_ptr);
                bld.arg(&conv_ptr);
                bld.arg(&proj_ptr);
                bld.arg(&x_stride_i);
                let cw = lw.conv1d_weight();
                let cb = lw.conv1d_bias();
                bld.arg(&cw);
                bld.arg(&cb);
                bld.arg(&b_i);
                bld.arg(&di_i);
                bld.arg(&dc_i);
                unsafe { bld.launch(grid_1d(b * di)) }
                    .map_err(|e| format!("conv1d_step+silu L{layer_idx}: {e:?}"))?;
            }

            // F5: x_proj GEMM (bf16 everywhere).
            let (xpw, xpw_dt) = lw.x_proj_w();
            gpu_gemm_typed_forward_raw(
                &engine.ctx,
                TypedPtr {
                    ptr: scratch.xdbl.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: scratch.u.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: xpw,
                    dtype: xpw_dt,
                },
                None,
                (b, di, xdbl_dim),
            )?;

            // F6: gather_cols typed — dt slice of xdbl into dt_gather (bf16).
            {
                let b_i = b as i32;
                let xdbl_i = xdbl_dim as i32;
                let dt_i = dt_rank as i32;
                let offset: i32 = 0;
                let mut bld = engine
                    .ctx
                    .stream
                    .launch_builder(k.gather_cols_typed.get(dt));
                let dtg_ptr = scratch.dt_gather.cached_ptr();
                let xdbl_ptr = scratch.xdbl.cached_ptr();
                bld.arg(&dtg_ptr);
                bld.arg(&xdbl_ptr);
                bld.arg(&b_i);
                bld.arg(&xdbl_i);
                bld.arg(&dt_i);
                bld.arg(&offset);
                unsafe { bld.launch(grid_1d(b * dt_rank)) }
                    .map_err(|e| format!("gather_cols dt L{layer_idx}: {e:?}"))?;
            }

            // F7: dt_proj GEMM (+ f32 bias) → delta bf16; the step kernel applies the softplus.
            let (dpw, dpw_dt) = lw.dt_proj_w();
            gpu_gemm_typed_forward_raw(
                &engine.ctx,
                TypedPtr {
                    ptr: scratch.delta.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: scratch.dt_gather.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: dpw,
                    dtype: dpw_dt,
                },
                Some(lw.dt_proj_b()),
                (b, dt_rank, di),
            )?;
            // F7-F10: the step kernel applies softplus to the dt_proj output,
            // reads B and C straight from xdbl, runs the recurrence, and
            // multiplies the gate (read from proj, SiLU recomputed) into y
            // before the store. Every folded value is spelled as the kernel
            // it replaces spelled it, so y keeps its bits.
            //
            // The step kernel keeps the state in registers sized at compile
            // time and silently returns without writing y beyond that cap;
            // T=1 decode has no parallel alternative to route to, so refuse
            // loudly.
            assert!(
                ds <= k.state_cap,
                "ssm_step_fwd_fused_typed: d_state {ds} exceeds the compiled \
                 state capacity {} (the fused kernel keeps the state in registers \
                 sized at compile time)",
                k.state_cap
            );
            {
                let b_i = b as i32;
                let di_i = di as i32;
                let ds_i = ds as i32;
                let xdbl_stride_i = xdbl_dim as i32;
                let gate_stride_i = (2 * di) as i32;
                let b_off = dt_rank as i32;
                let c_off = (dt_rank + ds) as i32;
                let dp = lw.d_param();
                let mut bld = engine
                    .ctx
                    .stream
                    .launch_builder(k.ssm_step_fwd_fused_typed.get(dt));
                let y_ssm_ptr = scratch.y.cached_ptr();
                let delta_ssm_ptr = scratch.delta.cached_ptr();
                let u_ssm_ptr = scratch.u.cached_ptr();
                let xdbl_ssm_ptr = scratch.xdbl.cached_ptr();
                let proj_ptr = scratch.proj.cached_ptr();
                bld.arg(&ssm_ptr);
                bld.arg(&y_ssm_ptr);
                bld.arg(&delta_ssm_ptr);
                bld.arg(&u_ssm_ptr);
                bld.arg(&xdbl_ssm_ptr);
                bld.arg(&proj_ptr);
                bld.arg(&gate_stride_i);
                bld.arg(&aneg_ptr);
                bld.arg(&dp);
                bld.arg(&b_i);
                bld.arg(&di_i);
                bld.arg(&ds_i);
                bld.arg(&xdbl_stride_i);
                bld.arg(&b_off);
                bld.arg(&c_off);
                unsafe { bld.launch(grid_1d(b * di)) }
                    .map_err(|e| format!("ssm_step fused L{layer_idx}: {e:?}"))?;
            }

            // F11: out_proj GEMM (bf16 y → bf16 temporal).
            let (opw, opw_dt) = lw.out_proj_w();
            gpu_gemm_typed_forward_raw(
                &engine.ctx,
                TypedPtr {
                    ptr: scratch.temporal.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: scratch.y.cached_ptr(),
                    dtype: dt,
                },
                TypedPtr {
                    ptr: opw,
                    dtype: opw_dt,
                },
                None,
                (b, di, dm),
            )?;
        }

        // The last layer's residual add is still pending. A full step folds
        // it into norm_f: residual_f32 += temporal, then temporal_bf16 <-
        // norm(residual) for the lm_head. A debug run that stopped early
        // applies the add on its own so the caller sees the residual as it
        // stands after the requested layer, without norm_f.
        if layer_limit > 0 && stop_after_layer.is_some() {
            let n = (b * dm) as i32;
            let mut bld = engine
                .ctx
                .stream
                .launch_builder(k.residual_add_f32_typed.get(dt));
            let r_ptr = scratch.residual.cached_ptr();
            let t_ptr = scratch.temporal.cached_ptr();
            bld.arg(&r_ptr);
            bld.arg(&r_ptr);
            bld.arg(&t_ptr);
            bld.arg(&n);
            unsafe { bld.launch(grid_1d(b * dm)) }
                .map_err(|e| format!("residual_add_f32 debug tail: {e:?}"))?;
        }
        if stop_after_layer.is_none() {
            let b_i = b as i32;
            let dm_i = dm as i32;
            let eps: f32 = cfg.rms_norm_eps;
            let mut bld = engine
                .ctx
                .stream
                .launch_builder(k.rmsnorm_fwd_resadd_typed.get(dt));
            let t_ptr = scratch.temporal.cached_ptr();
            let rms_ptr = scratch.rms_buf.cached_ptr();
            let res_ptr = scratch.residual.cached_ptr();
            bld.arg(&t_ptr);
            bld.arg(&rms_ptr);
            bld.arg(&res_ptr);
            bld.arg(&t_ptr); // branch = the last layer's out_proj
            let nfw = w.norm_f_weight();
            bld.arg(&nfw);
            bld.arg(&b_i);
            bld.arg(&dm_i);
            bld.arg(&eps);
            unsafe { bld.launch(grid_norm(b, dm)) }.map_err(|e| format!("norm_f_mixed: {e:?}"))?;
        }

        Ok(())
    }

    /// Run a mixed-native step with f32 input/output and internal bf16 activations.
    pub fn step_mixed_native(
        &self,
        input: &[f32],
        output: &mut [f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
    ) -> Result<(), String> {
        self.eager_mixed_native_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.engine.ctx.stream, input)?;
        if self.graph.is_some() {
            self.ensure_graph_path(MixedGraphPath::Native)?;
            if self.captured_gemm_route != Some(self.engine.ctx.gemm_route()) {
                return Err(
                    "mixed inference graph replay: GEMM route changed since capture".into(),
                );
            }
            self.ensure_graph_scratch()?;
            assert_eq!(state.conv.cached_ptr(), self.captured_state_ptr);
            assert_eq!(scratch.gpu_input.cached_ptr(), self.captured_scratch_ptr);
            self.launch_captured_graph(MixedGraphPath::Native)?;
        } else {
            if self.has_gemm_work(MixedGraphPath::Native) {
                prepare_inference_arch_rung(&self.engine.ctx)?;
            }
            let manifest = self
                .engine
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels_mixed_native(state, scratch))?;
            self.eager_mixed_native_gemm_manifest.set(Some(manifest));
        }
        self.engine
            .ctx
            .stream
            .synchronize()
            .map_err(|e| format!("sync: {e:?}"))?;
        scratch
            .temporal
            .download_f32(&self.engine.ctx.stream, output)?;
        Ok(())
    }

    /// GPU-only step (no output download) for chained mixed-native inference.
    pub fn step_gpu_only_mixed_native(
        &self,
        input: &[f32],
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
    ) -> Result<(), String> {
        self.eager_mixed_native_gemm_manifest.set(None);
        scratch.gpu_input.upload(&self.engine.ctx.stream, input)?;
        if self.graph.is_some() {
            self.ensure_graph_path(MixedGraphPath::Native)?;
            if self.captured_gemm_route != Some(self.engine.ctx.gemm_route()) {
                return Err(
                    "mixed inference graph replay: GEMM route changed since capture".into(),
                );
            }
            self.ensure_graph_scratch()?;
            assert_eq!(state.conv.cached_ptr(), self.captured_state_ptr);
            assert_eq!(scratch.gpu_input.cached_ptr(), self.captured_scratch_ptr);
            self.launch_captured_graph(MixedGraphPath::Native)?;
            Ok(())
        } else {
            if self.has_gemm_work(MixedGraphPath::Native) {
                prepare_inference_arch_rung(&self.engine.ctx)?;
            }
            let manifest = self
                .engine
                .ctx
                .record_eager_gemm_manifest(|| self.step_kernels_mixed_native(state, scratch))?;
            self.eager_mixed_native_gemm_manifest.set(Some(manifest));
            Ok(())
        }
    }

    /// Allocate a `GpuInferenceMixedScratch` sized for this engine.
    pub fn alloc_mixed_scratch(&self) -> Result<GpuInferenceMixedScratch, String> {
        GpuInferenceMixedScratch::new(
            &self.engine.ctx.stream,
            self.engine.batch,
            &self.engine.cfg,
            self.engine.input_dim,
            self.mixed_weights.bulk_dtype,
        )
    }

    /// Capture a CUDA Graph for the mixed-native pipeline.
    /// Requires and consumes this path's successful eager manifest. Mode or
    /// family drift and mismatching GEMM inventories are rejected. Transfers
    /// remain outside the body; the manifest inventories GEMMs only.
    ///
    /// # Safety
    ///
    /// `state`, `scratch`, and their views must remain unchanged until the
    /// graph is cleared and all replays complete. The engine context, stream,
    /// cuBLAS workspace, modules, functions, and weights must also stay fixed.
    pub unsafe fn capture_graph_mixed_native(
        &mut self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceMixedScratch,
    ) -> Result<(), String> {
        let manifest = self
            .eager_mixed_native_gemm_manifest
            .take()
            .ok_or_else(|| {
                "M1 mixed-native inference graph capture requires a successful eager step"
                    .to_string()
            })?;
        self.engine.ctx.presize_bi_scratch()?;
        self.engine
            .ctx
            .presize_mixed_graph_scratch_m1(&self.engine.dims(1), self.mixed_weights.bulk_dtype)?;
        let snap_state = state.conv.cached_ptr();
        let snap_scratch = scratch.gpu_input.cached_ptr();
        let snap_half_staging = self.engine.ctx.half_staging_ptr();
        let snap_bi_upcast = self.engine.ctx.bi_upcast_scratch_ptrs();
        let snap_gemm_route = self.engine.ctx.gemm_route();
        self.engine.ctx.freeze_graph_scratch();
        let (graph, captured_gemm_plan) = unsafe {
            capture_into_graph_with_gemm_plan(
                &self.engine.ctx,
                manifest.route_capacity,
                &manifest,
                || self.step_kernels_mixed_native(state, scratch),
            )
        }?;
        require_deterministic_gemm_graph_plan(
            &self.engine.ctx,
            self.has_gemm_work(MixedGraphPath::Native),
            captured_gemm_plan.as_ref(),
            "M1 mixed-native inference graph capture",
        )?;
        self.graph = Some(graph);
        self.captured_gemm_route = Some(snap_gemm_route);
        self.captured_gemm_plan = captured_gemm_plan;
        self.captured_path = Some(MixedGraphPath::Native);
        self.captured_state_ptr = snap_state;
        self.captured_scratch_ptr = snap_scratch;
        self.captured_half_staging_ptr = snap_half_staging;
        self.captured_bi_upcast_ptrs = snap_bi_upcast;
        self.engine.ctx.note_graph_capture();
        Ok(())
    }

    /// Capture the legacy F32-activation mixed pipeline after a successful
    /// eager legacy step on these fixed buffers. Consumes that path's GEMM-only
    /// manifest and rejects missing or changed work or GEMM policy. Native
    /// entries cannot replay this graph. H2D/D2H remain outside the graph.
    ///
    /// # Safety
    ///
    /// `state`, `scratch`, and their views must remain unchanged until the
    /// graph is cleared and all replays complete. The engine context, stream,
    /// cuBLAS workspace, modules, functions, and weights must also stay fixed.
    pub unsafe fn capture_graph(
        &mut self,
        state: &mut GpuInferenceState,
        scratch: &mut GpuInferenceScratch,
    ) -> Result<(), String> {
        let manifest = self.eager_legacy_gemm_manifest.take().ok_or_else(|| {
            "M1 legacy mixed inference graph capture requires a successful eager step".to_string()
        })?;
        self.engine.ctx.presize_bi_scratch()?;
        self.engine
            .ctx
            .presize_mixed_graph_scratch_m1(&self.engine.dims(1), self.mixed_weights.bulk_dtype)?;
        let snap_state = state.conv.cached_ptr();
        let snap_scratch = scratch.gpu_input.cached_ptr();
        let snap_half_staging = self.engine.ctx.half_staging_ptr();
        let snap_bi_upcast = self.engine.ctx.bi_upcast_scratch_ptrs();
        let snap_gemm_route = self.engine.ctx.gemm_route();
        self.engine.ctx.freeze_graph_scratch();
        let (graph, captured_gemm_plan) = unsafe {
            capture_into_graph_with_gemm_plan(
                &self.engine.ctx,
                manifest.route_capacity,
                &manifest,
                || self.step_kernels_mixed(state, scratch),
            )
        }?;
        require_deterministic_gemm_graph_plan(
            &self.engine.ctx,
            self.has_gemm_work(MixedGraphPath::Legacy),
            captured_gemm_plan.as_ref(),
            "M1 legacy mixed inference graph capture",
        )?;
        self.graph = Some(graph);
        self.captured_gemm_route = Some(snap_gemm_route);
        self.captured_gemm_plan = captured_gemm_plan;
        self.captured_path = Some(MixedGraphPath::Legacy);
        self.captured_state_ptr = snap_state;
        self.captured_scratch_ptr = snap_scratch;
        self.captured_half_staging_ptr = snap_half_staging;
        self.captured_bi_upcast_ptrs = snap_bi_upcast;
        self.engine.ctx.note_graph_capture();
        Ok(())
    }

    pub fn alloc_state(&self) -> Result<GpuInferenceState, String> {
        self.engine.alloc_state()
    }

    pub fn alloc_scratch(&self) -> Result<GpuInferenceScratch, String> {
        self.engine.alloc_scratch()
    }

    pub fn config(&self) -> &MambaConfig {
        &self.engine.cfg
    }

    pub fn batch(&self) -> usize {
        self.engine.batch
    }

    /// Access the retained engine context.
    ///
    /// Storage is reported by [`Self::bulk_dtype`]; inspect execution with
    /// [`GpuCtx::gemm_mode`] and [`GpuCtx::bi_gemm_family`] before graph capture.
    pub fn ctx(&self) -> &GpuCtx {
        &self.engine.ctx
    }

    pub fn stream(&self) -> &Arc<cudarc::driver::CudaStream> {
        &self.engine.ctx.stream
    }

    pub fn bulk_dtype(&self) -> WeightDtype {
        self.mixed_weights.bulk_dtype
    }

    pub fn has_graph(&self) -> bool {
        self.graph.is_some()
    }

    /// Access the underlying f32 inference engine (used by unified backbone
    /// wrapper to read cfg/batch/input_dim).
    pub fn engine_ref(&self) -> &GpuMambaInference {
        &self.engine
    }

    /// Access the mixed-precision weights (for parallel prefill path).
    pub fn weights_mixed_ref(&self) -> &GpuMambaMixedWeights {
        &self.mixed_weights
    }

    /// Access the precomputed a_neg arena (shared between engine paths).
    pub fn a_neg_all_ref(&self) -> &GpuBuffer {
        &self.a_neg_all
    }
}

// ---------------------------------------------------------------------------
// High-level GPU Backbone
// ---------------------------------------------------------------------------

/// High-level GPU Mamba backbone — owns engine + state + scratch.
///
/// Simple API for inference: `step()`, `reset()`, `capture_graph()`.
///
/// ```rust,no_run
/// use mamba_rs::MambaConfig;
/// use mamba_rs::gpu::inference::GpuMambaBackbone;
///
/// let cfg = MambaConfig::default();
/// let weights = mamba_rs::MambaWeights::init(&cfg, 128, 42);
/// let mut bb = GpuMambaBackbone::new(0, &weights, cfg, 128, 1).unwrap();
/// bb.capture_graph().unwrap(); // optional ~2-5x speedup
///
/// let input = vec![0.1f32; 128];
/// let mut output = vec![0.0f32; 128];
/// bb.step(&input, &mut output).unwrap();
/// bb.reset().unwrap();
/// ```
/// Internal engine variant. Box the larger variant to keep enum size small
/// (Mixed includes f32 engine + mixed weights arena + a_neg_all).
enum BackboneEngine {
    F32(Box<GpuMambaInference>),
    Mixed(Box<GpuMambaInferenceMixed>),
}

enum BackboneScratch {
    F32(GpuInferenceScratch),
    Mixed(GpuInferenceMixedScratch),
}

impl BackboneScratch {
    /// Raw device pointer of the final temporal output (post norm_f).
    /// f32 for F32 path, bf16/f16 for Mixed native path.
    fn temporal_ptr(&self) -> cudarc::driver::sys::CUdeviceptr {
        match self {
            BackboneScratch::F32(s) => s.temporal.cached_ptr(),
            BackboneScratch::Mixed(s) => s.temporal.cached_ptr(),
        }
    }

    /// Dtype of the final temporal output.
    fn temporal_dtype(&self) -> WeightDtype {
        match self {
            BackboneScratch::F32(_) => WeightDtype::F32,
            BackboneScratch::Mixed(s) => s.dtype,
        }
    }
}

/// High-level GPU Mamba backbone — unified API over f32 / bf16 / f16 storage.
///
/// Weights are uploaded to GPU in the requested dtype. Compute is always f32
/// (CUBLAS_COMPUTE_32F for GEMMs, f32 internally in kernels via upcast).
///
/// Activation storage tracks the weight dtype end-to-end on the Mixed path
/// (bf16/f16): the scratch carries half activations through the layer and
/// GEMMs write half output directly, no per-GEMM cast staging. The residual
/// stream and SSM recurrent state stay f32 for numerical stability (matches
/// HF `residual_in_fp32=True`). The F32 path keeps f32 scratch throughout.
///
/// ```rust,no_run
/// use mamba_rs::MambaConfig;
/// use mamba_rs::gpu::inference::GpuMambaBackbone;
/// use mamba_rs::mamba_ssm::gpu::dtype::WeightDtype;
///
/// let cfg = MambaConfig::default();
/// let weights = mamba_rs::MambaWeights::init(&cfg, 128, 42);
///
/// // f32 (default)
/// let mut bb = GpuMambaBackbone::new(0, &weights, cfg, 128, 1).unwrap();
///
/// // bf16 (half VRAM)
/// let mut bb_bf16 = GpuMambaBackbone::new_with_dtype(
///     0, &weights, cfg, 128, 1, WeightDtype::Bf16
/// ).unwrap();
///
/// bb.capture_graph().unwrap();
/// let mut output = vec![0.0f32; 128];
/// bb.step(&vec![0.1f32; 128], &mut output).unwrap();
/// ```
pub struct GpuMambaBackbone {
    engine: BackboneEngine,
    state: GpuInferenceState,
    scratch: BackboneScratch,
}

impl GpuMambaBackbone {
    /// Create an f32 GPU backbone in the GEMM mode `MAMBA_RS_GEMM_MODE` names
    /// (`deterministic` when unset).
    ///
    /// Storage remains f32; inspect the resolved route with [`Self::ctx`].
    /// Construction and later graph capture can return validation, CUDA, upload,
    /// allocation, or route errors.
    pub fn new(
        gpu_ordinal: usize,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
    ) -> Result<Self, String> {
        Self::new_with_dtype(
            gpu_ordinal,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            WeightDtype::F32,
        )
    }

    /// Create an f32 GPU backbone with an explicit GEMM execution mode.
    ///
    /// The mode is independent of f32 weight storage; `MAMBA_RS_GEMM_MODE` is
    /// ignored. Model validation, CUDA setup, upload, and allocation failures are
    /// returned. Captured graphs still require an unchanged complete GEMM route.
    pub fn new_with_mode(
        gpu_ordinal: usize,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        mode: GemmMode,
    ) -> Result<Self, String> {
        Self::new_with_dtype_and_mode(
            gpu_ordinal,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            WeightDtype::F32,
            mode,
        )
    }

    /// Create a GPU backbone with an explicit storage precision in the GEMM mode
    /// `MAMBA_RS_GEMM_MODE` names (`deterministic` when unset).
    ///
    /// `dtype` selects f32, tf32, bf16, or f16 storage independently of the mode;
    /// `Tf32` stores f32 and lets the deterministic GEMMs take the TF32 kernels
    /// where one is measured. Use [`Self::new_with_dtype_and_mode`] to pick the
    /// mode in code and [`Self::ctx`] to inspect the route that graph capture
    /// binds. Configuration, dtype, CUDA, upload, or allocation can also fail.
    pub fn new_with_dtype(
        gpu_ordinal: usize,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        dtype: WeightDtype,
    ) -> Result<Self, String> {
        Self::new_with_dtype_inner(gpu_ordinal, cpu_weights, cfg, input_dim, batch, dtype, None)
    }

    /// Create a GPU backbone with an explicit storage precision and GEMM mode.
    ///
    /// `dtype` controls weight/activation storage (`Tf32` stores f32 with
    /// deterministic TF32 products); `mode` controls GEMM execution and does not
    /// change that storage choice. `MAMBA_RS_GEMM_MODE` is ignored. Invalid
    /// configuration, dtype-specific setup, CUDA, upload, or allocation failures
    /// are returned; graph replay requires the construction route.
    pub fn new_with_dtype_and_mode(
        gpu_ordinal: usize,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        dtype: WeightDtype,
        mode: GemmMode,
    ) -> Result<Self, String> {
        Self::new_with_dtype_inner(
            gpu_ordinal,
            cpu_weights,
            cfg,
            input_dim,
            batch,
            dtype,
            Some(mode),
        )
    }

    fn new_with_dtype_inner(
        gpu_ordinal: usize,
        cpu_weights: &MambaWeights,
        cfg: MambaConfig,
        input_dim: usize,
        batch: usize,
        dtype: WeightDtype,
        mode: Option<GemmMode>,
    ) -> Result<Self, String> {
        let device = GpuDevice::new(gpu_ordinal)?;
        let (engine, state, scratch) = match dtype {
            WeightDtype::F32 | WeightDtype::Tf32 => {
                let e = GpuMambaInference::new_inner(
                    &device,
                    cpu_weights,
                    cfg,
                    input_dim,
                    batch,
                    mode,
                    dtype,
                )?;
                let s = e.alloc_state()?;
                let sc = BackboneScratch::F32(e.alloc_scratch()?);
                (BackboneEngine::F32(Box::new(e)), s, sc)
            }
            WeightDtype::Bf16 | WeightDtype::F16 => {
                let e = match mode {
                    Some(mode) => GpuMambaInferenceMixed::new_with_mode(
                        &device,
                        cpu_weights,
                        cfg,
                        input_dim,
                        batch,
                        dtype,
                        mode,
                    )?,
                    None => GpuMambaInferenceMixed::new(
                        &device,
                        cpu_weights,
                        cfg,
                        input_dim,
                        batch,
                        dtype,
                    )?,
                };
                let s = e.alloc_state()?;
                let sc = BackboneScratch::Mixed(e.alloc_mixed_scratch()?);
                (BackboneEngine::Mixed(Box::new(e)), s, sc)
            }
        };
        Ok(Self {
            engine,
            state,
            scratch,
        })
    }

    /// Storage dtype for this backbone's weights.
    pub fn dtype(&self) -> WeightDtype {
        match &self.engine {
            BackboneEngine::F32(e) => e.ctx.f32_storage_dtype(),
            BackboneEngine::Mixed(e) => e.bulk_dtype(),
        }
    }

    /// Run one inference step. `input`: `[batch * input_dim]`, `output`: `[batch * d_model]`.
    pub fn step(&mut self, input: &[f32], output: &mut [f32]) -> Result<(), String> {
        match (&self.engine, &mut self.scratch) {
            (BackboneEngine::F32(e), BackboneScratch::F32(sc)) => {
                e.step(input, output, &mut self.state, sc)
            }
            (BackboneEngine::Mixed(e), BackboneScratch::Mixed(sc)) => {
                e.step_mixed_native(input, output, &mut self.state, sc)
            }
            _ => Err("engine/scratch dtype mismatch (internal invariant)".to_string()),
        }
    }

    /// Reset recurrent state (episode/sequence boundary).
    pub fn reset(&mut self) -> Result<(), String> {
        let stream = match &self.engine {
            BackboneEngine::F32(e) => e.ctx.stream.clone(),
            BackboneEngine::Mixed(e) => e.stream().clone(),
        };
        self.state.reset(&stream)
    }

    /// Capture CUDA Graph for faster inference. Call after at least one warmup step.
    pub fn capture_graph(&mut self) -> Result<(), String> {
        let (input_dim, batch, d_model) = match &self.engine {
            BackboneEngine::F32(e) => (e.input_dim, e.batch, e.cfg.d_model),
            BackboneEngine::Mixed(e) => (
                e.engine_ref().input_dim,
                e.engine_ref().batch,
                e.engine_ref().cfg.d_model,
            ),
        };
        let input = vec![0.0f32; batch * input_dim];
        let mut output = vec![0.0f32; batch * d_model];
        self.step(&input, &mut output)?;
        self.reset()?;
        match (&mut self.engine, &mut self.scratch) {
            (BackboneEngine::F32(e), BackboneScratch::F32(sc)) => {
                // This owner drops its graph before the owned state and scratch.
                unsafe { e.capture_graph(&mut self.state, sc) }
            }
            (BackboneEngine::Mixed(e), BackboneScratch::Mixed(sc)) => {
                // This owner drops its graph before the owned state and scratch.
                unsafe { e.capture_graph_mixed_native(&mut self.state, sc) }
            }
            _ => Err("engine/scratch dtype mismatch".to_string()),
        }
    }

    pub fn config(&self) -> &MambaConfig {
        match &self.engine {
            BackboneEngine::F32(e) => e.config(),
            BackboneEngine::Mixed(e) => e.config(),
        }
    }

    pub fn batch(&self) -> usize {
        match &self.engine {
            BackboneEngine::F32(e) => e.batch(),
            BackboneEngine::Mixed(e) => e.batch(),
        }
    }

    pub fn has_graph(&self) -> bool {
        match &self.engine {
            BackboneEngine::F32(e) => e.has_graph(),
            BackboneEngine::Mixed(e) => e.has_graph(),
        }
    }

    /// Access the backbone's GPU execution context.
    ///
    /// Storage is reported separately by [`Self::dtype`]; inspect execution
    /// with [`GpuCtx::gemm_mode`] and [`GpuCtx::bi_gemm_family`]. A captured
    /// graph retains this complete route.
    pub fn ctx(&self) -> &GpuCtx {
        match &self.engine {
            BackboneEngine::F32(e) => &e.ctx,
            BackboneEngine::Mixed(e) => e.ctx(),
        }
    }

    /// Access the GPU stream.
    pub fn stream(&self) -> &std::sync::Arc<cudarc::driver::CudaStream> {
        match &self.engine {
            BackboneEngine::F32(e) => &e.ctx.stream,
            BackboneEngine::Mixed(e) => e.stream(),
        }
    }

    /// Run step, keep output on GPU (for chaining with lm_head).
    pub fn step_gpu_only(&mut self, input: &[f32]) -> Result<(), String> {
        match (&self.engine, &mut self.scratch) {
            (BackboneEngine::F32(e), BackboneScratch::F32(sc)) => {
                e.step_gpu_only(input, &mut self.state, sc)
            }
            (BackboneEngine::Mixed(e), BackboneScratch::Mixed(sc)) => {
                e.step_gpu_only_mixed_native(input, &mut self.state, sc)
            }
            _ => Err("engine/scratch dtype mismatch".to_string()),
        }
    }

    /// Debug-only: drive a single step from the current state, but stop after
    /// processing `layer_limit` layers. Downloads the post-layer residual
    /// (always f32 — for the F32 backbone this is `scratch.temporal` because
    /// the residual is copied into temporal at the start of every layer; for
    /// the Mixed backbone this is `scratch.residual` directly).
    #[doc(hidden)]
    pub fn debug_step_partial(
        &mut self,
        input: &[f32],
        layer_limit: usize,
        out: &mut [f32],
    ) -> Result<(), String> {
        match (&self.engine, &mut self.scratch) {
            (BackboneEngine::F32(e), BackboneScratch::F32(sc)) => {
                sc.gpu_input.upload(&e.ctx.stream, input)?;
                e.step_kernels_f32_debug(&mut self.state, sc, layer_limit)?;
                e.ctx
                    .stream
                    .synchronize()
                    .map_err(|err| format!("sync: {err:?}"))?;
                sc.temporal.download(&e.ctx.stream, out)
            }
            (BackboneEngine::Mixed(e), BackboneScratch::Mixed(sc)) => {
                sc.gpu_input.upload(e.stream(), input)?;
                e.step_kernels_mixed_native_debug(&mut self.state, sc, layer_limit)?;
                e.stream()
                    .synchronize()
                    .map_err(|err| format!("sync: {err:?}"))?;
                sc.residual.download(e.stream(), out)
            }
            _ => Err("engine/scratch dtype mismatch".to_string()),
        }
    }

    /// GPU temporal buffer pointer (valid after `step_gpu_only`).
    pub fn temporal_ptr(&self) -> cudarc::driver::sys::CUdeviceptr {
        self.scratch.temporal_ptr()
    }

    /// Dtype of the temporal buffer — f32 for the F32 engine, bf16/f16
    /// for the Mixed engine's native path.
    pub fn temporal_dtype(&self) -> WeightDtype {
        self.scratch.temporal_dtype()
    }

    /// Download temporal from GPU to CPU (always f32 output; Mixed path
    /// upcasts half → f32 on the fly).
    pub fn download_temporal(&self, output: &mut [f32]) -> Result<(), String> {
        self.stream()
            .synchronize()
            .map_err(|e| format!("sync: {e:?}"))?;
        match &self.scratch {
            BackboneScratch::F32(s) => s.temporal.download(self.stream(), output),
            BackboneScratch::Mixed(s) => s.temporal.download_f32(self.stream(), output),
        }
    }

    /// Run parallel prefill over T tokens using burnin kernels.
    ///
    /// `ip_out_flat`: `[B * T * d_model]` of pre-embedded prompt tokens (f32).
    /// After this call, the backbone's recurrent state holds position T, and
    /// the temporal buffer (accessible via `temporal_ptr()`) contains the
    /// last-timestep hidden state (ready for lm_head). Dtype matches
    /// `temporal_dtype()`.
    ///
    /// - F32 engine: uses `gpu_forward_inference_prefill` (f32 path).
    /// - Mixed engine: uses `gpu_forward_inference_prefill_mixed` — bf16/f16
    ///   activations end-to-end with f32 residual stream. Caller must pass a
    ///   `GpuMambaTargetMixedScratch` via the parallel `prefill_sequence_mixed`
    ///   entry point (f32 `GpuMambaTargetScratch` here is the F32-only path).
    pub fn prefill_sequence(
        &mut self,
        ip_out_flat: &GpuBuffer,
        prefill_scratch: &mut super::backward::GpuMambaTargetScratch,
    ) -> Result<(), String> {
        use super::prefill::{PrefillInputs, gpu_forward_inference_prefill};
        match (&self.engine, &mut self.scratch) {
            (BackboneEngine::F32(e), BackboneScratch::F32(sc)) => gpu_forward_inference_prefill(
                &e.ctx,
                &mut sc.temporal,
                PrefillInputs {
                    ip_out_flat,
                    weights: &e.weights,
                    a_neg_all: &e.a_neg_all,
                },
                &mut self.state,
                prefill_scratch,
            ),
            (BackboneEngine::Mixed(_), _) => {
                Err("mixed backbone: use prefill_sequence_mixed with a \
                 GpuMambaTargetMixedScratch (native bf16/f16 prefill)"
                    .to_string())
            }
            _ => Err("engine/scratch dtype mismatch".to_string()),
        }
    }

    /// Mixed-native prefill. Requires a Mixed backbone and matching mixed scratch.
    pub fn prefill_sequence_mixed(
        &mut self,
        ip_out_flat: &GpuBuffer,
        prefill_scratch: &mut super::backward::GpuMambaTargetMixedScratch,
    ) -> Result<(), String> {
        use super::prefill::{PrefillInputs, gpu_forward_inference_prefill_mixed};
        match (&self.engine, &mut self.scratch) {
            (BackboneEngine::Mixed(e), BackboneScratch::Mixed(sc)) => {
                gpu_forward_inference_prefill_mixed(
                    e.ctx(),
                    &sc.temporal,
                    PrefillInputs {
                        ip_out_flat,
                        weights: e.weights_mixed_ref(),
                        a_neg_all: e.a_neg_all_ref(),
                    },
                    &mut self.state,
                    prefill_scratch,
                )
            }
            _ => Err("prefill_sequence_mixed requires Mixed backbone + Mixed scratch".to_string()),
        }
    }

    /// Allocate a mixed-native prefill scratch matching this backbone.
    pub fn alloc_prefill_mixed_scratch(
        &self,
        seq_len: usize,
    ) -> Result<super::backward::GpuMambaTargetMixedScratch, String> {
        let dtype = match &self.engine {
            BackboneEngine::Mixed(e) => e.bulk_dtype(),
            BackboneEngine::F32(_) => {
                return Err("alloc_prefill_mixed_scratch: backbone is F32".to_string());
            }
        };
        let cfg = self.config();
        let dims = super::forward::GpuMambaDims {
            batch: self.batch(),
            seq_len,
            n_layers: cfg.n_layers,
            d_model: cfg.d_model,
            d_inner: cfg.d_inner(),
            d_state: cfg.d_state,
            d_conv: cfg.d_conv,
            dt_rank: cfg.dt_rank(),
            xdbl_dim: cfg.xdbl_dim(),
            mamba_input_dim: cfg.d_model,
            scan_mode: cfg.scan_mode,
            rms_norm_eps: cfg.rms_norm_eps,
        };
        super::backward::GpuMambaTargetMixedScratch::new(self.stream(), &dims, dtype)
    }

    /// Build a target scratch allocated for this backbone + seq_len.
    pub fn alloc_prefill_scratch(
        &self,
        seq_len: usize,
    ) -> Result<super::backward::GpuMambaTargetScratch, String> {
        let cfg = self.config();
        let dims = super::forward::GpuMambaDims {
            batch: self.batch(),
            seq_len,
            n_layers: cfg.n_layers,
            d_model: cfg.d_model,
            d_inner: cfg.d_inner(),
            d_state: cfg.d_state,
            d_conv: cfg.d_conv,
            dt_rank: cfg.dt_rank(),
            xdbl_dim: cfg.xdbl_dim(),
            mamba_input_dim: cfg.d_model, // HF LLM path: no input_proj, input_dim == d_model
            scan_mode: cfg.scan_mode,
            rms_norm_eps: cfg.rms_norm_eps,
        };
        super::backward::GpuMambaTargetScratch::new(self.stream(), &dims)
    }
}