mamba-rs 0.6.5

//! Compile and register Mamba-3 SISO CUDA kernels.
//!
//! The Mamba-3 kernel registry, compiled via NVRTC at runtime.
//! Separate from Mamba SSM's `MambaKernels` — different pipeline, no conv1d.

use crate::mamba_ssm::gpu::kernels::{HalfKernel, TypedKernel};
use cudarc::driver::{CudaContext, CudaFunction, CudaModule};
use std::sync::Arc;

/// All compiled Mamba-3 SISO CUDA kernels.
pub struct Mamba3Kernels {
    _module: Arc<CudaModule>,

    /// State-dimension capacity the kernels were compiled with (the
    /// per-thread register-array size). The engine and trainer
    /// constructors derive it from the model config, so a mismatched
    /// launch cannot be built through the public constructors; the M1
    /// launch-path asserts additionally compare against it, and the
    /// kernels carry their own capacity guards.
    pub state_cap: usize,

    // ── Sequential SSM (mamba3_ssd.cu) ──
    pub m3_step_fwd: CudaFunction,
    pub m3_burnin_fwd: CudaFunction,
    pub m3_burnin_fwd_nosave: CudaFunction,
    pub m3_backward_seq: CudaFunction,
    pub m3_reduce_d_d: CudaFunction,

    // ── Shared ops (mamba3_ops.cu) ──
    pub m3_split: CudaFunction,
    pub m3_split_bwd: CudaFunction,
    pub bcnorm_fwd: CudaFunction,
    pub bcnorm_bwd: CudaFunction,
    pub bc_bias_add: CudaFunction,
    pub bc_bias_add_bwd: CudaFunction,
    pub angle_dt_fwd: CudaFunction,
    pub m3_angle_dt_fwd_batch: CudaFunction,
    pub m3_angle_dt_fwd_seq: CudaFunction,
    /// Chunk-parallel angle accumulation trio — replaces the sequential
    /// kernel on multi-chunk windows (its serial fp64 chain dominates
    /// the prefill profile at production shapes).
    pub m3_angle_chunk_sums: CudaFunction,
    pub m3_angle_chunk_carries: CudaFunction,
    pub m3_angle_chunk_apply: CudaFunction,
    pub angle_dt_bwd: CudaFunction,
    pub m3_angle_dt_bwd_seq: CudaFunction,
    pub rope_fwd: CudaFunction,
    /// Fused bc_bias_add (B + C) + rope_fwd - one launch replaces the
    /// F4c/F4d/F4ef triple; biased saves still materialize for backward.
    pub m3_bias_rope_fwd: CudaFunction,
    pub rope_bwd: CudaFunction,
    pub m3_compute_abg: CudaFunction,
    pub m3_abg_bwd: CudaFunction,
    pub silu_gate_fwd: CudaFunction,
    pub silu_gate_bwd: CudaFunction,
    pub rmsnorm_gated_fwd: CudaFunction,
    pub rmsnorm_gated_bwd: CudaFunction,

    // ── Shared kernels from norms.cu + elementwise.cu (used by training pipeline) ──
    pub rmsnorm_fwd: CudaFunction,
    pub rmsnorm_bwd: CudaFunction,
    pub colsum_accumulate: CudaFunction,
    /// Deterministic axis-0 reduction — finalises Rule-B per-sample partials
    /// into f32 master-grad slots. Replaces atomicAdd accumulators used in
    /// M3 backward (dD from m3_dqkv, d_angles_raw/d_dt_angle from
    /// m3_angle_dt_bwd_seq, and d_scale from rmsnorm_bwd).
    pub reduce_sum_axis0: CudaFunction,
    pub vec_add_inplace: CudaFunction,
    pub elementwise_mul: CudaFunction,
    pub fill_scalar: CudaFunction,
    pub cast_f32_to_bf16: CudaFunction,
    pub cast_f32_to_f16: CudaFunction,
    pub cast_bf16_to_f32: CudaFunction,
    pub cast_f16_to_f32: CudaFunction,
    pub residual_add: CudaFunction,
    pub gather_last_timestep: CudaFunction,

    // ── AdamW optimizer (adamw.cu) ──
    pub adamw_step_f32: CudaFunction,
    /// CUDA-Graph-capturable variant: bias factors read from a device
    /// buffer instead of scalar args.
    pub adamw_step_f32_capturable: CudaFunction,
    pub adamw_step_multi: TypedKernel,

    // ── Chunked parallel scan (mamba3_chunked.cu) ──
    pub m3_preprocess_chunks: CudaFunction,
    pub m3_da_cumsum: CudaFunction,
    pub m3_chunk_state_fwd: CudaFunction,
    pub m3_state_passing_fwd: CudaFunction,
    /// Trapezoidal boundary fold for state-carrying prefill: seeds the
    /// entering SSM state with `v_state (x) k_state * dt0 * (1 - trap0)`
    /// so the beta term's one-step reach-back across the window seam is
    /// honored (reference: mamba3_siso_fwd.py HAS_INITIAL_STATES).
    pub m3_chunk_entering_state: CudaFunction,
    pub m3_writeback_parallel_states: CudaFunction,
    pub m3_chunk_scan_fwd: CudaFunction,
    /// Cooperative-block twin: one head per 128-thread block, triangle
    /// tile + staged operands in dynamic smem. Routed when the smem total
    /// fits the 48 KB budget; bit-identical to the original.
    pub m3_chunk_scan_fwd_coop: CudaFunction,
    // m3_chunk_scan_bwd, m3_state_passing_bwd, m3_chunk_state_bwd,
    // m3_cumsum_bwd removed — dead code replaced by monolithic m3_dqkv +
    // m3_dqktheta path below.
    pub m3_extract_da_cs_sum: CudaFunction,
    pub m3_dqkv: CudaFunction,
    /// Per-chunk B terms of the reverse d_state recurrence (f32/typed
    /// activation reads; f32 output).
    pub m3_dqkv_state_terms_typed: TypedKernel,
    /// Serial per-(b,h,p,n) reverse fold producing each chunk's
    /// entering d_state.
    pub m3_dstate_passing_bwd: CudaFunction,
    pub m3_dqktheta: CudaFunction,
    pub m3_ddt_dtrap: CudaFunction,
    pub m3_final_grads: CudaFunction,

    // ── Typed variants for end-to-end bf16/f16 inference ──
    /// 8-way split + fused softplus/sigmoid, bf16/f16 proj and activation outputs.
    /// Coefficient outputs (dt, a_val, trap, angles, raw saves) stay f32.
    pub m3_split_typed: TypedKernel,
    /// RMSNorm on B/C per group, half I/O, f32 rms_val and weight.
    pub bcnorm_fwd_typed: TypedKernel,
    /// Fused B+C variant of bcnorm_fwd_typed (2× grid via blockIdx.y).
    pub bcnorm_fwd_bc_typed: TypedKernel,
    /// Fused B+C norm, f32 lane (the decode step merges its two bcnorm
    /// launches through it).
    pub bcnorm_fwd_bc_f32: CudaFunction,
    /// Per-head bias add, half I/O, f32 bias.
    pub bc_bias_add_typed: TypedKernel,
    /// Fused B+C variant of bc_bias_add_typed (2× grid via blockIdx.y).
    pub bc_bias_add_bc_typed: TypedKernel,
    /// RoPE rotation, half B/C, f32 angle_cumsum.
    pub rope_fwd_typed: TypedKernel,
    /// Fused typed twin of bias + rope (round-trip contract preserved).
    pub m3_bias_rope_fwd_typed: TypedKernel,
    /// Plain SiLU gate (no norm), half I/O.
    pub silu_gate_fwd_typed: TypedKernel,
    /// SiLU-gate backward (typed twins share the f32 kernel's argument
    /// order and factored d_silu form).
    pub silu_gate_bwd_typed: TypedKernel,
    /// RMSNorm-gated output (half I/O, f32 weight/rms_vals).
    pub rmsnorm_gated_fwd_typed: TypedKernel,
    /// M3 SSM step — shared with training, already templated in mamba3_ssd.cu.
    pub m3_step_fwd_typed: TypedKernel,
    /// M3 burnin forward (training) — sequential T-loop SSM with activation
    /// saves. Typed x/k/q/y; f32 state + alpha/beta/gamma + D + saves.
    pub m3_burnin_fwd_typed_bf16: CudaFunction,
    pub m3_burnin_fwd_typed_f16: CudaFunction,

    // -- Typed M3 sequential backward kernels --
    /// RmsNorm over B/C groups, typed dy → typed d_B; f32 rms + weight +
    /// d_weight master-grad accumulator.
    pub bcnorm_bwd_typed: TypedKernel,
    /// Per-head bias reduction from typed d_B_biased → typed d_B_normed
    /// (expand groups backward). Bias grad handled by reduce_bias_typed.
    pub bc_bias_add_bwd_typed: TypedKernel,
    /// RoPE rotation backward, typed B/C grads + saved B/C, f32 angle saves.
    pub rope_bwd_typed: TypedKernel,
    /// 8-way split backward: assemble typed d_proj from typed d_z/d_x/
    /// d_B_raw/d_C_raw plus f32 dd_dt/dd_a/trap/angles.
    pub m3_split_bwd_typed: TypedKernel,
    /// RMSNorm-gated backward (`out = RMSNorm(y) * weight * SiLU(z)`) —
    /// typed d_y/d_z/d_out/y/z; f32 weight + d_weight (per-sample,
    /// reduced later).
    pub rmsnorm_gated_bwd_typed: TypedKernel,
    // -- Typed M3 "final grad" kernels (HIGHEST RISK) --
    /// Huge dqkv kernel with smem tiles — typed Q_rot/K_scaled/V_in/dO;
    /// all 6 grad outputs stay f32 (atomicAdd on dD; master grads on others).
    pub m3_dqkv_typed: TypedKernel,
    /// Inverse-RoPE + bias backward — typed Q_raw/K_raw; 7 grad outputs f32.
    pub m3_dqktheta_typed: TypedKernel,

    // m3_chunk_scan_bwd_typed + m3_chunk_state_bwd_typed
    // removed (dead — the typed backward path, like the f32 path, now routes
    // through m3_dqkv_typed + m3_dqktheta_typed monolithic kernels).

    // -- Typed M3 chunked parallel forward kernels --
    /// Per-chunk gamma/scale + qk_dot + K prescale. Typed K/Q/K_scaled,
    /// f32 DT/trap_sig/qk_dot/scale/gamma (these are scalars computed in
    /// float and reused by the scan path).
    pub m3_preprocess_chunks_typed: TypedKernel,
    /// Per-chunk SSM state matmul (typed x, typed K_scaled, f32 dA_cumsum,
    /// f32 states_out — BPTT state MUST remain f32 per Tri Dao invariant).
    pub m3_chunk_state_fwd_typed: TypedKernel,
    /// Persist final states to ssm_state/k_state/v_state (all f32 persistent
    /// buffers). Typed inputs k_flat/x_flat.
    pub m3_writeback_parallel_states_typed: TypedKernel,
    /// Intra-chunk output: typed y_out/x/Q/K_scaled; f32 qk_dot/dA_cumsum/
    /// prev_states/D.
    pub m3_chunk_scan_fwd_typed: TypedKernel,
    /// Cooperative typed twin (see `m3_chunk_scan_fwd_coop`).
    pub m3_chunk_scan_fwd_coop_typed: TypedKernel,

    /// Shared from M1: f32 residual → half post-norm (identical kernel, reused).
    pub rmsnorm_fwd_f32in_typed: HalfKernel,
    /// Shared from M1: f32 residual += half branch (stays f32).
    pub residual_add_f32_typed: HalfKernel,
    /// Shared from M1: gather last timestep of B×T×D into B×D, dtype-preserving.
    pub gather_last_timestep_typed: TypedKernel,
}

impl Mamba3Kernels {
    /// Compile all 47 Mamba-3 CUDA kernels from source. Takes ~100-200ms.
    /// Compile with the default state capacity of 64. Models with a
    /// larger `d_state` use [`Self::compile_with_state_cap`].
    pub fn compile(ctx: &Arc<CudaContext>, arch: &'static str) -> Result<Self, String> {
        Self::compile_with_state_cap(ctx, arch, 64)
    }

    /// Compile all Mamba-3 kernels. `state_cap` sizes the per-thread
    /// state register arrays (see
    /// [`crate::mamba_ssm::gpu::kernels::state_capacity`]); raising it
    /// past the register budget makes the compiler spill to local
    /// memory — correct, slower, accepted as a first-class capacity
    /// knob rather than a separate slow path.
    pub fn compile_with_state_cap(
        ctx: &Arc<CudaContext>,
        arch: &'static str,
        state_cap: usize,
    ) -> Result<Self, String> {
        let sources = [
            // Inline the prelude first so each source file's
            // `#include "_typed_prelude.cuh"` can be safely stripped below.
            include_str!("../../../kernels/_typed_prelude.cuh"),
            include_str!("../../../kernels/mamba3_ssd.cu"),
            include_str!("../../../kernels/mamba3_ops.cu"),
            include_str!("../../../kernels/mamba3_chunked.cu"),
            // Shared kernels needed by training pipeline
            include_str!("../../../kernels/norms.cu"),
            include_str!("../../../kernels/elementwise.cu"),
            include_str!("../../../kernels/adamw.cu"),
        ];

        let combined: String = sources
            .iter()
            .map(|s| {
                s.lines()
                    .filter(|l| !l.trim().starts_with("#include \"_typed_prelude.cuh\""))
                    .collect::<Vec<_>>()
                    .join("\n")
            })
            .collect::<Vec<_>>()
            .join("\n");
        let option_strings = vec![
            "--fmad=true".to_string(),
            "--extra-device-vectorization".to_string(),
            // Mirrors the M1 compiler: strip device assert() trap
            // checks (none in the m3 sources today, but shared
            // headers may grow them); asserts compute no values.
            "-DNDEBUG".to_string(),
            format!("-DMAMBA_RS_STATE_CAP={state_cap}"),
        ];
        let opts = cudarc::nvrtc::CompileOptions {
            arch: Some(arch),
            options: option_strings.clone(),
            include_paths: crate::mamba_ssm::gpu::kernels::cuda_include_paths(),
            ..Default::default()
        };

        // PTX disk cache with the M1 loader's key law (source blob + arch
        // + options + NVRTC version): a hit skips the NVRTC half of the
        // boot tax; identical PTX by construction. Any hit-path failure
        // deletes the entry and falls through to a real compile; a
        // failed compile is never cached. The M3 loader paid a full
        // NVRTC compile of 7 sources on EVERY process boot before this.
        let (nv_major, nv_minor) = crate::mamba_ssm::gpu::kernels::nvrtc_version();
        // include_paths participate in the key (two-toolkit boxes).
        let key = crate::mamba_ssm::gpu::kernels::cache_key(&format!(
            "{combined}\u{1f}{arch}\u{1f}{option_strings:?}\u{1f}{:?}\u{1f}nvrtc{nv_major}.{nv_minor}",
            opts.include_paths
        ));
        let cache_path = crate::mamba_ssm::gpu::kernels::kernel_cache_dir()
            .map(|d| d.join(format!("mamba3-kernels-{key}.ptx")));

        // Hit path loads the module directly; a torn or rotted entry is
        // deleted and falls through to the real compile.
        let mut module = None;
        if let Some(path) = &cache_path
            && let Ok(src) = std::fs::read_to_string(path)
        {
            match ctx.load_module(cudarc::nvrtc::Ptx::from_src(src)) {
                Ok(m) => module = Some(m),
                Err(_) => {
                    let _ = std::fs::remove_file(path);
                }
            }
        }
        let module = match module {
            Some(m) => m,
            None => {
                let ptx = cudarc::nvrtc::compile_ptx_with_opts(combined, opts).map_err(|e| {
                    format!(
                        "NVRTC M3 compile failed: {}",
                        format!("{e:?}").replace("\\n", "\n")
                    )
                })?;
                if let Some(path) = &cache_path
                    && let Some(dir) = path.parent()
                    && std::fs::create_dir_all(dir).is_ok()
                {
                    // Atomic publish: write-then-rename so a concurrent
                    // boot never reads a torn entry; failures non-fatal.
                    let tmp = path.with_extension(format!("tmp-{}", std::process::id()));
                    if std::fs::write(&tmp, ptx.to_src()).is_ok()
                        && std::fs::rename(&tmp, path).is_err()
                    {
                        let _ = std::fs::remove_file(&tmp);
                    }
                }
                ctx.load_module(ptx)
                    .map_err(|e| format!("M3 module load failed: {e:?}"))?
            }
        };

        let get = |name: &str| -> Result<CudaFunction, String> {
            module
                .load_function(name)
                .map_err(|e| format!("M3 kernel '{name}' not found: {e:?}"))
        };

        let kernels = Self {
            state_cap,
            // Sequential SSM
            m3_step_fwd: get("m3_step_fwd")?,
            m3_burnin_fwd: get("m3_burnin_fwd")?,
            m3_burnin_fwd_nosave: get("m3_burnin_fwd_nosave")?,
            m3_backward_seq: get("m3_backward_seq")?,
            m3_reduce_d_d: get("m3_reduce_d_D")?,

            // Shared ops
            m3_split: get("m3_split")?,
            m3_split_bwd: get("m3_split_bwd")?,
            bcnorm_fwd: get("bcnorm_fwd")?,
            bcnorm_bwd: get("bcnorm_bwd")?,
            bc_bias_add: get("bc_bias_add")?,
            bc_bias_add_bwd: get("bc_bias_add_bwd")?,
            angle_dt_fwd: get("angle_dt_fwd")?,
            m3_angle_dt_fwd_batch: get("m3_angle_dt_fwd_batch")?,
            m3_angle_dt_fwd_seq: get("m3_angle_dt_fwd_seq")?,
            m3_angle_chunk_sums: get("m3_angle_chunk_sums")?,
            m3_angle_chunk_carries: get("m3_angle_chunk_carries")?,
            m3_angle_chunk_apply: get("m3_angle_chunk_apply")?,
            angle_dt_bwd: get("angle_dt_bwd")?,
            m3_angle_dt_bwd_seq: get("m3_angle_dt_bwd_seq")?,
            rope_fwd: get("rope_fwd")?,
            m3_bias_rope_fwd: get("m3_bias_rope_fwd")?,
            rope_bwd: get("rope_bwd")?,
            m3_compute_abg: get("m3_compute_abg")?,
            m3_abg_bwd: get("m3_abg_bwd")?,
            silu_gate_fwd: get("silu_gate_fwd")?,
            silu_gate_bwd: get("silu_gate_bwd")?,
            rmsnorm_gated_fwd: get("rmsnorm_gated_forward")?,
            rmsnorm_gated_bwd: get("rmsnorm_gated_backward")?,

            // Shared (norms.cu + elementwise.cu)
            rmsnorm_fwd: get("rmsnorm_forward")?,
            rmsnorm_bwd: get("rmsnorm_backward")?,
            colsum_accumulate: get("colsum_accumulate")?,
            reduce_sum_axis0: get("reduce_sum_axis0")?,
            vec_add_inplace: get("vec_add_inplace")?,
            elementwise_mul: get("elementwise_mul")?,
            fill_scalar: get("fill_scalar")?,
            cast_f32_to_bf16: get("cast_f32_to_bf16")?,
            cast_f32_to_f16: get("cast_f32_to_f16")?,
            cast_bf16_to_f32: get("cast_bf16_to_f32")?,
            cast_f16_to_f32: get("cast_f16_to_f32")?,
            residual_add: get("residual_add")?,
            gather_last_timestep: get("gather_last_timestep")?,

            // AdamW optimizer
            adamw_step_f32: get("adamw_step_f32")?,
            adamw_step_f32_capturable: get("adamw_step_f32_capturable")?,
            adamw_step_multi: TypedKernel {
                f32: get("adamw_step_multi_f32")?,
                bf16: get("adamw_step_multi_bf16")?,
                f16: get("adamw_step_multi_f16")?,
            },

            // Chunked parallel scan
            m3_preprocess_chunks: get("m3_preprocess_chunks")?,
            m3_da_cumsum: get("m3_dA_cumsum")?,
            m3_chunk_state_fwd: get("m3_chunk_state_fwd")?,
            m3_state_passing_fwd: get("m3_state_passing_fwd")?,
            m3_chunk_entering_state: get("m3_chunk_entering_state")?,
            m3_writeback_parallel_states: get("m3_writeback_parallel_states")?,
            m3_chunk_scan_fwd: get("m3_chunk_scan_fwd")?,
            m3_chunk_scan_fwd_coop: get("m3_chunk_scan_fwd_coop")?,
            m3_extract_da_cs_sum: get("m3_extract_da_cs_sum")?,
            m3_dqkv: get("m3_dqkv")?,
            m3_dqkv_state_terms_typed: TypedKernel {
                f32: get("m3_dqkv_state_terms")?,
                bf16: get("m3_dqkv_state_terms_bf16")?,
                f16: get("m3_dqkv_state_terms_f16")?,
            },
            m3_dstate_passing_bwd: get("m3_dstate_passing_bwd")?,
            m3_dqktheta: get("m3_dqktheta")?,
            m3_ddt_dtrap: get("m3_ddt_dtrap")?,
            m3_final_grads: get("m3_final_grads")?,

            // Typed variants for mixed-dtype inference
            m3_split_typed: TypedKernel {
                f32: get("m3_split")?,
                bf16: get("m3_split_bf16")?,
                f16: get("m3_split_f16")?,
            },
            bcnorm_fwd_typed: TypedKernel {
                f32: get("bcnorm_fwd")?,
                bf16: get("bcnorm_fwd_bf16")?,
                f16: get("bcnorm_fwd_f16")?,
            },
            bcnorm_fwd_bc_typed: TypedKernel {
                f32: get("bcnorm_fwd_bc_f32")?,
                bf16: get("bcnorm_fwd_bc_bf16")?,
                f16: get("bcnorm_fwd_bc_f16")?,
            },
            bcnorm_fwd_bc_f32: get("bcnorm_fwd_bc_f32")?,
            bc_bias_add_typed: TypedKernel {
                f32: get("bc_bias_add")?,
                bf16: get("bc_bias_add_bf16")?,
                f16: get("bc_bias_add_f16")?,
            },
            bc_bias_add_bc_typed: TypedKernel {
                f32: get("bc_bias_add")?,
                bf16: get("bc_bias_add_bc_bf16")?,
                f16: get("bc_bias_add_bc_f16")?,
            },
            rope_fwd_typed: TypedKernel {
                f32: get("rope_fwd")?,
                bf16: get("rope_fwd_bf16")?,
                f16: get("rope_fwd_f16")?,
            },
            m3_bias_rope_fwd_typed: TypedKernel {
                f32: get("m3_bias_rope_fwd")?,
                bf16: get("m3_bias_rope_fwd_bf16")?,
                f16: get("m3_bias_rope_fwd_f16")?,
            },
            silu_gate_fwd_typed: TypedKernel {
                f32: get("silu_gate_fwd")?,
                bf16: get("silu_gate_fwd_bf16")?,
                f16: get("silu_gate_fwd_f16")?,
            },
            silu_gate_bwd_typed: TypedKernel {
                f32: get("silu_gate_bwd")?,
                bf16: get("silu_gate_bwd_bf16")?,
                f16: get("silu_gate_bwd_f16")?,
            },
            rmsnorm_gated_fwd_typed: TypedKernel {
                f32: get("rmsnorm_gated_forward")?,
                bf16: get("rmsnorm_gated_forward_bf16")?,
                f16: get("rmsnorm_gated_forward_f16")?,
            },
            m3_step_fwd_typed: TypedKernel {
                f32: get("m3_step_fwd")?,
                bf16: get("m3_step_fwd_bf16")?,
                f16: get("m3_step_fwd_f16")?,
            },
            m3_burnin_fwd_typed_bf16: get("m3_burnin_fwd_bf16")?,
            m3_burnin_fwd_typed_f16: get("m3_burnin_fwd_f16")?,
            bcnorm_bwd_typed: TypedKernel {
                f32: get("bcnorm_bwd")?,
                bf16: get("bcnorm_bwd_bf16")?,
                f16: get("bcnorm_bwd_f16")?,
            },
            bc_bias_add_bwd_typed: TypedKernel {
                f32: get("bc_bias_add_bwd")?,
                bf16: get("bc_bias_add_bwd_bf16")?,
                f16: get("bc_bias_add_bwd_f16")?,
            },
            rope_bwd_typed: TypedKernel {
                f32: get("rope_bwd")?,
                bf16: get("rope_bwd_bf16")?,
                f16: get("rope_bwd_f16")?,
            },
            m3_split_bwd_typed: TypedKernel {
                f32: get("m3_split_bwd")?,
                bf16: get("m3_split_bwd_bf16")?,
                f16: get("m3_split_bwd_f16")?,
            },
            rmsnorm_gated_bwd_typed: TypedKernel {
                f32: get("rmsnorm_gated_backward")?,
                bf16: get("rmsnorm_gated_backward_bf16")?,
                f16: get("rmsnorm_gated_backward_f16")?,
            },
            m3_dqkv_typed: TypedKernel {
                f32: get("m3_dqkv")?,
                bf16: get("m3_dqkv_bf16")?,
                f16: get("m3_dqkv_f16")?,
            },
            m3_dqktheta_typed: TypedKernel {
                f32: get("m3_dqktheta")?,
                bf16: get("m3_dqktheta_bf16")?,
                f16: get("m3_dqktheta_f16")?,
            },

            m3_preprocess_chunks_typed: TypedKernel {
                f32: get("m3_preprocess_chunks")?,
                bf16: get("m3_preprocess_chunks_bf16")?,
                f16: get("m3_preprocess_chunks_f16")?,
            },
            m3_chunk_state_fwd_typed: TypedKernel {
                f32: get("m3_chunk_state_fwd")?,
                bf16: get("m3_chunk_state_fwd_bf16")?,
                f16: get("m3_chunk_state_fwd_f16")?,
            },
            m3_writeback_parallel_states_typed: TypedKernel {
                f32: get("m3_writeback_parallel_states")?,
                bf16: get("m3_writeback_parallel_states_bf16")?,
                f16: get("m3_writeback_parallel_states_f16")?,
            },
            m3_chunk_scan_fwd_coop_typed: TypedKernel {
                f32: get("m3_chunk_scan_fwd_coop")?,
                bf16: get("m3_chunk_scan_fwd_coop_bf16")?,
                f16: get("m3_chunk_scan_fwd_coop_f16")?,
            },
            m3_chunk_scan_fwd_typed: TypedKernel {
                f32: get("m3_chunk_scan_fwd")?,
                bf16: get("m3_chunk_scan_fwd_bf16")?,
                f16: get("m3_chunk_scan_fwd_f16")?,
            },

            rmsnorm_fwd_f32in_typed: HalfKernel {
                bf16: get("rmsnorm_forward_f32in_bf16")?,
                f16: get("rmsnorm_forward_f32in_f16")?,
            },
            residual_add_f32_typed: HalfKernel {
                bf16: get("residual_add_f32_bf16")?,
                f16: get("residual_add_f32_f16")?,
            },
            gather_last_timestep_typed: TypedKernel {
                f32: get("gather_last_timestep_f32")?,
                bf16: get("gather_last_timestep_bf16")?,
                f16: get("gather_last_timestep_f16")?,
            },

            _module: module,
        };

        // The chunked-backward kernel's dynamic shared memory grows
        // linearly with d_state (two chunk-by-d_state operand tiles
        // dominate) and exceeds the 48 KB default past d_state ~ 70.
        // Opt these functions in to the device's extended budget so a
        // larger state runs on the same code path. Best effort: on a
        // device without the budget the attribute call fails here and
        // an oversized launch later fails loudly with its own error —
        // never silently.
        // Unconditional: the chunked-backward tiles exceed the 48 KB
        // default from ordinary shapes too (e.g. d_state 64 with headdim
        // 32 needs ~67 KB), not only at raised state capacities.
        {
            use cudarc::driver::sys::CUfunction_attribute_enum as FnAttr;
            // 99 KB: consumer parts (sm_89/sm_120) cap the per-block opt-in
            // near 99-100 KB; the triangle-packed pair matrices fit
            // comfortably at CS=64 tile sizes.
            let budget: i32 = 99 * 1024;
            for f in [
                &kernels.m3_dqkv,
                &kernels.m3_dqkv_typed.f32,
                &kernels.m3_dqkv_typed.bf16,
                &kernels.m3_dqkv_typed.f16,
            ] {
                let _ = f.set_attribute(
                    FnAttr::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                    budget,
                );
            }
        }
        Ok(kernels)
    }
}

/// Launch geometry for the chunk-scan forward: the cooperative kernel (one
/// head per 128-thread block, triangle tile + staged operands in dynamic
/// smem) whenever its smem total fits the 48 KB default budget; wider
/// shapes keep the original two-head static-tile kernel. Returns
/// `(use_coop, cfg)` - the argument list is identical for both kernels.
pub(crate) fn chunk_scan_cfg(
    batch: usize,
    n_chunks: usize,
    nh: usize,
    hd: usize,
    ds: usize,
    chunk_size: usize,
) -> (bool, cudarc::driver::LaunchConfig) {
    let smem_floats = chunk_size * (chunk_size - 1) / 2
        + 2 * chunk_size * ds
        + chunk_size * hd
        + hd * ds
        + 2 * chunk_size;
    let smem_bytes = smem_floats * std::mem::size_of::<f32>();
    if chunk_size <= 64 && smem_bytes <= 48 * 1024 {
        (
            true,
            cudarc::driver::LaunchConfig {
                grid_dim: ((batch * n_chunks) as u32, nh as u32, 1),
                block_dim: (128, 1, 1),
                shared_mem_bytes: smem_bytes as u32,
            },
        )
    } else {
        (
            false,
            cudarc::driver::LaunchConfig {
                grid_dim: ((batch * n_chunks) as u32, nh.div_ceil(2) as u32, 1),
                block_dim: (hd as u32, 2, 1),
                shared_mem_bytes: 0,
            },
        )
    }
}