mamba-rs 0.7.1

//! GPU Mamba SSM forward/backward — mirrors CPU mamba/optimized.rs.
//!
//! Key difference from CPU: batches across ALL B samples simultaneously.
//! - SGEMM calls use batch=B*T (all samples, all timesteps)
//! - SSM recurrence + conv1d: sequential across T, parallel across B*d_inner
//!
//! ## Forward pipeline (per layer):
//! F1: RmsNorm → F2: in_proj SGEMM → F3: split+SiLU(gate)
//! → F4a: conv1d burnin+SiLU → F4b: x_proj SGEMM → F4c: dt_proj+softplus
//! → F4d: SSM burnin forward → F4e: gating → F5: out_proj SGEMM → F6: residual
//!
//! ## Backward pipeline (per layer):
//! B1: out_proj bwd → B2: gating bwd → B3: SSM BPTT + reductions
//! → B4: softplus bwd + dt_proj bwd → B5: x_proj bwd → B6: SiLU+conv1d bwd
//! → B7: in_proj bwd → B8: RmsNorm bwd + residual
//!
//! Source: CPU reference in train/forward.rs

use super::blas::gpu_gemm_bi_backward_grad_raw;
use super::buffers::GpuBuffer;
use super::context::GpuCtx;
use super::forward::{GpuMambaBackboneActs, GpuMambaLayerActs, GpuMambaScratch};
use super::launch::{grid_1d, grid_norm};
use super::weights::{
    GpuMambaGrads, GpuMambaLayerGrads, GpuMambaTrainLayerWeights, GpuMambaTrainWeights,
};
use cudarc::driver::PushKernelArg;
use std::sync::Arc;

// ---------------------------------------------------------------------------

pub fn gpu_backward_mamba_layer(
    ctx: &GpuCtx,
    d_temporal: &mut GpuBuffer,
    d_lw: &GpuMambaLayerGrads,
    acts: &GpuMambaLayerActs,
    lw: &GpuMambaTrainLayerWeights,
    a_neg_ptr: cudarc::driver::sys::CUdeviceptr,
    scratch: &mut GpuMambaScratch,
) -> Result<(), String> {
    let dims = scratch.dims; // Copy (GpuMambaDims is Copy)
    let bt = dims.bt();
    let dm = dims.d_model;
    let di = dims.d_inner;
    let ds = dims.d_state;
    let dt_rank = dims.dt_rank;
    let xdbl_dim = dims.xdbl_dim;
    let b = dims.batch;
    let t = dims.seq_len;
    let d_conv = dims.d_conv;

    // ===================================================================
    // B1: Batch out_proj backward
    // ===================================================================
    gpu_gemm_bi_backward_grad_raw(
        ctx,
        &mut scratch.d_gated,
        (&d_lw.out_proj_w, None),
        d_temporal,
        &acts.gated,
        lw.out_proj_w.cached_ptr(),
        (bt, di, dm),
    )?;

    // ===================================================================
    // B2: Gating backward
    // ===================================================================
    // gating_backward(d_y, d_gate_pre, d_gated, y, gate_pre, n) - the
    // post-SiLU activation is recomputed inside the kernel.
    {
        let n = (bt * di) as i32;
        let mut builder = ctx.stream.launch_builder(&ctx.kernels.gating_backward);
        let di_i = di as i32;
        let proj_stride = (2 * di) as i32;
        let gate_off = di as i32;
        builder.arg(scratch.d_y.inner_mut());
        // Writes the gate half of d_proj directly.
        builder.arg(scratch.d_proj.inner_mut());
        builder.arg(scratch.d_gated.inner());
        builder.arg(acts.y.inner());
        builder.arg(acts.proj.inner()); // gate half, read through the same geometry
        builder.arg(&n);
        builder.arg(&di_i);
        builder.arg(&proj_stride);
        builder.arg(&gate_off);
        unsafe { builder.launch(grid_1d(bt * di)) }
            .map_err(|e| format!("gating_backward mamba: {:?}", e))?;
    }

    // ===================================================================
    // B3: SSM BPTT + reductions
    // ===================================================================
    // Gather B and C from xdbl for ssm_backward_local
    {
        let bt_i = bt as i32;
        let xdbl_i = xdbl_dim as i32;
        let ds_i = ds as i32;
        let b_offset = dt_rank as i32;
        let c_offset = (dt_rank + ds) as i32;
        // Fused gather B+C from xdbl (saves 1 kernel launch). The
        // parallel route gathers T-major so the scan's per-(d, n) lane
        // reads contiguous t-runs; identical values either way.
        let tmajor = dims.scan_mode.use_parallel(t, ds);
        let kernel = if tmajor {
            &ctx.kernels.gather_bc_cols_tmajor
        } else {
            &ctx.kernels.gather_bc_cols
        };
        let t_i = t as i32;
        let mut builder = ctx.stream.launch_builder(kernel);
        builder.arg(scratch.d_b_reduced.inner_mut());
        builder.arg(scratch.d_c_reduced.inner_mut());
        builder.arg(acts.xdbl.inner());
        builder.arg(&bt_i);
        if tmajor {
            builder.arg(&t_i);
        }
        builder.arg(&xdbl_i);
        builder.arg(&ds_i);
        builder.arg(&b_offset);
        builder.arg(&c_offset);
        unsafe { builder.launch(grid_1d(bt * ds)) }
            .map_err(|e| format!("gather_bc_cols bwd mamba: {:?}", e))?;
    }

    // No zeroing on either route now: the SEQUENTIAL kernel `=`-stores
    // the full domain from a register accumulator, and the PARALLEL fold
    // writes one partial SLOT per chunk (`=`, every slot covered) that
    // the chunked reducer folds afterwards.

    // ssm_backward_local(h_saved, delta_saved, u_saved, B_saved, C_saved, a_neg, D,
    //   dy, d_delta, d_u, d_B_local, d_C_local, d_D_local, d_a_log_local,
    //   batch, T, d_inner, d_state)
    //
    // Dispatch mirrors the mixed-precision path: the sequential kernel keeps
    // per-(b,d) state in registers and contains a `d_state > 64` early-return
    // guard, so d_state in (64, 256] MUST take the parallel reverse-scan
    // kernel (it used to silently no-op, leaving stale scratch as gradients).
    // Long T also prefers the parallel kernel for wall-clock.
    {
        let b_i = b as i32;
        let t_i = t as i32;
        let di_i = di as i32;
        let ds_i = ds as i32;
        let use_parallel = dims.scan_mode.use_parallel(t, ds);
        let use_fold = use_parallel && di.is_multiple_of(super::launch::SCAN_BWD_DGROUP);
        let kernel = if use_fold {
            ctx.kernels
                .ssm_parallel_bwd_fold_typed
                .get(super::dtype::WeightDtype::F32)
        } else if use_parallel {
            ctx.kernels
                .ssm_parallel_bwd_typed
                .get(super::dtype::WeightDtype::F32)
        } else {
            &ctx.kernels.ssm_backward_local
        };
        if use_parallel && !use_fold {
            // The ungrouped parallel kernel accumulates its per-sample
            // d_a_log rows chunk by chunk with additions; it needs a zero
            // slate every step, unlike the fold, which assigns every slot.
            scratch
                .d_a_log_local
                .zero(&ctx.stream)
                .map_err(|e| format!("zero d_a_log_local: {e:?}"))?;
        }
        let mut builder = ctx.stream.launch_builder(kernel);
        builder.arg(acts.h_saved.inner());
        builder.arg(acts.delta.inner());
        builder.arg(acts.u.inner());
        builder.arg(scratch.d_b_reduced.inner()); // B_saved
        builder.arg(scratch.d_c_reduced.inner()); // C_saved
        builder.arg(&a_neg_ptr); // raw ptr at layer offset
        let dp_ptr = lw.d_param.cached_ptr();
        builder.arg(&dp_ptr);
        builder.arg(scratch.d_y.inner()); // dy (from gating backward)
        if use_fold {
            // The fold's epilogue applies the softplus derivative inline
            // (round-first) and emits the PRE-softplus dt gradient - the
            // separate softplus backward launch is skipped on this route.
            builder.arg(scratch.d_delta_raw.inner_mut());
            builder.arg(acts.delta_raw.inner());
        } else {
            builder.arg(scratch.d_delta.inner_mut());
        }
        builder.arg(scratch.d_u.inner_mut());
        builder.arg(scratch.d_b_local.inner_mut());
        builder.arg(scratch.d_c_local.inner_mut());
        builder.arg(scratch.d_d_local.inner_mut());
        builder.arg(scratch.d_a_log_local.inner_mut());
        builder.arg(&b_i);
        builder.arg(&t_i);
        builder.arg(&di_i);
        builder.arg(&ds_i);
        // The slim tape rides the h_saved buffer; the flag picks the
        // kernel's replay path. Sequential kernel keeps its signature.
        let tape_p = acts.h_saved.cached_ptr();
        let slim_i: i32 = i32::from(super::launch::scan_tape_slim());
        if use_parallel {
            builder.arg(&tape_p);
            builder.arg(&slim_i);
        }
        let cfg = if use_fold {
            super::launch::grid_parallel_scan_bwd_fold(b, di, ds, 4)
        } else if use_parallel {
            super::launch::grid_parallel_scan_bwd(b, di)
        } else {
            grid_1d(b * di)
        };
        unsafe { builder.launch(cfg) }
            .map_err(|e| format!("ssm bwd (parallel={use_parallel}): {e:?}"))?;
    }

    // Reductions: sum per-sample gradients across batch/d_inner.
    // The fused dB+dC kernel `=`-stores the full domain (0.0f + sum),
    // so the reduction targets need no zeroing even though they held
    // gathered B/C values from the SSM backward — full overwrite.
    {
        let b_i = b as i32;
        let t_i = t as i32;
        let di_i = di as i32;
        let ds_i = ds as i32;
        // Fused d_B + d_C reduction. The PARALLEL route writes its
        // locals T-major (the parallel-route tape layout) and takes the tmajor twin;
        // the sequential route keeps the historical layout + reducer.
        // Values and output layout are identical either way.
        let use_parallel = dims.scan_mode.use_parallel(t, ds);
        let use_fold = use_parallel && di.is_multiple_of(super::launch::SCAN_BWD_DGROUP);
        let reduce_bc = if use_parallel {
            ctx.kernels
                .ssm_reduce_d_bc_tmajor_typed
                .get(super::dtype::WeightDtype::F32)
        } else {
            ctx.kernels
                .ssm_reduce_d_bc_typed
                .get(super::dtype::WeightDtype::F32)
        };
        // Under the fold the locals hold one pre-summed row per d group,
        // so the reducer's d depth is the group count.
        let reduce_di_i = if use_fold {
            (di / super::launch::SCAN_BWD_DGROUP) as i32
        } else {
            di_i
        };
        let mut builder = ctx.stream.launch_builder(reduce_bc);
        builder.arg(scratch.d_b_reduced.inner_mut());
        builder.arg(scratch.d_c_reduced.inner_mut());
        builder.arg(scratch.d_b_local.inner());
        builder.arg(scratch.d_c_local.inner());
        builder.arg(&b_i);
        builder.arg(&t_i);
        builder.arg(&reduce_di_i);
        builder.arg(&ds_i);
        unsafe { builder.launch(grid_1d(bt * ds)) }
            .map_err(|e| format!("ssm_reduce_d_BC mamba: {:?}", e))?;
        // d_D reduction
        let mut builder = ctx.stream.launch_builder(&ctx.kernels.ssm_reduce_d_d);
        let _p = d_lw.d_param.ptr();
        builder.arg(&_p);
        builder.arg(scratch.d_d_local.inner());
        builder.arg(&b_i);
        builder.arg(&di_i);
        unsafe { builder.launch(grid_1d(di)) }
            .map_err(|e| format!("ssm_reduce_d_D mamba: {:?}", e))?;
        // d_a_log reduction - the parallel fold hands over chunk-partial
        // rows; the sequential kernel keeps the per-sample accumulator
        // layout. The chunked reducer folds one sample's slots first and
        // only then adds across the batch, reproducing the retired
        // accumulate-then-reduce association exactly.
        let _p = d_lw.a_log.ptr();
        // Only the fold writes chunk-partial rows; the ungrouped parallel
        // kernel and the sequential kernel both leave one accumulated row
        // per sample, which the flat reducer folds.
        let use_fold =
            dims.scan_mode.use_parallel(t, ds) && di.is_multiple_of(super::launch::SCAN_BWD_DGROUP);
        if use_fold {
            let nc = t.div_ceil(super::launch::SCAN_CHUNK).max(1) as i32;
            let mut builder = ctx
                .stream
                .launch_builder(&ctx.kernels.ssm_reduce_d_a_log_chunks);
            builder.arg(&_p);
            builder.arg(scratch.d_a_log_local.inner());
            builder.arg(&b_i);
            builder.arg(&nc);
            builder.arg(&di_i);
            builder.arg(&ds_i);
            unsafe { builder.launch(grid_1d(di * ds)) }
                .map_err(|e| format!("ssm_reduce_d_a_log_chunks mamba: {:?}", e))?;
        } else {
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.ssm_reduce_d_a_log);
            builder.arg(&_p);
            builder.arg(scratch.d_a_log_local.inner());
            builder.arg(&b_i);
            builder.arg(&di_i);
            builder.arg(&ds_i);
            unsafe { builder.launch(grid_1d(di * ds)) }
                .map_err(|e| format!("ssm_reduce_d_a_log mamba: {:?}", e))?;
        }
    }

    // d_xdbl assembly moved to ONE pack_xdbl_cols launch after the
    // dt_proj backward below: the dt|B|C ranges exactly tile the row, so
    // the old zero + three scatter_adds collapse once all three sources
    // (d_dt_input, d_b_reduced, d_c_reduced) are ready. Nothing reads
    // d_xdbl before the x_proj backward, and nothing overwrites the two
    // reduce outputs in between.

    // ===================================================================
    // B4: Softplus backward + dt_proj backward
    // ===================================================================
    // softplus_backward(dx, x_saved, dy, n) - only the non-fold routes
    // need it; the fold epilogue already emitted d_delta_raw.
    if !(dims.scan_mode.use_parallel(t, ds) && di.is_multiple_of(super::launch::SCAN_BWD_DGROUP)) {
        let n = (bt * di) as i32;
        let mut builder = ctx.stream.launch_builder(&ctx.kernels.softplus_bwd);
        builder.arg(scratch.d_delta_raw.inner_mut()); // dx output
        builder.arg(acts.delta_raw.inner()); // x_saved (pre-softplus activation)
        builder.arg(scratch.d_delta.inner()); // dy (upstream gradient from SSM)
        builder.arg(&n);
        unsafe { builder.launch(grid_1d(bt * di)) }
            .map_err(|e| format!("softplus_bwd mamba: {:?}", e))?;
    }

    // Gather xdbl dt portion for x_saved (separate buffer from dx output)
    {
        let bt_i = bt as i32;
        let xdbl_i = xdbl_dim as i32;
        let dt_i = dt_rank as i32;
        let offset: i32 = 0;
        let mut builder = ctx.stream.launch_builder(&ctx.kernels.gather_cols);
        builder.arg(scratch.dt_xdbl_buf.inner_mut());
        builder.arg(acts.xdbl.inner());
        builder.arg(&bt_i);
        builder.arg(&xdbl_i);
        builder.arg(&dt_i);
        builder.arg(&offset);
        unsafe { builder.launch(grid_1d(bt * dt_rank)) }
            .map_err(|e| format!("gather dt x_saved mamba: {:?}", e))?;
    }

    gpu_gemm_bi_backward_grad_raw(
        ctx,
        &mut scratch.d_dt_input, // dx [B*T*dt_rank]
        (&d_lw.dt_proj_w, Some(&d_lw.dt_proj_b)),
        &scratch.d_delta_raw, // dy [B*T*d_inner]
        &scratch.dt_xdbl_buf, // x_saved [B*T*dt_rank]
        lw.dt_proj_w.cached_ptr(),
        (bt, dt_rank, di),
    )?;

    // One-kernel d_xdbl assembly: dt | B | C ranges tile the row.
    {
        let bt_i = bt as i32;
        let dt_i = dt_rank as i32;
        let ds_i = ds as i32;
        let mut builder = ctx.stream.launch_builder(
            ctx.kernels
                .pack_xdbl_cols_typed
                .get(super::dtype::WeightDtype::F32),
        );
        builder.arg(scratch.d_xdbl.inner_mut());
        builder.arg(scratch.d_dt_input.inner());
        builder.arg(scratch.d_b_reduced.inner());
        builder.arg(scratch.d_c_reduced.inner());
        builder.arg(&bt_i);
        builder.arg(&dt_i);
        builder.arg(&ds_i);
        unsafe { builder.launch(grid_1d(bt * xdbl_dim)) }
            .map_err(|e| format!("pack_xdbl_cols mamba: {:?}", e))?;
    }

    // ===================================================================
    // B5: x_proj backward
    // ===================================================================
    gpu_gemm_bi_backward_grad_raw(
        ctx,
        &mut scratch.d_u_xproj,
        (&d_lw.x_proj_w, None),
        &scratch.d_xdbl,
        &acts.u,
        lw.x_proj_w.cached_ptr(),
        (bt, di, xdbl_dim),
    )?;

    // Accumulate d_u += d_u_xproj
    {
        let n = (bt * di) as i32;
        let mut builder = ctx.stream.launch_builder(&ctx.kernels.vec_add_inplace);
        builder.arg(scratch.d_u.inner_mut());
        builder.arg(scratch.d_u_xproj.inner());
        builder.arg(&n);
        unsafe { builder.launch(grid_1d(bt * di)) }
            .map_err(|e| format!("vec_add d_u mamba: {:?}", e))?;
    }

    // ===================================================================
    // B6: SiLU backward + Conv1d backward (fused burnin kernel)
    // ===================================================================
    // Rule B (no atomicAdd): two-stage launch.
    // Stage 1: conv1d_burnin_bwd writes per-(b,d) partials into axis0_partials
    //          split as [weight_partials | bias_partials]:
    //            weight at offset 0, size B * d_inner * d_conv
    //            bias   at offset B*d_inner*d_conv, size B * d_inner
    // Stage 2: two reduce_sum_axis0 launches reduce across B → d_lw grad slices.
    {
        let b_i = b as i32;
        let t_i = t as i32;
        let di_i = di as i32;
        let dc_i = d_conv as i32;
        let n_tiles = t.div_ceil(128);
        let weight_partials_elems = b * n_tiles * di * d_conv;
        let bias_offset_bytes = (weight_partials_elems * std::mem::size_of::<f32>()) as u64;
        let rows_i = (b * n_tiles) as i32;
        let axis0_base = scratch.axis0_partials.cached_ptr();
        let wp_ptr = axis0_base;
        let bp_ptr = axis0_base + bias_offset_bytes;
        // Stage 1: one tiled kernel walks d_x, the tap partials and the
        // bias partials together, recomputing the pre-activation from the
        // x window (no tape). The x half of d_proj is written in place
        // (stride 2*d_inner, offset 0).
        {
            let f32k = super::dtype::WeightDtype::F32;
            let mut builder = ctx
                .stream
                .launch_builder(ctx.kernels.conv1d_bwd_tiled_typed.get(f32k));
            builder.arg(scratch.d_proj.inner_mut());
            builder.arg(&wp_ptr); // d_weight_partials
            builder.arg(&bp_ptr); // d_bias_partials
            builder.arg(scratch.d_u.inner());
            builder.arg(acts.proj.inner()); // x half, read through the row stride
            builder.arg(acts.conv_states.inner()); // carry-in window (conv_init)
            let cw_ptr = lw.conv1d_weight.cached_ptr();
            let cb_ptr = lw.conv1d_bias.cached_ptr();
            builder.arg(&cw_ptr);
            builder.arg(&cb_ptr);
            builder.arg(&b_i);
            builder.arg(&t_i);
            builder.arg(&di_i);
            builder.arg(&dc_i);
            let proj_stride = (2 * di) as i32;
            let x_off = 0i32;
            builder.arg(&proj_stride);
            builder.arg(&x_off);
            builder.arg(&proj_stride); // x row stride
            unsafe { builder.launch(super::launch::grid_conv_tiled(b, di, t)) }
                .map_err(|e| format!("conv1d_bwd_tiled mamba: {:?}", e))?;
        }
        // Stage 2a: reduce weight partials [B, d_inner*d_conv] → d_lw.conv1d_weight.
        {
            let block_dim = (b as u32).next_power_of_two().clamp(32, 256);
            let accumulate_i: i32 = 1;
            let dim_w = (di * d_conv) as i32;
            let p = d_lw.conv1d_weight.ptr();
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.reduce_sum_axis0);
            builder.arg(&p);
            builder.arg(&wp_ptr);
            builder.arg(&rows_i);
            builder.arg(&dim_w);
            builder.arg(&accumulate_i);
            let cfg = cudarc::driver::LaunchConfig {
                grid_dim: ((di * d_conv) as u32, 1, 1),
                block_dim: (block_dim, 1, 1),
                shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
            };
            unsafe { builder.launch(cfg) }
                .map_err(|e| format!("conv1d_burnin_bwd mamba weight final: {:?}", e))?;
        }
        // Stage 2b: reduce bias partials [B, d_inner] → d_lw.conv1d_bias.
        {
            let block_dim = (b as u32).next_power_of_two().clamp(32, 256);
            let accumulate_i: i32 = 1;
            let p = d_lw.conv1d_bias.ptr();
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.reduce_sum_axis0);
            builder.arg(&p);
            builder.arg(&bp_ptr);
            builder.arg(&rows_i);
            builder.arg(&di_i);
            builder.arg(&accumulate_i);
            let cfg = cudarc::driver::LaunchConfig {
                grid_dim: (di as u32, 1, 1),
                block_dim: (block_dim, 1, 1),
                shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
            };
            unsafe { builder.launch(cfg) }
                .map_err(|e| format!("conv1d_burnin_bwd mamba bias final: {:?}", e))?;
        }
    }

    // ===================================================================
    // B7: Batch in_proj backward
    // ===================================================================
    // d_proj is already complete: the gating backward wrote its gate
    // half and the conv dx pass wrote its x half, both in place.

    gpu_gemm_bi_backward_grad_raw(
        ctx,
        &mut scratch.d_norm,
        (&d_lw.in_proj_w, None),
        &scratch.d_proj,
        &acts.post_norm,
        lw.in_proj_w.cached_ptr(),
        (bt, dm, 2 * di),
    )?;

    // ===================================================================
    // B8: RmsNorm backward + residual
    // ===================================================================
    // Rule B (no atomicAdd): two-stage launch.
    // Stage 1: rmsnorm_bwd writes per-sample per-dim partials to axis0_partials[bt * dm].
    // Stage 2: reduce_sum_axis0 reduces across bt → d_lw.norm_weight (accumulate=1).
    {
        let bt_i = bt as i32;
        let dm_i = dm as i32;
        let nw_ptr = lw.norm_weight.cached_ptr();
        let axis0_ptr = scratch.axis0_partials.cached_ptr();
        // Stage 1
        {
            // dx = d_temporal with accumulate=1: the residual-path add
            // (old separate vec_add_inplace) is folded into the store.
            let accumulate_dx: i32 = 1;
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.rmsnorm_bwd);
            builder.arg(d_temporal.inner_mut());
            builder.arg(&axis0_ptr); // d_scale_partials
            builder.arg(scratch.d_norm.inner());
            builder.arg(acts.residual.inner()); // x = input before norm
            builder.arg(&nw_ptr);
            builder.arg(acts.rms_vals.inner());
            builder.arg(&bt_i);
            builder.arg(&dm_i);
            builder.arg(&accumulate_dx);
            // No typed consumer on the pure-f32 lane.
            let no_mirror: cudarc::driver::sys::CUdeviceptr = 0;
            builder.arg(&no_mirror);
            unsafe { builder.launch(grid_norm(bt, dm)) }
                .map_err(|e| format!("rmsnorm_bwd mamba partial: {:?}", e))?;
        }
        // Stage 2
        {
            let block_dim = (bt as u32).next_power_of_two().clamp(32, 256);
            let accumulate_i: i32 = 1;
            let p = d_lw.norm_weight.ptr();
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.reduce_sum_axis0);
            builder.arg(&p);
            builder.arg(&axis0_ptr);
            builder.arg(&bt_i);
            builder.arg(&dm_i);
            builder.arg(&accumulate_i);
            let cfg = cudarc::driver::LaunchConfig {
                grid_dim: (dm as u32, 1, 1),
                block_dim: (block_dim, 1, 1),
                shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
            };
            unsafe { builder.launch(cfg) }
                .map_err(|e| format!("rmsnorm_bwd mamba final: {:?}", e))?;
        }
    }

    // (Residual add folded into the rmsnorm_bwd accumulate store above.)

    Ok(())
}

/// GPU Mamba backbone backward: layers in reverse + input_proj backward.
///
/// Mirrors CPU `backward_mamba_backbone_batched` from train/forward.rs.
///
/// **IMPORTANT**: weight gradients in `d_mamba` are **accumulated** via
/// `beta=1.0` in `gpu_gemm_bi_backward_dw_grad`. The caller MUST call
/// [`GpuMambaGrads::zero`] before each training step if the buffer is
/// reused across iterations; otherwise gradients from step N−1 pollute
/// step N and the optimizer sees doubled updates.
pub fn gpu_backward_mamba_backbone(
    ctx: &GpuCtx,
    d_temporal: &mut GpuBuffer,
    d_mamba: &GpuMambaGrads,
    acts: &GpuMambaBackboneActs,
    mamba_w: &GpuMambaTrainWeights,
    a_neg_all: &GpuBuffer,
    scratch: &mut GpuMambaScratch,
) -> Result<(), String> {
    let dims = scratch.dims; // Copy (GpuMambaDims is Copy)
    let bt = dims.bt();

    // norm_f backward — before reverse layer loop.
    // Rule B (no atomicAdd): two-stage launch.
    // Stage 1: rmsnorm_bwd writes per-sample per-dim partials to axis0_partials.
    // Stage 2: reduce_sum_axis0 reduces across bt → d_mamba.norm_f_weight (accumulate=1).
    {
        let bt_i = bt as i32;
        let dm_i = dims.d_model as i32;
        let axis0_ptr = scratch.axis0_partials.cached_ptr();
        // Stage 1
        {
            // In-place: dx aliases dy (= d_temporal). Kernel-safe: the
            // sum pass reads all dy before the barrier, and the write
            // pass reads dy[off+i] before storing the same element.
            // Kills the d_norm temp + the copy-back.
            let dt_ptr = d_temporal.cached_ptr();
            let accumulate_dx: i32 = 0;
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.rmsnorm_bwd);
            builder.arg(&dt_ptr); // dx (in place)
            builder.arg(&axis0_ptr); // d_scale_partials
            builder.arg(&dt_ptr); // dy (upstream gradient)
            builder.arg(acts.norm_f_input.inner()); // saved pre-norm input
            let nf_ptr = mamba_w.norm_f_weight.cached_ptr();
            builder.arg(&nf_ptr); // scale
            builder.arg(acts.norm_f_rms.inner()); // saved rms
            builder.arg(&bt_i);
            builder.arg(&dm_i);
            builder.arg(&accumulate_dx);
            let no_mirror: cudarc::driver::sys::CUdeviceptr = 0;
            builder.arg(&no_mirror);
            unsafe { builder.launch(grid_norm(bt, dims.d_model)) }
                .map_err(|e| format!("rmsnorm_bwd norm_f partial: {:?}", e))?;
        }
        // Stage 2
        {
            let block_dim = (bt as u32).next_power_of_two().clamp(32, 256);
            let accumulate_i: i32 = 1;
            let p = d_mamba.norm_f_weight.ptr();
            let mut builder = ctx.stream.launch_builder(&ctx.kernels.reduce_sum_axis0);
            builder.arg(&p);
            builder.arg(&axis0_ptr);
            builder.arg(&bt_i);
            builder.arg(&dm_i);
            builder.arg(&accumulate_i);
            let cfg = cudarc::driver::LaunchConfig {
                grid_dim: (dims.d_model as u32, 1, 1),
                block_dim: (block_dim, 1, 1),
                shared_mem_bytes: (block_dim as usize * std::mem::size_of::<f32>()) as u32,
            };
            unsafe { builder.launch(cfg) }
                .map_err(|e| format!("rmsnorm_bwd norm_f final: {:?}", e))?;
        }
        // (dx written in place into d_temporal — no copy-back.)
    }

    // Mamba layers in reverse — per-layer a_neg offset
    let a_neg_per_layer = dims.d_inner * dims.d_state;
    for layer_idx in (0..dims.n_layers).rev() {
        let base = a_neg_all.raw_ptr(&ctx.stream);
        let a_neg_ptr = base + (layer_idx * a_neg_per_layer * std::mem::size_of::<f32>()) as u64;

        gpu_backward_mamba_layer(
            ctx,
            d_temporal,
            &d_mamba.layers[layer_idx],
            &acts.layers[layer_idx],
            &mamba_w.layers[layer_idx],
            a_neg_ptr,
            scratch,
        )?;
    }

    // Input projection backward (dx discarded — input embedding detached)
    gpu_gemm_bi_backward_grad_raw(
        ctx,
        &mut scratch.d_input_proj_dx,
        (&d_mamba.input_proj_w, Some(&d_mamba.input_proj_b)),
        d_temporal,
        &acts.input_proj_inputs,
        mamba_w.input_proj_w.cached_ptr(),
        (bt, dims.mamba_input_dim, dims.d_model),
    )?;

    Ok(())
}

// ---------------------------------------------------------------------------
// GPU Mamba target forward (batched B*T, nosave kernels)
// ---------------------------------------------------------------------------

/// Scratch buffers for GPU Mamba target forward (batched B*T pipeline).
///
/// Uses nosave burnin kernels (no h_saved/conv_states writes).
/// Per-layer conv/SSM state via flat buffers with layer offsets (matches online).
///
/// Uses batched B*T pipeline instead of per-sample step-by-step approach.
pub struct GpuMambaTargetScratch {
    // Batched B*T scratch (reusable from online forward scratch)
    pub proj_flat: GpuBuffer,   // [B*T*2*d_inner]
    pub u: GpuBuffer,           // [B*T*d_inner]
    pub xdbl: GpuBuffer,        // [B*T*xdbl_dim]
    pub dt_gather: GpuBuffer,   // [B*T*dt_rank]
    pub delta: GpuBuffer,       // [B*T*d_inner]
    pub gated: GpuBuffer,       // [B*T*d_inner] the scan's gated store
    pub out_flat: GpuBuffer,    // [B*T*d_model]
    pub residual: GpuBuffer,    // [B*T*d_model] (saved before RmsNorm)
    pub rms_discard: GpuBuffer, // [B*T] (RmsNorm scalars, discarded)
    pub b_gathered: GpuBuffer,  // [B*T*d_state]
    pub c_gathered: GpuBuffer,  // [B*T*d_state]
    // Per-layer state (flat, with layer offsets — same pattern as online)
    pub conv_states: GpuBuffer, // [B*n_layers*d_inner*d_conv]
    pub ssm_states: GpuBuffer,  // [B*n_layers*d_inner*d_state]
    /// Dimensions this scratch was allocated for.
    pub dims: super::forward::GpuMambaDims,
}

impl GpuMambaTargetScratch {
    /// Allocate scratch buffers for batched Mamba target forward.
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        dims: &super::forward::GpuMambaDims,
    ) -> Result<Self, String> {
        let batch = dims.batch;
        let d_model = dims.d_model;
        let d_inner = dims.d_inner;
        let d_state = dims.d_state;
        let d_conv = dims.d_conv;
        let dt_rank = dims.dt_rank;
        let n_layers = dims.n_layers;
        let bt = batch * dims.seq_len;
        let xdbl_dim = dt_rank + 2 * d_state;

        Ok(Self {
            proj_flat: GpuBuffer::zeros(stream, bt * 2 * d_inner)?,
            u: GpuBuffer::zeros(stream, bt * d_inner)?,
            xdbl: GpuBuffer::zeros(stream, bt * xdbl_dim)?,
            dt_gather: GpuBuffer::zeros(stream, bt * dt_rank)?,
            delta: GpuBuffer::zeros(stream, bt * d_inner)?,
            gated: GpuBuffer::zeros(stream, bt * d_inner)?,
            out_flat: GpuBuffer::zeros(stream, bt * d_model)?,
            residual: GpuBuffer::zeros(stream, bt * d_model)?,
            rms_discard: GpuBuffer::zeros(stream, bt)?,
            b_gathered: GpuBuffer::zeros(stream, bt * d_state)?,
            c_gathered: GpuBuffer::zeros(stream, bt * d_state)?,
            conv_states: GpuBuffer::zeros(stream, batch * n_layers * d_inner * d_conv)?,
            ssm_states: GpuBuffer::zeros(stream, batch * n_layers * d_inner * d_state)?,
            dims: *dims,
        })
    }
}

// ---------------------------------------------------------------------------
// Mixed-precision target scratch for end-to-end bf16/f16 prefill.
//
// All activation-layer tensors are DtypedBuf (bf16/f16), residual stays f32
// (HF residual_in_fp32), rms_discard stays f32 (per-batch stats).
// ---------------------------------------------------------------------------

use super::buffers::DtypedBuf;
use super::dtype::WeightDtype;

pub struct GpuMambaTargetMixedScratch {
    pub proj_flat: DtypedBuf,
    pub u: DtypedBuf,
    pub xdbl: DtypedBuf,
    pub dt_gather: DtypedBuf,
    pub delta: DtypedBuf,
    pub gated: DtypedBuf,
    pub out_flat: DtypedBuf,
    /// Residual accumulator — f32 across layers (HF residual_in_fp32).
    pub residual: GpuBuffer,
    pub rms_discard: GpuBuffer,
    pub b_gathered: DtypedBuf,
    pub c_gathered: DtypedBuf,
    pub dims: super::forward::GpuMambaDims,
    pub dtype: WeightDtype,
}

impl GpuMambaTargetMixedScratch {
    pub fn new(
        stream: &Arc<cudarc::driver::CudaStream>,
        dims: &super::forward::GpuMambaDims,
        dtype: WeightDtype,
    ) -> Result<Self, String> {
        if matches!(dtype, WeightDtype::F32) {
            return Err("GpuMambaTargetMixedScratch requires bf16 or f16 dtype".to_string());
        }
        let batch = dims.batch;
        let d_model = dims.d_model;
        let d_inner = dims.d_inner;
        let d_state = dims.d_state;
        let dt_rank = dims.dt_rank;
        let bt = batch * dims.seq_len;
        let xdbl_dim = dt_rank + 2 * d_state;

        Ok(Self {
            proj_flat: DtypedBuf::zeros(stream, bt * 2 * d_inner, dtype)?,
            u: DtypedBuf::zeros(stream, bt * d_inner, dtype)?,
            xdbl: DtypedBuf::zeros(stream, bt * xdbl_dim, dtype)?,
            dt_gather: DtypedBuf::zeros(stream, bt * dt_rank, dtype)?,
            delta: DtypedBuf::zeros(stream, bt * d_inner, dtype)?,
            gated: DtypedBuf::zeros(stream, bt * d_inner, dtype)?,
            out_flat: DtypedBuf::zeros(stream, bt * d_model, dtype)?,
            residual: GpuBuffer::zeros(stream, bt * d_model)?,
            rms_discard: GpuBuffer::zeros(stream, bt)?,
            b_gathered: DtypedBuf::zeros(stream, bt * d_state, dtype)?,
            c_gathered: DtypedBuf::zeros(stream, bt * d_state, dtype)?,
            dims: *dims,
            dtype,
        })
    }
}