mamba-rs 0.7.3

//! CUDA-Graph-captured training step for Mamba SSM mixed precision.
//!
//! Captures **forward + backward + AdamW + master→compute sync** as a
//! single CUDA Graph. On replay, the entire training step launches with
//! a single `cuGraphLaunch` instead of ~150 individual kernel launches —
//! cuts launch overhead from ~50–100 µs per step to ~5 µs.
//!
//! ## Scope
//! - **bf16 directly**; **f16 via the trainer's `capture_graph_f16`**,
//!   which extends this body with the in-graph loss-scaler kernels
//!   (`check_inf_nan` + `scale_grads_skip`) so AdamW can run
//!   unconditionally and the CPU reads the overflow flag only AFTER
//!   replay. The structs in this module stay bf16-only by assert.
//! - **Per (batch, seq_len) shape**: a captured graph bakes in tensor
//!   sizes via the kernel launch configs. A new shape requires a new
//!   capture; the [`GpuMambaTrainingStepGraph`] holder records its dims
//!   and can be invalidated explicitly.
//!
//! ## What's inside the graph
//!   1. `grads.flat.zero()` — async memset on stream
//!   2. `gpu_forward_mamba_backbone_mixed`
//!   3. `gpu_backward_mamba_backbone_mixed`
//!   4. `step_m1_capturable` (per-tensor AdamW, bias factors via device buf)
//!   5. `sync_master_to_compute` (per-tensor cast f32 → bf16 / D2D copy)
//!
//! ## What stays outside the graph
//!   - H2D upload of `mamba_input` (per-step input data)
//!   - H2D upload of `d_temporal` (per-step loss gradient)
//!   - H2D upload of `(bc1, bc2)` into [`AdamWBiasFactors`] — caller
//!     calls `bias.write(stream, bc1, bc2, lr)` BEFORE each replay, after
//!     `adam.advance()` to bump the step counter.
//!   - State management (zeroing or carrying recurrent state)
//!
//! ## Pointer-stability invariant
//! Captured pointers are stored at capture time. On replay, the holder
//! asserts that every input buffer's `cached_ptr()` still matches what
//! was captured. Any reallocation panics: mutating buffers between
//! capture and replay silently corrupts outputs, as a race on the 130m
//! model once showed.
//!
//! ## Reference
//! Same pattern as `GpuInferenceEngine::capture_graph` (inference.rs:307),
//! extended to backward + optimizer. Mirrors PyTorch's
//! `torch.cuda.CUDAGraph` + `capturable=True` optimizer mode (PyTorch
//! 2.5, `torch/cuda/graphs.py` and `_multi_tensor_adamw`).

use cudarc::driver::{CudaGraph, PushKernelArg};
use std::rc::Rc;

use crate::mamba_ssm::gpu::adamw::{AdamWBiasFactors, AdamWMultiPlan, GpuAdamW, step_multi};
use crate::mamba_ssm::gpu::backward::gpu_backward_mamba_backbone;
use crate::mamba_ssm::gpu::backward_mixed::gpu_backward_mamba_backbone_mixed;
use crate::mamba_ssm::gpu::buffers::GpuBuffer;
use crate::mamba_ssm::gpu::context::GpuCtx;
use crate::mamba_ssm::gpu::dtype::WeightDtype;
use crate::mamba_ssm::gpu::forward::{
    GpuMambaBackboneActs, GpuMambaScratch, GpuRecurrentState, gpu_forward_mamba_backbone,
};
use crate::mamba_ssm::gpu::forward_mixed::{
    GpuMambaBackboneMixedActs, GpuMambaMixedTrainScratch, gpu_forward_mamba_backbone_mixed,
};
use crate::mamba_ssm::gpu::graph_capture::{
    capture_into_graph_with_gemm_plan, require_deterministic_gemm_graph_plan,
    with_validated_gemm_graph_launch,
};
use crate::mamba_ssm::gpu::kernel_identity::{CapturedGemmGraphPlan, PreparedGemmCaptureManifest};
use crate::mamba_ssm::gpu::launch::grid_1d;
use crate::mamba_ssm::gpu::weights::{
    GpuMambaGrads, GpuMambaTrainLayerWeights, GpuMambaTrainWeights,
};
use crate::mamba_ssm::gpu::weights_mixed_train::GpuMambaTrainMixedWeights;

/// Capture-side variant of the trainer's `recompute_a_neg_all` helper.
/// Launches `exp_negate` per layer from `master_layers[l].a_log` into
/// BOTH `a_neg_all[l*per_layer..]` (backward) and
/// `state_a_neg_all[l*per_layer..]` (forward). These launches must be
/// inside the captured body so replay picks up the freshly-updated
/// `a_log` from AdamW. Without them, every replay runs the SSM with the
/// A-matrix from graph-capture time, effectively freezing the decay.
fn recompute_a_neg_captured(
    ctx: &GpuCtx,
    master_layers: &[GpuMambaTrainLayerWeights],
    a_neg_all: &GpuBuffer,
    state_a_neg_all: &GpuBuffer,
    d_inner: usize,
    d_state: usize,
) -> Result<(), String> {
    let per_layer = d_inner * d_state;
    if per_layer == 0 {
        return Ok(());
    }
    let n_i32 = per_layer as i32;
    for (li, mw) in master_layers.iter().enumerate() {
        let src = mw.a_log.cached_ptr();
        // Two-destination refresh — twin of the eager path.
        let dst_a = a_neg_all.inner_at(li * per_layer);
        let dst_s = state_a_neg_all.inner_at(li * per_layer);
        let mut b1 = ctx.stream.launch_builder(&ctx.kernels.exp_negate2);
        b1.arg(&dst_a);
        b1.arg(&dst_s);
        b1.arg(&src);
        b1.arg(&n_i32);
        unsafe { b1.launch(grid_1d(per_layer)) }
            .map_err(|e| format!("exp_negate2 captured a_neg mirrors L{li}: {e:?}"))?;
    }
    Ok(())
}

/// Buffers and parameter state borrowed into the captured M1 bf16
/// training-step body (forward + backward + AdamW + master→compute sync
/// + a_neg recompute).
pub struct MambaMixedCapture<'a> {
    pub train_w: &'a mut GpuMambaTrainMixedWeights,
    pub adam: &'a GpuAdamW,
    pub bias: &'a AdamWBiasFactors,
    /// Fused multi-tensor AdamW chunk table (built at construction).
    pub multi_plan: &'a AdamWMultiPlan,
    pub grads: &'a mut GpuMambaGrads,
    pub acts: &'a mut GpuMambaBackboneMixedActs,
    pub scratch: &'a mut GpuMambaMixedTrainScratch,
    /// Backward-side `a_neg_all` — a SEPARATE buffer from `state.a_neg_all`.
    pub a_neg_all: &'a GpuBuffer,
    pub mamba_input: &'a GpuBuffer,
    pub d_temporal: &'a mut GpuBuffer,
    pub state: &'a mut GpuRecurrentState,
}

/// Live buffers checked against the captured pointers on every M1 bf16
/// graph replay.
#[derive(Clone, Copy)]
pub struct MambaMixedReplay<'a> {
    pub train_w: &'a GpuMambaTrainMixedWeights,
    pub adam: &'a GpuAdamW,
    pub bias: &'a AdamWBiasFactors,
    pub grads: &'a GpuMambaGrads,
    /// Backward-side `a_neg_all` — a SEPARATE buffer from `state.a_neg_all`.
    pub a_neg_all: &'a GpuBuffer,
    pub mamba_input: &'a GpuBuffer,
    pub d_temporal: &'a GpuBuffer,
    pub state: &'a GpuRecurrentState,
}

/// CUDA-Graph holder for a single Mamba SSM training step (bf16).
///
/// The raw driver graph stays private so replay cannot bypass the captured
/// context, route, and pointer checks.
///
/// ```compile_fail
/// # use mamba_rs::mamba_ssm::gpu::training_graph::GpuMambaTrainingStepGraph;
/// # fn raw_launch(graph: &GpuMambaTrainingStepGraph) {
/// let _ = &graph.graph;
/// # }
/// ```
pub struct GpuMambaTrainingStepGraph {
    graph: CudaGraph,
    ctx_resources: Rc<crate::mamba_ssm::gpu::context::GpuCtxResources>,
    pub batch: usize,
    pub seq_len: usize,
    pub dtype: WeightDtype,

    // Captured-pointer assertions. If any of these change between
    // capture and replay, the graph is silently writing to the original
    // address — panic immediately.
    captured_input_ptr: u64,
    captured_d_temporal_ptr: u64,
    captured_grads_flat_ptr: u64,
    captured_adam_m_ptr: u64,
    captured_adam_v_ptr: u64,
    captured_bias_factors_ptr: u64,
    // ALL three GpuRecurrentState fields are accessed by forward (conv_states
    // and a_neg_all are read directly from `state.*`, ssm_states is the
    // recurrent SSM working buffer). Each pointer is baked into the graph;
    // any reallocation between capture and replay corrupts silently.
    captured_state_ssm_states_ptr: u64,
    captured_state_conv_states_ptr: u64,
    captured_state_a_neg_all_ptr: u64,
    // The backward-side `a_neg_all` is a SEPARATE parameter (not the
    // `state.a_neg_all` field above — they happen to share a name but are
    // distinct buffers in the public API).
    captured_a_neg_all_ptr: u64,
    // Weight-stability proxy: snapshot BOTH the first-allocated and the
    // last-allocated master tensor, plus their compute slices. If the
    // weight set is rebuilt (e.g. checkpoint reload), at least one of the
    // four is virtually guaranteed to land at a different address (the
    // allocator can't reuse all original slots simultaneously).
    captured_master_input_proj_w_ptr: u64,
    captured_master_norm_f_ptr: u64,
    captured_compute_input_proj_w_ptr: u64,
    captured_compute_norm_f_ptr: u64,
    // Defensive lazy-grow guard: today the typed mixed forward path uses
    // `gpu_gemm_typed_forward_raw` which doesn't call `ensure_half_staging`,
    // so this guard is currently inert. Kept (with cheap presize) so that
    // any future refactor that routes a captured kernel through
    // `gpu_gemm_forward_dispatch` doesn't silently bake a freed pointer
    // (CUDA_ERROR_ILLEGAL_ADDRESS).
    captured_half_staging_ptr: u64,
    // Same guard for the batch-invariant typed-GEMM upcast scratch triple:
    // every bf16 bi GEMM without a native typed bucket (all Big/split-K
    // training shapes) reads/writes these buffers inside the captured body.
    captured_bi_upcast_ptrs: [u64; 3],
    // Complete GEMM route captured with the graph.
    captured_gemm_flags: crate::mamba_ssm::gpu::context::GemmRoute,
    has_gemm_work: bool,
    captured_gemm_plan: Option<CapturedGemmGraphPlan>,
    captured_ctx_token: u64,
    captured_stream_token: usize,
}

impl GpuMambaTrainingStepGraph {
    /// Capture the training step into a CUDA Graph.
    ///
    /// Caller owns all the input/output buffers. The graph records the
    /// kernel launches; pointers are baked in. After capture, [`Self::replay`]
    /// re-runs the entire step with a single `cuGraphLaunch`.
    ///
    /// ## Required ordering before this call
    ///   1. Allocate `acts`, `scratch`, `grads`, `adam` (m/v), `bias`,
    ///      `state`, `mamba_input`, `d_temporal`, all weight tensors.
    ///   2. Run a warmup forward+backward+adamw+sync ONCE without graph
    ///      capture — this primes cuBLAS workspace selection and any
    ///      lazy CUDA module loading.
    ///   3. Then call this function. The pre-capture stream sync inside
    ///      ensures all warmup + allocations are complete (the "130m
    ///      race-fix invariant" — see `GpuMambaBackboneMixedActs::new`).
    ///
    /// `bias` must already hold the bias-correction factors for step **1**
    /// (or whatever step number you'll start replaying at). The captured
    /// AdamW reads from `bias`'s device pointer; CPU-side `adam.advance()`
    /// + `bias.write()` between replays drives the per-step values.
    /// # Safety
    ///
    /// The context, stream, cuBLAS handle, module functions, multi-plan, and
    /// every allocation or view used by the captured body must remain unchanged
    /// until this holder is destroyed and every replay has completed. Replay
    /// pointer checks diagnose drift but do not extend any CUDA lifetime.
    pub unsafe fn capture(
        ctx: &GpuCtx,
        cfg: &crate::config::MambaConfig,
        cap: MambaMixedCapture<'_>,
        batch: usize,
        seq_len: usize,
        manifest: &PreparedGemmCaptureManifest,
    ) -> Result<Self, String> {
        let MambaMixedCapture {
            train_w,
            adam,
            bias,
            multi_plan,
            grads,
            acts,
            scratch,
            a_neg_all,
            mamba_input,
            d_temporal,
            state,
        } = cap;
        // Result-returning capture must not panic on a caller mistake -
        // surface the contract as an error.
        if !matches!(train_w.dtype, WeightDtype::Bf16) {
            return Err(format!(
                "bf16-only graph capture: got {:?}; f16 needs the in-graph \
                 overflow check",
                train_w.dtype
            ));
        }
        if acts.dtype != WeightDtype::Bf16 {
            return Err(format!(
                "bf16-only graph capture: acts dtype {:?} != Bf16",
                acts.dtype
            ));
        }

        // CRITICAL: presize the half-precision staging buffer BEFORE capture
        // so `ensure_half_staging` inside the body is a no-op. Otherwise a
        // lazy grow during the captured forward bakes a freed pointer into
        // the graph (CUDA_ERROR_ILLEGAL_ADDRESS on replay).
        ctx.presize_half_staging_for_train(cfg, batch, seq_len, train_w.dtype)?;
        // Same contract for the bi upcast scratch: when the batch-invariant
        // flag is on, the captured body routes its typed GEMMs through
        // `with_bi_upcast_scratch` — grow it to the step's maximum now.
        // input_dim-aware: a trainable patch-embed input_proj can exceed
        // every backbone dim and would otherwise grow the scratch INSIDE
        // capture.
        let input_dim = mamba_input.len() / (batch * seq_len);
        ctx.presize_bi_upcast_scratch_for_train_with_input(
            cfg,
            batch,
            seq_len,
            input_dim,
            train_w.dtype,
        )?;
        // Snapshot pointers BEFORE capture so we can stash them after the
        // helper consumes the &mut borrows.
        let snap_input = mamba_input.cached_ptr();
        let snap_d_temporal = d_temporal.cached_ptr();
        let snap_grads_flat = grads.flat.cached_ptr();
        let snap_adam_m = adam.m.cached_ptr();
        let snap_adam_v = adam.v.cached_ptr();
        let snap_bias = bias.ptr();
        let snap_state_ssm = state.ssm_states.cached_ptr();
        let snap_state_conv = state.conv_states.cached_ptr();
        let snap_state_a_neg = state.a_neg_all.cached_ptr();
        let snap_a_neg = a_neg_all.cached_ptr();
        let snap_master_input = train_w.master.input_proj_w.cached_ptr();
        let snap_master_norm_f = train_w.master.norm_f_weight.cached_ptr();
        let snap_compute_input = train_w.compute.input_proj_w.ptr();
        let snap_compute_norm_f = train_w.compute.norm_f_weight.ptr();
        let snap_half_staging = ctx.half_staging_ptr();
        let snap_bi_upcast = ctx.bi_upcast_scratch_ptrs();
        let route_capacity = manifest.route_capacity;
        let has_gemm_work = batch != 0
            && seq_len != 0
            && (train_w.compute.input_proj_w.len_elems() != 0 || cfg.n_layers != 0);

        ctx.freeze_graph_scratch();
        let (graph, captured_gemm_plan) = unsafe {
            capture_into_graph_with_gemm_plan(ctx, route_capacity, manifest, || {
                grads.zero(&ctx.stream)?;
                gpu_forward_mamba_backbone_mixed(
                    ctx,
                    acts,
                    &train_w.compute,
                    mamba_input,
                    state,
                    scratch,
                )?;
                gpu_backward_mamba_backbone_mixed(
                    ctx,
                    d_temporal,
                    grads,
                    acts,
                    &train_w.compute,
                    a_neg_all,
                    scratch,
                )?;
                // ONE fused launch: every master tensor updated, the four
                // bulk typed shadows written in the same kernel (the old
                // per-tensor walk was 243 launches + a 4-cast/layer sync).
                step_multi(
                    ctx,
                    ctx.kernels.adamw_step_multi.get(train_w.dtype),
                    multi_plan,
                    adam,
                    bias.ptr(),
                )?;
                // f32-stays-f32 tensors still ride the (trimmed) sync walk.
                train_w.sync_master_to_compute(ctx)?;
                // Recompute a_neg = -exp(a_log) into BOTH a_neg_all buffers
                // used by forward and backward. Without these launches baked
                // into the captured body the graph replays a stale A-matrix
                // forever — a_log moves per AdamW step but the SSM kernel
                // reads the initial a_neg values until re-capture.
                recompute_a_neg_captured(
                    ctx,
                    &train_w.master.layers,
                    a_neg_all,
                    &state.a_neg_all,
                    cfg.d_inner(),
                    cfg.d_state,
                )?;
                Ok(())
            })
        }?;
        require_deterministic_gemm_graph_plan(
            ctx,
            has_gemm_work,
            captured_gemm_plan.as_ref(),
            "M1 bf16 training graph capture",
        )?;

        Ok(Self {
            graph,
            ctx_resources: ctx.resource_anchor(),
            batch,
            seq_len,
            dtype: train_w.dtype,
            captured_input_ptr: snap_input,
            captured_d_temporal_ptr: snap_d_temporal,
            captured_grads_flat_ptr: snap_grads_flat,
            captured_adam_m_ptr: snap_adam_m,
            captured_adam_v_ptr: snap_adam_v,
            captured_bias_factors_ptr: snap_bias,
            captured_state_ssm_states_ptr: snap_state_ssm,
            captured_state_conv_states_ptr: snap_state_conv,
            captured_state_a_neg_all_ptr: snap_state_a_neg,
            captured_a_neg_all_ptr: snap_a_neg,
            captured_master_input_proj_w_ptr: snap_master_input,
            captured_master_norm_f_ptr: snap_master_norm_f,
            captured_compute_input_proj_w_ptr: snap_compute_input,
            captured_compute_norm_f_ptr: snap_compute_norm_f,
            captured_half_staging_ptr: snap_half_staging,
            captured_bi_upcast_ptrs: snap_bi_upcast,
            captured_gemm_flags: {
                ctx.note_graph_capture();
                ctx.gemm_route()
            },
            has_gemm_work,
            captured_gemm_plan,
            captured_ctx_token: ctx.instance_token(),
            captured_stream_token: ctx.stream_token(),
        })
    }

    /// Replay the captured graph. Caller must have already:
    ///   1. Uploaded fresh `mamba_input` content
    ///   2. Computed loss + uploaded fresh `d_temporal`
    ///   3. Called `adam.advance()` and `bias.write(stream, bc1, bc2, lr)`
    ///      with the new step number's bias factors
    ///   4. Optionally zeroed `state` (or carried forward; user choice)
    ///
    /// Asserts every captured pointer still matches the live buffer's
    /// `cached_ptr()` — any reallocation since capture is a panic.
    pub fn replay(&self, ctx: &GpuCtx, rp: &MambaMixedReplay<'_>) -> Result<(), String> {
        if ctx.instance_token() != self.captured_ctx_token {
            return Err(
                "training_graph replay: GpuCtx differs from capture; re-capture instead".into(),
            );
        }
        if ctx.stream_token() != self.captured_stream_token {
            return Err(
                "training_graph replay: GpuCtx stream differs from capture; re-capture instead"
                    .into(),
            );
        }
        let MambaMixedReplay {
            train_w,
            adam,
            bias,
            grads,
            a_neg_all,
            mamba_input,
            d_temporal,
            state,
        } = *rp;
        assert_eq!(
            mamba_input.cached_ptr(),
            self.captured_input_ptr,
            "training_graph replay: mamba_input pointer changed since capture"
        );
        assert_eq!(
            d_temporal.cached_ptr(),
            self.captured_d_temporal_ptr,
            "training_graph replay: d_temporal pointer changed since capture"
        );
        assert_eq!(
            grads.flat.cached_ptr(),
            self.captured_grads_flat_ptr,
            "training_graph replay: grads.flat pointer changed since capture"
        );
        assert_eq!(
            adam.m.cached_ptr(),
            self.captured_adam_m_ptr,
            "training_graph replay: adam.m pointer changed since capture"
        );
        assert_eq!(
            adam.v.cached_ptr(),
            self.captured_adam_v_ptr,
            "training_graph replay: adam.v pointer changed since capture"
        );
        assert_eq!(
            bias.ptr(),
            self.captured_bias_factors_ptr,
            "training_graph replay: bias_factors pointer changed since capture"
        );
        assert_eq!(
            state.ssm_states.cached_ptr(),
            self.captured_state_ssm_states_ptr,
            "training_graph replay: state.ssm_states pointer changed since capture"
        );
        assert_eq!(
            state.conv_states.cached_ptr(),
            self.captured_state_conv_states_ptr,
            "training_graph replay: state.conv_states pointer changed since capture"
        );
        assert_eq!(
            state.a_neg_all.cached_ptr(),
            self.captured_state_a_neg_all_ptr,
            "training_graph replay: state.a_neg_all pointer changed since capture"
        );
        assert_eq!(
            a_neg_all.cached_ptr(),
            self.captured_a_neg_all_ptr,
            "training_graph replay: standalone a_neg_all pointer changed since capture"
        );
        assert_eq!(
            train_w.master.input_proj_w.cached_ptr(),
            self.captured_master_input_proj_w_ptr,
            "training_graph replay: master input_proj_w pointer changed since capture"
        );
        assert_eq!(
            train_w.master.norm_f_weight.cached_ptr(),
            self.captured_master_norm_f_ptr,
            "training_graph replay: master norm_f_weight pointer changed since capture"
        );
        assert_eq!(
            train_w.compute.input_proj_w.ptr(),
            self.captured_compute_input_proj_w_ptr,
            "training_graph replay: compute input_proj_w pointer changed since capture"
        );
        assert_eq!(
            train_w.compute.norm_f_weight.ptr(),
            self.captured_compute_norm_f_ptr,
            "training_graph replay: compute norm_f_weight pointer changed since capture"
        );
        assert_eq!(
            ctx.half_staging_ptr(),
            self.captured_half_staging_ptr,
            "training_graph replay: half_staging pointer changed since capture \
             (lazy grow during a previous step?)"
        );
        assert_eq!(
            ctx.bi_upcast_scratch_ptrs(),
            self.captured_bi_upcast_ptrs,
            "training_graph replay: bi_upcast_scratch pointer changed since \
             capture (a larger typed bi GEMM regrew the scratch after this \
             graph was captured — re-capture or presize for the larger shape)"
        );
        assert_eq!(
            ctx.gemm_route(),
            self.captured_gemm_flags,
            "training_graph replay: GEMM route changed since capture; re-capture instead"
        );
        with_validated_gemm_graph_launch(
            ctx,
            self.has_gemm_work,
            self.captured_gemm_plan.as_ref(),
            "training_graph replay",
            || {
                self.graph
                    .launch()
                    .map_err(|e| format!("training_graph launch: {e:?}"))
            },
        )
    }
}

impl Drop for GpuMambaTrainingStepGraph {
    fn drop(&mut self) {
        let _ = self.ctx_resources.stream.synchronize();
    }
}

// ════════════════════════════════════════════════════════════════════════
// f32 training step graph (no master/compute split, no half_staging).
// ════════════════════════════════════════════════════════════════════════

/// Buffers and parameter state borrowed into the captured M1 f32
/// training-step body (`grads.zero + forward + backward + AdamW` +
/// a_neg recompute).
pub struct MambaF32Capture<'a> {
    pub weights: &'a mut GpuMambaTrainWeights,
    pub adam: &'a GpuAdamW,
    pub bias: &'a AdamWBiasFactors,
    /// Fused multi-tensor AdamW chunk table (built at construction).
    pub multi_plan: &'a AdamWMultiPlan,
    pub grads: &'a mut GpuMambaGrads,
    pub acts: &'a mut GpuMambaBackboneActs,
    pub scratch: &'a mut GpuMambaScratch,
    /// Backward-side `a_neg_all` — a SEPARATE buffer from `state.a_neg_all`.
    pub a_neg_all: &'a GpuBuffer,
    /// Written by the captured forward — pointer baked in.
    pub temporal: &'a mut GpuBuffer,
    pub mamba_input: &'a GpuBuffer,
    pub d_temporal: &'a mut GpuBuffer,
    pub state: &'a mut GpuRecurrentState,
}

/// Live buffers checked against the captured pointers on every M1 f32
/// graph replay.
#[derive(Clone, Copy)]
pub struct MambaF32Replay<'a> {
    pub weights: &'a GpuMambaTrainWeights,
    pub adam: &'a GpuAdamW,
    pub bias: &'a AdamWBiasFactors,
    pub grads: &'a GpuMambaGrads,
    pub temporal: &'a GpuBuffer,
    /// Backward-side `a_neg_all` — a SEPARATE buffer from `state.a_neg_all`.
    pub a_neg_all: &'a GpuBuffer,
    pub mamba_input: &'a GpuBuffer,
    pub d_temporal: &'a GpuBuffer,
    pub state: &'a GpuRecurrentState,
}

/// CUDA-Graph holder for a single Mamba SSM f32 training step. Captures
/// `grads.zero + forward + backward + AdamW`. There's no
/// `sync_master_to_compute` because f32 training has no compute shadow —
/// weights are read directly during the next step's forward.
///
/// ```compile_fail
/// # use mamba_rs::mamba_ssm::gpu::training_graph::GpuMambaF32TrainingStepGraph;
/// # fn raw_launch(graph: &GpuMambaF32TrainingStepGraph) {
/// let _ = &graph.graph;
/// # }
/// ```
pub struct GpuMambaF32TrainingStepGraph {
    graph: CudaGraph,
    ctx_resources: Rc<crate::mamba_ssm::gpu::context::GpuCtxResources>,
    pub batch: usize,
    pub seq_len: usize,

    captured_input_ptr: u64,
    captured_d_temporal_ptr: u64,
    captured_grads_flat_ptr: u64,
    captured_adam_m_ptr: u64,
    captured_adam_v_ptr: u64,
    captured_bias_factors_ptr: u64,
    // ALL three GpuRecurrentState fields are accessed by forward.
    captured_state_ssm_states_ptr: u64,
    captured_state_conv_states_ptr: u64,
    captured_state_a_neg_all_ptr: u64,
    // `temporal` is written by forward; standalone `a_neg_all` (separate
    // from `state.a_neg_all` above) is read by backward.
    captured_temporal_ptr: u64,
    captured_a_neg_all_ptr: u64,
    captured_weights_input_proj_w_ptr: u64,
    captured_weights_norm_f_ptr: u64,
    // The trainer checks this route before replay.
    captured_gemm_flags: crate::mamba_ssm::gpu::context::GemmRoute,
    has_gemm_work: bool,
    captured_gemm_plan: Option<CapturedGemmGraphPlan>,
    captured_ctx_token: u64,
    captured_stream_token: usize,
}

impl GpuMambaF32TrainingStepGraph {
    /// Complete GEMM route captured with the graph.
    pub fn captured_gemm_flags(&self) -> crate::mamba_ssm::gpu::context::GemmRoute {
        self.captured_gemm_flags
    }

    /// Capture the f32 training step. Caller responsible for the same
    /// warmup contract as the mixed variant: run one eager step before
    /// calling this so cuBLAS has selected its kernels and any lazy
    /// resources have settled.
    /// # Safety
    ///
    /// The context, stream, cuBLAS handle, module functions, multi-plan, and
    /// every allocation or view used by the captured body must remain unchanged
    /// until this holder is destroyed and every replay has completed. Replay
    /// pointer checks diagnose drift but do not extend any CUDA lifetime.
    pub unsafe fn capture(
        ctx: &GpuCtx,
        cfg: &crate::config::MambaConfig,
        cap: MambaF32Capture<'_>,
        batch: usize,
        seq_len: usize,
        manifest: &PreparedGemmCaptureManifest,
    ) -> Result<Self, String> {
        let MambaF32Capture {
            weights,
            adam,
            bias,
            multi_plan,
            grads,
            acts,
            scratch,
            a_neg_all,
            temporal,
            mamba_input,
            d_temporal,
            state,
        } = cap;
        let snap_input = mamba_input.cached_ptr();
        let snap_d_temporal = d_temporal.cached_ptr();
        let snap_grads_flat = grads.flat.cached_ptr();
        let snap_adam_m = adam.m.cached_ptr();
        let snap_adam_v = adam.v.cached_ptr();
        let snap_bias = bias.ptr();
        let snap_state_ssm = state.ssm_states.cached_ptr();
        let snap_state_conv = state.conv_states.cached_ptr();
        let snap_state_a_neg = state.a_neg_all.cached_ptr();
        let snap_temporal = temporal.cached_ptr();
        let snap_a_neg = a_neg_all.cached_ptr();
        let snap_input_proj = weights.input_proj_w.cached_ptr();
        let snap_norm_f = weights.norm_f_weight.cached_ptr();
        let route_capacity = manifest.route_capacity;
        let has_gemm_work = batch != 0 && seq_len != 0;

        let cfg_local = *cfg;
        let (graph, captured_gemm_plan) = unsafe {
            capture_into_graph_with_gemm_plan(ctx, route_capacity, manifest, || {
                grads.zero(&ctx.stream)?;
                gpu_forward_mamba_backbone(
                    ctx,
                    temporal,
                    acts,
                    weights,
                    mamba_input,
                    state,
                    scratch,
                )?;
                gpu_backward_mamba_backbone(
                    ctx, d_temporal, grads, acts, weights, a_neg_all, scratch,
                )?;
                step_multi(
                    ctx,
                    ctx.kernels
                        .adamw_step_multi
                        .get(crate::mamba_ssm::gpu::dtype::WeightDtype::F32),
                    multi_plan,
                    adam,
                    bias.ptr(),
                )?;
                // Recompute a_neg after AdamW — see mixed graph above for
                // rationale. Without this the f32 SSM runs on a stale A-matrix
                // across every replay.
                recompute_a_neg_captured(
                    ctx,
                    &weights.layers,
                    a_neg_all,
                    &state.a_neg_all,
                    cfg_local.d_inner(),
                    cfg_local.d_state,
                )?;
                Ok(())
            })
        }?;
        require_deterministic_gemm_graph_plan(
            ctx,
            has_gemm_work,
            captured_gemm_plan.as_ref(),
            "M1 f32 training graph capture",
        )?;

        Ok(Self {
            graph,
            ctx_resources: ctx.resource_anchor(),
            batch,
            seq_len,
            captured_input_ptr: snap_input,
            captured_d_temporal_ptr: snap_d_temporal,
            captured_grads_flat_ptr: snap_grads_flat,
            captured_adam_m_ptr: snap_adam_m,
            captured_adam_v_ptr: snap_adam_v,
            captured_bias_factors_ptr: snap_bias,
            captured_state_ssm_states_ptr: snap_state_ssm,
            captured_state_conv_states_ptr: snap_state_conv,
            captured_state_a_neg_all_ptr: snap_state_a_neg,
            captured_temporal_ptr: snap_temporal,
            captured_a_neg_all_ptr: snap_a_neg,
            captured_weights_input_proj_w_ptr: snap_input_proj,
            captured_weights_norm_f_ptr: snap_norm_f,
            captured_gemm_flags: {
                ctx.note_graph_capture();
                ctx.gemm_route()
            },
            has_gemm_work,
            captured_gemm_plan,
            captured_ctx_token: ctx.instance_token(),
            captured_stream_token: ctx.stream_token(),
        })
    }

    pub fn replay(&self, ctx: &GpuCtx, rp: &MambaF32Replay<'_>) -> Result<(), String> {
        if ctx.instance_token() != self.captured_ctx_token {
            return Err(
                "f32 training_graph replay: GpuCtx differs from capture; re-capture instead".into(),
            );
        }
        if ctx.stream_token() != self.captured_stream_token {
            return Err(
                "f32 training_graph replay: GpuCtx stream differs from capture; \
                 re-capture instead"
                    .into(),
            );
        }
        self.captured_gemm_flags
            .ensure_current(ctx.gemm_route(), "f32 training_graph replay")?;
        let MambaF32Replay {
            weights,
            adam,
            bias,
            grads,
            temporal,
            a_neg_all,
            mamba_input,
            d_temporal,
            state,
        } = *rp;
        assert_eq!(
            mamba_input.cached_ptr(),
            self.captured_input_ptr,
            "f32 training_graph replay: mamba_input pointer changed since capture"
        );
        assert_eq!(
            d_temporal.cached_ptr(),
            self.captured_d_temporal_ptr,
            "f32 training_graph replay: d_temporal pointer changed since capture"
        );
        assert_eq!(
            grads.flat.cached_ptr(),
            self.captured_grads_flat_ptr,
            "f32 training_graph replay: grads.flat pointer changed since capture"
        );
        assert_eq!(
            adam.m.cached_ptr(),
            self.captured_adam_m_ptr,
            "f32 training_graph replay: adam.m pointer changed since capture"
        );
        assert_eq!(
            adam.v.cached_ptr(),
            self.captured_adam_v_ptr,
            "f32 training_graph replay: adam.v pointer changed since capture"
        );
        assert_eq!(
            bias.ptr(),
            self.captured_bias_factors_ptr,
            "f32 training_graph replay: bias_factors pointer changed since capture"
        );
        assert_eq!(
            state.ssm_states.cached_ptr(),
            self.captured_state_ssm_states_ptr,
            "f32 training_graph replay: state.ssm_states pointer changed since capture"
        );
        assert_eq!(
            state.conv_states.cached_ptr(),
            self.captured_state_conv_states_ptr,
            "f32 training_graph replay: state.conv_states pointer changed since capture"
        );
        assert_eq!(
            state.a_neg_all.cached_ptr(),
            self.captured_state_a_neg_all_ptr,
            "f32 training_graph replay: state.a_neg_all pointer changed since capture"
        );
        assert_eq!(
            temporal.cached_ptr(),
            self.captured_temporal_ptr,
            "f32 training_graph replay: temporal pointer changed since capture"
        );
        assert_eq!(
            a_neg_all.cached_ptr(),
            self.captured_a_neg_all_ptr,
            "f32 training_graph replay: standalone a_neg_all pointer changed since capture"
        );
        assert_eq!(
            weights.input_proj_w.cached_ptr(),
            self.captured_weights_input_proj_w_ptr,
            "f32 training_graph replay: input_proj_w pointer changed since capture"
        );
        assert_eq!(
            weights.norm_f_weight.cached_ptr(),
            self.captured_weights_norm_f_ptr,
            "f32 training_graph replay: norm_f_weight pointer changed since capture"
        );
        with_validated_gemm_graph_launch(
            ctx,
            self.has_gemm_work,
            self.captured_gemm_plan.as_ref(),
            "f32 training_graph replay",
            || {
                self.graph
                    .launch()
                    .map_err(|e| format!("f32 training_graph launch: {e:?}"))
            },
        )
    }
}

impl Drop for GpuMambaF32TrainingStepGraph {
    fn drop(&mut self) {
        let _ = self.ctx_resources.stream.synchronize();
    }
}