inferencelayer 0.2.13

//! CUDA decode engine for dense Qwen3 — the end-to-end arm of BEATING_VLLM.md.
//!
//! The kernel spikes proved both halves in isolation on an H100: cuBLAS `GemmEx` reaches
//! 558 TFLOP/s (prefill parity with vLLM) and the `q4_0` GEMV reaches 1829 GB/s = 91% of HBM peak
//! (a 986 tok/s decode ceiling vs vLLM's measured 155.9). This module turns the decode half into a
//! real forward pass so the comparison stops being a kernel number.
//!
//! Scope, deliberately: **dense Qwen3, decode path, greedy**. Prefill runs as sequential decode,
//! which is slow but irrelevant to the metric under test (decode tok/s) and keeps the KV cache
//! logic single-path. MoE (routed experts + shared-expert gate) is an extension of this once the
//! dense path is validated — not a prerequisite for proving the thesis.
//!
//! PORTABILITY: `#[cfg(feature = "cudarc")]` like every other CUDA file here, and cudarc dlopens
//! libcuda/libcublas/libnvrtc at runtime. A default build never compiles this; Metal/CPU/Vulkan
//! hosts are untouched and keep the WGSL engine.
use anyhow::{Context as _, Result};
use std::path::Path;
use std::sync::Arc;

use cudarc::cublas::{result as blas, sys as blas_sys, CudaBlas};
use cudarc::cublaslt::sys as lt_sys;
use cudarc::driver::{
    CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
    PushKernelArg,
};
use cudarc::nvrtc::{compile_ptx_with_opts, CompileOptions};
use half::f16;

use crate::weights::quantize_q4_0;

/// Dynamic shared bytes for `fa_prefill`: Q tile + DOUBLE-buffered K and V^T tiles.
/// (64*136 + 2*32*136 + 2*128*40) halves. Above 48 KB, so the function needs the opt-in below.
const FA_SHARED: u32 = ((64 * 136 + 2 * 32 * 136 + 2 * 128 * 40) * 2) as u32;

/// Dynamic shared bytes for `fa_prefill_hop`, the Hopper wgmma variant. The GMMA canonical layout
/// is unpadded (bank behaviour comes from the 128-byte core matrix, not from a pad), so the tiles
/// are exactly Q[2][64][128] + 2 x K[64][128] + 2 x V^T[128][64] halves, plus 128 B of alignment
/// slack because a matrix descriptor addresses shared memory in 16-byte units.
///
/// Q carries TWO query heads because the block does: see the GQA note on the kernel. That is the
/// whole point of the layout — the K and V tiles below are paid for once and used by both.
const FH_SHARED: u32 = ((2 * 64 * 128 + 2 * 64 * 128 + 2 * 128 * 64) * 2 + 128) as u32;

/// Dynamic shared bytes for `fa_prefill_ws` (HOP=2): Q[2][64][128] + 2 x K[128][128] +
/// 2 x V^T[128][128] halves. 160 KB — only reachable at one block per SM, which is exactly the
/// trade setmaxnreg pays for.
// +256: up to 126 B of alignment rounding for `base`, then 128 B holding the four mbarriers
/// before the tiles.
const FW_SHARED: u32 = ((2 * 64 * 128 + 2 * 128 * 128 + 2 * 128 * 128) * 2 + 256) as u32;

/// Everything the forward pass needs from `config.json`.
#[derive(Debug, Clone)]
pub struct Cfg {
    pub hidden: usize,
    pub layers: usize,
    pub n_heads: usize,
    pub n_kv: usize,
    pub head_dim: usize,
    pub intermediate: usize,
    pub eps: f32,
    pub rope_theta: f32,
    pub vocab: usize,
    pub tied: bool,
    /// MoE (Qwen3-30B-A3B): 0 => dense. Qwen3 MoE has NO shared expert (unlike Qwen2-MoE).
    pub num_experts: usize,
    pub topk: usize,
    pub moe_inter: usize,
    pub norm_topk: bool,
    /// Qwen3 applies per-head RMSNorm to q and k; Llama/Mistral do not. RMSNorm with weight 1 is
    /// NOT a no-op (it still divides by the RMS), so this has to be a real branch.
    pub qk_norm: bool,
    /// RoPE frequencies, precomputed on the host. Keeping the scaling maths here rather than in
    /// the kernel is what makes llama3 / linear scaling a config detail instead of a third kernel.
    pub inv_freq: Vec<f32>,
    // ---- Gemma 4 ----
    /// GeGLU: gelu_pytorch_tanh(gate) * up instead of silu(gate) * up.
    pub gelu: bool,
    /// Sandwich norms: x += rms(branch)·post_w (norm BEFORE the residual add), with a
    /// pre-norm on each branch input — four norms per layer, plus a whole-stream layer_scalar.
    pub sandwich: bool,
    /// Gemma scales the embedding output by sqrt(hidden). NOT folded into the table because the
    /// LM head is tied to the SAME table and must see unscaled rows.
    pub embed_scale: f32,
    /// Gemma4 text attention runs at scaling = 1.0 (q_norm bounds Q); everyone else 1/sqrt(hd).
    pub attn_scale_one: bool,
    /// Gemma4 applies a SCALELESS RMSNorm to V on every layer (v_norm has no weight).
    pub v_norm: bool,
    // ---- Qwen3.5/3.6 + Ornith: gated DeltaNet hybrid ----
    /// Per layer: true = linear_attention (gated DeltaNet), false = full attention. Empty for
    /// every non-hybrid model, which is what keeps the dense path a single branch.
    pub layer_linear: Vec<bool>,
    /// DeltaNet geometry: key/value head counts and dims, and the causal conv kernel width.
    pub dn_nk: usize,
    pub dn_nv: usize,
    pub dn_dk: usize,
    pub dn_dv: usize,
    pub dn_conv_kernel: usize,
    /// Qwen3.5/3.6 MoE carries a SHARED expert alongside the routed ones (Qwen3-30B-A3B does
    /// not). 0 = none.
    pub shared_expert_inter: usize,
    /// Partial rotary (Qwen3.5/3.6: 0.25): only the FIRST `rot_dim` dims of each head are
    /// rotated, the rest pass through. 0 = full rope over head_dim, which is everyone else.
    pub rot_dim: usize,
    /// Qwen3.5/3.6 gate the attention output: q_proj emits 2x and the second half is a
    /// sigmoid gate applied to the attention result before o_proj.
    pub attn_out_gate: bool,
    /// Per-layer attention geometry for layer-heterogeneous models (Gemma4: sliding layers are
    /// 16 kv-heads x hd 256, global layers 4 kv-heads x hd 512 with V sharing k_proj). Empty =
    /// uniform (head_dim/n_kv apply to every layer). head_dim/n_kv then hold the MAXES so every
    /// shared buffer is sized for the widest layer.
    pub layer_hd: Vec<usize>,
    pub layer_nkv: Vec<usize>,
    /// Sliding window per layer (0 = full attention).
    pub layer_window: Vec<usize>,
    /// Which rope table the layer uses: false = inv_freq, true = inv_freq2.
    pub layer_rope2: Vec<bool>,
    /// Second rope schedule (Gemma4 full-attention layers: "proportional" = the first
    /// partial_rotary_factor·hd/2 frequencies of the FULL-dim schedule, zero-padded — a zero
    /// frequency is cos=1/sin=0, i.e. identity on the unrotated pairs).
    pub inv_freq2: Vec<f32>,
    /// Gemma4 full-attention layers have NO v_proj — V is the raw k_proj output (then v_norm).
    /// Handled at LOAD time by fusing q|k|k so the buffer layout never changes; recorded here
    /// only so the loader knows which layers to duplicate.
    pub layer_v_from_k: Vec<bool>,
}

impl Cfg {
    fn from_json(root: &serde_json::Value) -> Result<Self> {
        // The `*ForConditionalGeneration` wrappers (Qwen3.5/3.6, Mistral3, Gemma 4) nest the
        // language model's dimensions under `text_config` and keep a vision tower beside it.
        // Reading the root would silently pick up the WRONG hidden size or miss the layer count.
        let v = root.get("text_config").unwrap_or(root);
        let g = |k: &str| -> Option<usize> { v.get(k)?.as_u64().map(|x| x as usize) };
        let hidden = g("hidden_size").context("hidden_size")?;
        let n_heads = g("num_attention_heads").context("num_attention_heads")?;
        let mut c = Cfg {
            hidden,
            layers: g("num_hidden_layers").context("num_hidden_layers")?,
            n_heads,
            n_kv: g("num_key_value_heads").unwrap_or(n_heads),
            head_dim: g("head_dim").unwrap_or(hidden / n_heads),
            // Pure-MoE checkpoints (Qwen3.6-35B-A3B, Ornith-35B) carry NO `intermediate_size` at
            // all — every FFN is expert-sized, so they define `moe_intermediate_size` only. Failing
            // here made those two report a config error instead of their real blocker, which is a
            // misleading refusal: the architecture check downstream is what should reject them.
            intermediate: g("intermediate_size")
                .or_else(|| g("moe_intermediate_size"))
                .context("neither intermediate_size nor moe_intermediate_size")?,
            eps: v.get("rms_norm_eps").and_then(|x| x.as_f64()).unwrap_or(1e-6) as f32,
            rope_theta: v.get("rope_theta").and_then(|x| x.as_f64()).unwrap_or(10000.0) as f32,
            vocab: g("vocab_size").context("vocab_size")?,
            tied: v
                .get("tie_word_embeddings")
                .and_then(|x| x.as_bool())
                .unwrap_or(false),
            num_experts: g("num_experts").unwrap_or(0),
            topk: g("num_experts_per_tok").unwrap_or(0),
            moe_inter: g("moe_intermediate_size").unwrap_or(0),
            norm_topk: v.get("norm_topk_prob").and_then(|x| x.as_bool()).unwrap_or(true),
            qk_norm: true,       // corrected in load() from the actual tensor list
            inv_freq: Vec::new(), // filled below; needs head_dim, hence a second pass
            gelu: false,
            sandwich: false,
            embed_scale: 1.0,
            layer_linear: Vec::new(),
            dn_nk: 0, dn_nv: 0, dn_dk: 0, dn_dv: 0, dn_conv_kernel: 0,
            shared_expert_inter: 0,
            rot_dim: 0, attn_out_gate: false,
            attn_scale_one: false,
            v_norm: false,
            layer_hd: Vec::new(),
            layer_nkv: Vec::new(),
            layer_window: Vec::new(),
            layer_rope2: Vec::new(),
            inv_freq2: Vec::new(),
            layer_v_from_k: Vec::new(),
        };
        if v.get("model_type").and_then(|x| x.as_str()) == Some("gemma4_text") {
            Self::gemma4(&mut c, v)?;
            return Ok(c);
        }
        c.inv_freq = Self::rope_inv_freq(&c, v.get("rope_scaling"))?;
        // ---- Qwen3.5/3.6 + Ornith: gated DeltaNet hybrid ----
        // layer_types marks which layers are the recurrent primitive. The DeltaNet dims live
        // under their own keys because they are INDEPENDENT of the full-attention geometry:
        // 16 key heads x 128 against the full layers' 16 heads x 256, in the same checkpoint.
        if let Some(lt) = v.get("layer_types").and_then(|x| x.as_array())
            && lt.iter().any(|k| k.as_str() == Some("linear_attention"))
        {
            c.layer_linear = lt.iter().map(|k| k.as_str() == Some("linear_attention")).collect();
            let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
            c.dn_nk = g("linear_num_key_heads");
            c.dn_nv = g("linear_num_value_heads");
            c.dn_dk = g("linear_key_head_dim");
            c.dn_dv = g("linear_value_head_dim");
            c.dn_conv_kernel = g("linear_conv_kernel_dim");
            anyhow::ensure!(c.dn_nk > 0 && c.dn_nv > 0 && c.dn_dk > 0 && c.dn_dv > 0
                            && c.dn_conv_kernel > 0,
                "layer_types says linear_attention but the linear_* dims are missing");
            anyhow::ensure!(c.dn_nv % c.dn_nk == 0,
                "DeltaNet value heads ({}) must be a multiple of key heads ({})", c.dn_nv, c.dn_nk);
        }
        c.shared_expert_inter = v.get("shared_expert_intermediate_size")
            .and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        // partial rotary: HF stores a FRACTION of head_dim; the kernels want the dim count, and
        // it must be even (rope rotates pairs).
        if let Some(pr) = v.get("partial_rotary_factor").and_then(|x| x.as_f64())
            && pr > 0.0 && pr < 1.0
        {
            c.rot_dim = (((c.head_dim as f64) * pr) as usize) & !1usize;
            anyhow::ensure!(c.rot_dim > 0, "partial_rotary_factor {pr} rounds to zero rope dims");
        }
        c.attn_out_gate = v.get("attn_output_gate").and_then(|x| x.as_bool()).unwrap_or(false);
        Ok(c)

    }

    /// Gemma 4 (dense): 5-sliding-1-global layer pattern with PER-TYPE geometry — sliding layers
    /// 16 kv-heads x hd 256 (theta 1e4, full rotary, window 1024), global layers 4 kv-heads x
    /// hd 512 (theta 1e6, "proportional" partial rotary 0.25, no v_proj: V = k_proj). Everything
    /// read from the checkpoint's transformers config, cross-checked against
    /// modeling_gemma4.py — the load-bearing surprises: RMSNorm is plain x̂·w (NOT Gemma-3's
    /// (1+w)), attention scaling is 1.0, V gets a scaleless RMSNorm on EVERY layer, and each
    /// layer multiplies the whole residual stream by its layer_scalar.
    fn gemma4(c: &mut Cfg, v: &serde_json::Value) -> Result<()> {
        let types = v.get("layer_types").and_then(|x| x.as_array()).context("layer_types")?;
        anyhow::ensure!(types.len() == c.layers, "layer_types len != num_hidden_layers");
        anyhow::ensure!(
            v.get("num_kv_shared_layers").and_then(|x| x.as_u64()).unwrap_or(0) == 0
                && !v.get("enable_moe_block").and_then(|x| x.as_bool()).unwrap_or(false)
                && v.get("hidden_size_per_layer_input").and_then(|x| x.as_u64()).unwrap_or(0) == 0,
            "gemma4 variant features (kv-shared / moe block / per-layer input) not implemented"
        );
        let hd_s = c.head_dim; // config head_dim = sliding layers
        let hd_g = v.get("global_head_dim").and_then(|x| x.as_u64()).map(|x| x as usize).unwrap_or(hd_s);
        let nkv_s = c.n_kv;
        let nkv_g = v.get("num_global_key_value_heads").and_then(|x| x.as_u64())
            .map(|x| x as usize).unwrap_or(nkv_s);
        let win = v.get("sliding_window").and_then(|x| x.as_u64()).context("sliding_window")? as usize;
        let k_eq_v = v.get("attention_k_eq_v").and_then(|x| x.as_bool()).unwrap_or(false);
        for t in types {
            let sliding = t.as_str() == Some("sliding_attention");
            c.layer_hd.push(if sliding { hd_s } else { hd_g });
            c.layer_nkv.push(if sliding { nkv_s } else { nkv_g });
            c.layer_window.push(if sliding { win } else { 0 });
            c.layer_rope2.push(!sliding);
            c.layer_v_from_k.push(!sliding && k_eq_v);
        }
        // head_dim/n_kv become the buffer-sizing MAXES; per-layer launches use the vecs.
        c.head_dim = hd_s.max(hd_g);
        c.n_kv = nkv_s.max(nkv_g);
        let rp = v.get("rope_parameters").context("rope_parameters")?;
        let theta = |k: &str| -> Result<f64> {
            rp.get(k).and_then(|x| x.get("rope_theta")).and_then(|x| x.as_f64())
                .with_context(|| format!("rope_parameters.{k}.rope_theta"))
        };
        // sliding: default rope over hd_s
        let th_s = theta("sliding_attention")?;
        c.inv_freq = (0..hd_s / 2)
            .map(|i| th_s.powf(-((2 * i) as f64) / hd_s as f64) as f32)
            .collect();
        // full: proportional — rope_angles real frequencies from the FULL hd_g schedule, then
        // zeros (identity rotation) out to hd_g/2. transformers
        // _compute_proportional_rope_parameters, verbatim.
        let fa = rp.get("full_attention").context("rope_parameters.full_attention")?;
        let th_g = theta("full_attention")?;
        let prf = fa.get("partial_rotary_factor").and_then(|x| x.as_f64()).unwrap_or(1.0);
        let factor = fa.get("factor").and_then(|x| x.as_f64()).unwrap_or(1.0);
        anyhow::ensure!(fa.get("rope_type").and_then(|x| x.as_str()) == Some("proportional")
            || prf == 1.0, "unexpected full_attention rope_type");
        let rot = ((prf * hd_g as f64) as usize) / 2;
        c.inv_freq2 = (0..hd_g / 2)
            .map(|i| if i < rot {
                (th_g.powf(-((2 * i) as f64) / hd_g as f64) / factor) as f32
            } else { 0.0 })
            .collect();
        c.gelu = v.get("hidden_activation").and_then(|x| x.as_str())
            .is_some_and(|a| a.contains("gelu"));
        c.sandwich = true;
        // HF multiplies by embed_scale CAST TO THE WEIGHT DTYPE (bf16): sqrt(5376) = 73.32
        // becomes 73.5. Match the rounding or every residual stream carries a 0.24% bias on its
        // embedding component relative to the reference.
        c.embed_scale = half::bf16::from_f32((c.hidden as f32).sqrt()).to_f32();
        c.attn_scale_one = true;
        c.v_norm = true;
        Ok(())
    }

    /// Per-layer attention geometry: (head_dim, n_kv, window, second-rope?).
    pub fn lgeom(&self, li: usize) -> (usize, usize, usize, bool) {
        if self.layer_hd.is_empty() {
            (self.head_dim, self.n_kv, 0, false)
        } else {
            (self.layer_hd[li], self.layer_nkv[li], self.layer_window[li], self.layer_rope2[li])
        }
    }

    /// max over layers of (nh + 2*nkv) * hd — the qkv scratch row width.
    pub fn max_qkv_width(&self) -> usize {
        // + n_heads again when the attention output gate is on: the gate rows ride the same
        // fused buffer, past q|k|v (see the loader).
        let g = if self.attn_out_gate { self.n_heads } else { 0 };
        (0..self.layers)
            .map(|li| { let (hd, nkv, _, _) = self.lgeom(li); (self.n_heads + 2 * nkv + g) * hd })
            .max()
            .unwrap_or((self.n_heads + 2 * self.n_kv) * self.head_dim)
    }

    /// theta^(-2i/hd), with HF's rope_scaling applied. Unknown scaling types are a hard error:
    /// silently ignoring one produces a model that loads, runs fast, and generates rubbish.
    fn rope_inv_freq(c: &Cfg, scaling: Option<&serde_json::Value>) -> Result<Vec<f32>> {
        let hd = c.head_dim;
        let mut f: Vec<f32> = (0..hd / 2)
            .map(|i| (c.rope_theta as f64).powf(-((2 * i) as f64) / hd as f64) as f32)
            .collect();
        let Some(sc) = scaling.filter(|x| !x.is_null()) else { return Ok(f) };
        let ty = sc.get("rope_type").or_else(|| sc.get("type"))
            .and_then(|x| x.as_str()).unwrap_or("default");
        let factor = sc.get("factor").and_then(|x| x.as_f64()).unwrap_or(1.0) as f32;
        match ty {
            "default" => {}
            "linear" => f.iter_mut().for_each(|v| *v /= factor),
            "llama3" => {
                let gf = |k: &str, d: f32| sc.get(k).and_then(|x| x.as_f64()).unwrap_or(d as f64) as f32;
                let (lo, hi) = (gf("low_freq_factor", 1.0), gf("high_freq_factor", 4.0));
                let orig = gf("original_max_position_embeddings", 8192.0);
                let (low_wl, high_wl) = (orig / lo, orig / hi);
                for v in f.iter_mut() {
                    let wl = 2.0 * std::f32::consts::PI / *v;
                    if wl > low_wl {
                        *v /= factor;
                    } else if wl >= high_wl {
                        let sm = ((orig / wl) - lo) / (hi - lo);
                        *v = (1.0 - sm) * (*v / factor) + sm * *v;
                    }
                }
            }
            other => anyhow::bail!("rope_scaling type {other:?} not implemented — refusing to \
                load a model whose positions would be silently wrong"),
        }
        Ok(f)
    }
}

/// A `q4_0` matrix resident on the device: `[rows, cols]`, 32-weight blocks, f16 scale + 4 u32.
impl Q4Dev {
    /// Alias the same device buffers. Needed only to satisfy the borrow checker when a method
    /// takes &mut self and also needs one of self's weights.
    fn clone_ref(&self) -> Q4Dev {
        Q4Dev { scales: self.scales.clone(), quants: self.quants.clone(),
                quants_p: self.quants_p.clone(), rows: self.rows, nblk: self.nblk }
    }
}

pub struct Q4Dev {
    scales: CudaSlice<f16>,
    quants: CudaSlice<u32>,
    /// Same weights in the GEMM's fragment-friendly permutation; built once, on demand.
    quants_p: Option<CudaSlice<u32>>,
    rows: usize,
    nblk: usize,
}

/// Gated-DeltaNet weights for ONE linear-attention layer (Qwen3.5/3.6, Ornith).
///
/// The four projections are kept as bf16 mirrors because the batched decode path multiplies
/// bf16 (the q4 arms are for the dense-attention weights); at 30 linear layers x ~34M weights
/// that is ~2 GB, against the MoE experts which must stay q4. `a`/`b` are [nv, hidden] — 32
/// rows — so they ride the same GEMM rather than earning their own kernel.
struct DnLayer {
    /// [conv_dim, hidden], conv_dim = 2*nk*dk + nv*dv — q, k and v fused on the row axis.
    qkv: CudaSlice<u16>,
    /// [nv*dv, hidden] — the z gate of the output RMSNorm.
    z: CudaSlice<u16>,
    /// [nv, hidden] each: the beta (b) and decay (a) projections.
    bp: CudaSlice<u16>,
    ap: CudaSlice<u16>,
    /// [hidden, nv*dv].
    out: CudaSlice<u16>,
    /// [conv_dim, kernel] depthwise causal conv, f32 (tiny).
    conv_w: CudaSlice<f32>,
    a_log: CudaSlice<f32>,
    dt_bias: CudaSlice<f32>,
    /// [dv], shared across heads.
    norm_w: CudaSlice<f32>,
}

struct LayerDev {
    attn_norm: CudaSlice<f32>,
    ffn_norm: CudaSlice<f32>,
    /// Sandwich norms (Gemma4): rms BEFORE each residual add. None on pre-norm-only models.
    post_attn_norm: Option<CudaSlice<f32>>,
    post_ffw_norm: Option<CudaSlice<f32>>,
    /// Whole-residual-stream scalar applied at layer end (Gemma4 layer_scalar; 1.0 elsewhere).
    ls: f32,
    /// This layer's attention geometry (uniform models: cfg.head_dim / cfg.n_kv / 0 / false).
    hd: usize,
    nkv: usize,
    window: usize,
    rope2: bool,
    /// q|k|v CONCATENATED on the row axis (q4_0 quantises row-wise, so this is exactly the three
    /// separate matrices) — one GEMV instead of three, and a 3x bigger kernel per launch.
    qkv: Q4Dev,
    o: Q4Dev,
    q_norm: CudaSlice<f32>,
    k_norm: CudaSlice<f32>,
    gate: Q4Dev,
    up: Q4Dev,
    down: Q4Dev,
    /// MoE: router [E, hidden] f32, and experts CONCATENATED — expert e owns rows
    /// [e*moe_inter, (e+1)*moe_inter) of w1/w3 and [e*hidden, (e+1)*hidden) of w2.
    router: Option<CudaSlice<f32>>,
    /// Shared expert (Qwen3.5/3.6): a dense SwiGLU MLP every token runs IN ADDITION to its
    /// routed experts, scaled by a sigmoid gate. `sw_gate` is [1, hidden].
    sw_gu: Option<Q4Dev>,
    sw_down: Option<Q4Dev>,
    sw_gate: Option<CudaSlice<f32>>,
    w1: Option<Q4Dev>,   // gate|up fused
    w2: Option<Q4Dev>,
    /// f16 mirrors, materialised once by `init_batch`. Batched decode is compute-bound, so it wants
    /// tensor cores; dequantising per step costs ~4.4 GB/step regardless of M (measured 5x slower),
    /// so the dequant happens ONCE and the GEMM reads f16 directly. Single-stream keeps q4_0.
    f16_qkv: Option<CudaSlice<u16>>,
    f16_o: Option<CudaSlice<u16>>,
    f16_gate: Option<CudaSlice<u16>>,
    /// gate and up dequantised into ONE [2*inter, hidden] buffer: one GEMM instead of two, and
    /// one activation conversion instead of two.
    f16_gu: Option<CudaSlice<u16>>,
    f16_up: Option<CudaSlice<u16>>,
    f16_down: Option<CudaSlice<u16>>,
}

/// Fixed FP8 activation scale. A per-step absmax would need a host readback inside the decode loop,
/// which serialises the stream and rules out graph capture; the e4m3 conversion saturates at +-448,
/// so a slightly tight scale clips outliers rather than corrupting the tensor.
const FP8_ACT_SCALE: f32 = 24.0;

/// cuBLASLt descriptors for one FP8 shape, kept alive for the lifetime of the engine.
struct Fp8Plan {
    cand: Vec<lt_sys::cublasLtMatmulAlgo_t>,
    op: lt_sys::cublasLtMatmulDesc_t,
    la: lt_sys::cublasLtMatrixLayout_t,
    lb: lt_sys::cublasLtMatrixLayout_t,
    ld: lt_sys::cublasLtMatrixLayout_t,
    algo: lt_sys::cublasLtMatmulAlgo_t,
}
unsafe impl Send for Fp8Plan {}

pub struct CudaQwen3 {
    _ctx: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    cfg: Cfg,
    embed: CudaSlice<f32>, // [vocab, hidden] — lookup source, kept f32 (read once per token)
    layers: Vec<LayerDev>,
    final_norm: CudaSlice<f32>,
    lm_head: Q4Dev,
    f16_lm: Option<CudaSlice<u16>>,
    // KV cache: [layers][max_seq][n_kv * head_dim]
    kc: Vec<CudaSlice<u16>>,
    vc: Vec<CudaSlice<u16>>,
    max_seq: usize,
    // scratch
    x: CudaSlice<f32>,
    h: CudaSlice<f32>,
    qkvbuf: CudaSlice<f32>, // [q | k | v]
    abuf: CudaSlice<f32>,
    gbuf: CudaSlice<f32>,
    ubuf: CudaSlice<f32>,
    // FlashDecoding partials: per (head, chunk) running max / sum / weighted-V
    pm: CudaSlice<f32>,
    pl: CudaSlice<f32>,
    pacc: CudaSlice<f32>,
    nsplit: usize,
    logits: CudaSlice<f32>,
    argm: CudaSlice<u32>,
    xs: CudaSlice<f32>, // [max|x|, 1/max|x|] for the GEMV currently in flight
    // MoE scratch
    rlog: CudaSlice<f32>,
    inv_freq: CudaSlice<f32>,
    /// Second rope schedule (Gemma4 global layers); 1-element dummy on uniform models.
    inv_freq2: CudaSlice<f32>,
    qkni: i32,   // router logits [E]
    eidx: CudaSlice<u32>,   // selected expert ids [k]
    ewt: CudaSlice<f32>,    // selected weights [k]
    exs: CudaSlice<f32>,    // per-slot (max,1/max) [2k]
    eg: CudaSlice<f32>,     // gate out [k, moe_inter]
    eu: CudaSlice<f32>,     // up out   [k, moe_inter]
    edn: CudaSlice<f32>,    // down out [k, hidden]
    // device-resident step state (so the whole decode step is capturable as ONE graph)
    d_token: CudaSlice<u32>,
    /// Per-slot copy of each prefill's sampled token, so a burst of admissions can replay all
    /// their prefill graphs back-to-back and sync ONCE instead of once per request.
    d_firsts: CudaSlice<u32>,
    /// Host->device staging for a whole admission burst's prompt ids: one synchronizing H2D for
    /// the burst instead of one per request.
    d_stage: CudaSlice<u32>,
    f_slot_ctl: CudaFunction,
    f_set_i32: CudaFunction,
    /// Gemma4: scaleless V-head RMSNorm (in the f32 qkv scratch, pre-rope).
    f_vderive: CudaFunction,
    /// Gemma4 sandwich: x = (x + rms(h)·postw)·ls; out16 = rms(x)·prew.
    f_sandwich: CudaFunction,
    /// Wide-M q4 GEMM on wgmma (raw q4 traffic, bf16 smem operand).
    f_q4wg: CudaFunction,
    /// Nibble repack for the wgmma GEMM's shift-free f16x2 extraction.
    f_q4rep_wg: CudaFunction,
    /// BN=64 twin of the wgmma GEMM, for decode widths (grid fills the machine).
    f_q4wg64: CudaFunction,
    /// Deterministic split-K partial reduction for the wgmma GEMM.
    f_gwksum: CudaFunction,
    /// Marlin-structure q4 GEMM: register-resident B dequant, 4-stage cp.async pipeline.
    /// Three tile widths; the host picks by batch width so the grid fills the machine.
    f_q4ml: CudaFunction,
    f_q4ml64: CudaFunction,
    f_q4ml32: CudaFunction,
    /// f16-emitting activation producers for the Marlin decode arm.
    f_rms_hf: CudaFunction,
    f_attn_red_h16: CudaFunction,
    f_silu_hf: CudaFunction,
    /// MARLIN=1: decode multiplies raw q4 weights through gemm_q4_ml instead of bf16 mirrors.
    ml_arm: bool,
    /// Split-K partial buffer, grown on demand.
    gwpart: Option<CudaSlice<f32>>,
    /// Split-K partials for the Marlin kernel.
    mlpart: Option<CudaSlice<f32>>,
    d_pos: CudaSlice<i32>,
    /// Prefill KV write offset, device-resident so one graph per length serves all slots.
    d_kvoff: CudaSlice<i32>,
    /// Per-segment KV offsets for wave prefill (one entry per admitted request in the wave).
    d_kvoffs: CudaSlice<i32>,
    /// Wave-prefill tail scratch: gathered last rows, their logits, their argmaxes.
    wave_x: CudaSlice<f32>,
    wave_logits: CudaSlice<f32>,
    wave_first: CudaSlice<u32>,
    /// score-wave sandwich tail: gathered per-segment last rows of x and h + the bf16 out rows
    wave_h: CudaSlice<f32>,
    wave_b16: CudaSlice<u16>,
    f_attn_wave: CudaFunction,
    d_out: CudaSlice<u32>,
    graph: Option<cudarc::driver::CudaGraph>,
    /// One graph per batch width — the grid dimensions are baked into a capture, so a narrower
    /// batch needs its own. Widths are powers of two, so this stays small.
    bgraphs: std::collections::HashMap<(usize, bool), cudarc::driver::CudaGraph>,
    /// (token count, graph) — a prefill graph is only valid for the length it was captured at.
    pfgraph: Option<(usize, cudarc::driver::CudaGraph)>,
    /// Prefill graphs for the SERVER path, keyed by (prompt length, slot). The bench prefills via a
    /// captured graph while cuda-serve launched ~280 kernels per request — worth ~6 ms, which is
    /// both the prefill-through-HTTP gap and the TTFT gap.
    pfslot_graphs: std::collections::HashMap<usize, cudarc::driver::CudaGraph>,
    // kernels
    f_gemv: CudaFunction,
    f_rms: CudaFunction,
    f_rms_add: CudaFunction,
    f_qkrope: CudaFunction,
    f_attn_part: CudaFunction,
    f_attn_red: CudaFunction,
    f_silu: CudaFunction,
    f_add: CudaFunction,
    f_embed: CudaFunction,
    f_argmax: CudaFunction,
    f_absmax: CudaFunction,
    f_gemv_f32: CudaFunction,
    f_topk: CudaFunction,
    f_topk_b: CudaFunction,
    f_combine_b: CudaFunction,
    f_bfr: CudaFunction,
    f_moe_count: CudaFunction,
    f_moe_scan: CudaFunction,
    f_moe_scatter: CudaFunction,
    f_gemv_moe_g: CudaFunction,
    mcnt: CudaSlice<i32>,
    moff: CudaSlice<i32>,
    mcur: CudaSlice<i32>,
    mplist: CudaSlice<u32>,
    /// bf16 router copies per layer (batched router GEMM); built by init_batch for MoE models
    router16: Vec<CudaSlice<u16>>,
    /// batched-MoE scratch, grown to `moe_rows` rows (max of batch width and prefill T)
    moe_rows: usize,
    mrlog: CudaSlice<f32>,
    meidx: CudaSlice<u32>,
    mewt: CudaSlice<f32>,
    meg: CudaSlice<f32>,
    meu: CudaSlice<f32>,
    medn: CudaSlice<f32>,
    mexs: CudaSlice<f32>,
    f_absmax_slots: CudaFunction,
    f_gemv_moe: CudaFunction,
    f_combine: CudaFunction,
    f_silu_moe: CudaFunction,
    f_kv_write: CudaFunction,
    f_advance: CudaFunction,
    // batched (concurrency) path
    f_gemv_b: CudaFunction,
    f_embed_b: CudaFunction,
    f_rms_add_b: CudaFunction,
    f_rms_bf: CudaFunction,
    f_rope_b: CudaFunction,
    /// fp8-KV twins (FP8KV=1): same math, e4m3-byte slabs.
    f_rope_b_f8: CudaFunction,
    f_rope_p_f8: CudaFunction,
    f_attn_part_b_f8: CudaFunction,
    /// e4m3 KV slabs: half the decode step's dominant traffic. Forces the plain partial_b
    /// attention arm and the batch path (seq/flash/grouped/chunk/mix are bf16-slab readers).
    fp8kv: bool,
    /// Block target behind attn_partial_b's Y grid (ATTNB_TGT, default 1024).
    attnb_tgt: usize,
    /// ATTNW=1: warp-private decode attention arm (no block barriers in the chunk loop).
    attnw: bool,
    f_attn_part_w: CudaFunction,
    /// Per layer: Some(..) on linear-attention layers, None on full-attention ones.
    dn: Vec<Option<DnLayer>>,
    f_attn_out_gate: CudaFunction,
    f_shared_expert: CudaFunction,
    f_add_inplace: CudaFunction,
    f_dn_conv: CudaFunction,
    f_dn_step: CudaFunction,
    /// MOET=1: tensor-core grouped MoE expert GEMM (replaces the per-pair GEMVs).
    moet: bool,
    /// MOET_PF=0 keeps PREFILL on the per-pair GEMV while decode uses the tensor arm —
    /// the discriminator for which of the two paths a numeric regression lives in.
    moet_pf: bool,
    f_moe_t: CudaFunction,
    f_moe_t64: CudaFunction,
    f_moe_th2: CudaFunction,
    f_moe_t64h2: CudaFunction,
    f_moe_t64h4: CudaFunction,
    f_moe_th4: CudaFunction,
    f_moe_wl: CudaFunction,
    f_moe_sort: CudaFunction,
    /// compact (expert, pair-tile) work list + its count
    mwork: CudaSlice<u32>,
    mwn: CudaSlice<i32>,
    /// ATTNT: tensor-core (mma) decode attention arm — one warp per 64-key chunk.
    /// "1" forces it, "0" disables; UNSET = AUTO: on when the live context reaches ATTNT_CTX
    /// (default 512). Measured: +14% at m=8/2k and +1.2% at M=64 short, but -1.7% at M=32
    /// short (143 KB blocks, mostly-idle warps at 3 chunks) — hence the context switch.
    attnt: bool,
    attnt_auto: bool,
    attnt_ctx: usize,
    /// Max context across the CURRENT live set, provided by the scheduler each iteration —
    /// unlike ctx_hint it decays when long sequences retire, so the graph key (width, arm)
    /// flips back to the scalar arm for short work.
    ctx_live: usize,
    f_attn_part_t: CudaFunction,
    f_attn_part_b: CudaFunction,
    f_attn_part_g: CudaFunction,
    f_attn_part_f: CudaFunction,
    f_attn_strip: CudaFunction,
    f_kv_gather: CudaFunction,
    f_attn_red_b: CudaFunction,
    f_attn_red_b16: CudaFunction,
    f_argmax_b: CudaFunction,
    f_advance_b: CudaFunction,
    f_embed_p: CudaFunction,
    f_rope_p: CudaFunction,
    f_softmax_c: CudaFunction,
    f_scale_out: CudaFunction,
    f_silu_bf: CudaFunction,
    f_silu_gu: CudaFunction,
    f_fa: CudaFunction,
    f_fa_hop: CudaFunction,
    f_fa_ws: CudaFunction,
    f_q4gemm: CudaFunction,
    f_q4v2: [CudaFunction; 4],
    f_q4red: CudaFunction,
    f_q4rep: CudaFunction,
    f_tofp8: CudaFunction,
    f_absmax_bf: CudaFunction,
    /// producer-side absmax variants + the read-and-reset scale kernel (decode fp8 chains)
    f_rms_bf_ax: CudaFunction,
    f_silu_bf_ax: CudaFunction,
    f_attn_red_b16_ax: CudaFunction,
    f_amax2s: CudaFunction,
    fp8_amax: CudaSlice<f32>,
    /// fastest Lt algo per (m, rows, cols) — timed once per shape, shared by every layer
    /// Algo cache keyed (m-or-t, rows, cols, delayed_scale) — the dw descriptor restricts
    /// which algos execute, so the two configs must never share a cached winner.
    fp8_algo_cache: std::collections::HashMap<(usize, usize, usize, bool), lt_sys::cublasLtMatmulAlgo_t>,
    /// FP8DW auto-disable for this run: set when any decode shape has no executable
    /// delayed-scale algo (the step degrades to the amax-fold chains).
    fp8dw_dead: bool,
    /// delayed-scale slots [layers][4 sites][amax, amax/448, next-amax]; site 0=qkv 1=o 2=gu 3=down
    fp8_dsc: CudaSlice<f32>,
    /// tuned bf16 decode plans, keyed (m, rows, cols) — shapes repeat across layers
    bf16_plans: std::collections::HashMap<(usize, usize, usize), Fp8Plan>,
    /// longest live context, host side. The decode attention arm is chosen from it: KV traffic
    /// is ~7% of a short-context step but ~35% at 2k, and the flash kernel only pays there.
    ctx_hint: usize,
    f_rms_fp8: CudaFunction,
    f_silu_fp8: CudaFunction,
    f_attn_red_fp8: CudaFunction,
    f_sc_rot: CudaFunction,
    /// f32 flavours for the score-path attention output (p.attn is f32 there)
    f_amax_f_part: CudaFunction,
    f_tofp8_f: CudaFunction,
    f_mkscale: CudaFunction,
    f_amax_part: CudaFunction,
    f_amax_fin: CudaFunction,
    fp8_part: CudaSlice<f32>,
    /// prefill FP8: plans are width-specific, so prefill needs its own set keyed on T
    fp8_pplans: Vec<Fp8Plan>,
    fp8_pplans_qkv: Vec<Fp8Plan>,
    fp8_pplans_o: Vec<Fp8Plan>,
    fp8_pplans_down: Vec<Fp8Plan>,
    fp8_pplan_t: usize,
    fp8_px: CudaSlice<u8>,
    fp8_psx: Vec<CudaSlice<f32>>,
    fp8_psmax: Vec<CudaSlice<f32>>,
    /// FP8 e4m3 mirror of gate|up per layer, with its per-tensor scale
    fp8_gu: Vec<(CudaSlice<u8>, CudaSlice<f32>)>,
    /// FP8 e4m3 mirrors of the REST of the step's weight bytes (qkv / o / down) — same
    /// per-tensor-scale scheme as fp8_gu; together they halve decode weight traffic vs bf16
    fp8_qkv: Vec<(CudaSlice<u8>, CudaSlice<f32>)>,
    fp8_o: Vec<(CudaSlice<u8>, CudaSlice<f32>)>,
    fp8_down: Vec<(CudaSlice<u8>, CudaSlice<f32>)>,
    fp8_plans: Vec<Fp8Plan>,
    fp8_plans_qkv: Vec<Fp8Plan>,
    fp8_plans_o: Vec<Fp8Plan>,
    fp8_plans_down: Vec<Fp8Plan>,
    /// batch width the plans were built for — layouts are width-specific
    fp8_plan_m: usize,
    fp8_sxmax: CudaSlice<f32>,
    fp8_plan: Option<Fp8Plan>,
    fp8_x: CudaSlice<u8>,
    fp8_sx: CudaSlice<f32>,
    /// Permuted q4 copies of the five projections per layer, for the q4 decode arm.
    q4p: Vec<[CudaSlice<u32>; 5]>,
    /// Merged gate|up: permuted quants and concatenated scales, so one GEMM serves both.
    q4gu: Vec<(CudaSlice<u32>, CudaSlice<f16>)>,
    /// unused-pointer placeholders for the rope kernel's two cache flavours
    dummy_f32: CudaSlice<f32>,
    dummy_f32b: CudaSlice<f32>,
    dummy_u16: CudaSlice<u16>,
    dummy_u16b: CudaSlice<u16>,
    /// Runtime toggle for the q4 decode arm, so both arms can be diffed in one process.
    q4_on: bool,
    /// Q4FREE=1 dropped the original q4 quants — every path that multiplies them is dead.
    q4free: bool,
    q4part: Option<CudaSlice<f32>>,
    f_vt: CudaFunction,
    /// Flash-attention prefill. On by default; FA=0 falls back to the cuBLAS score-matrix path,
    /// which is what the two are diffed against.
    fa: bool,
    /// cc 9.x: use the wgmma prefill attention kernel instead of the mma.sync one.
    hop: bool,
    /// HOP=2: the warp-specialised (setmaxnreg + producer warpgroup) variant.
    ws: bool,
    /// Work-queue counter for the persistent-block ws kernel; zeroed before each launch.
    fw_ticket: CudaSlice<u32>,
    /// Device SM count — the persistent ws kernel launches exactly this many blocks.
    sm_count: i32,
    pf: Option<Prefill>,
    batch: Option<Batch>,
    /// Per-slot staged prompt ids for mixed prefill/decode steps. Staged ONCE at admission so a
    /// mixed step's chunk read is a device slice, not a pageable htod that would drain the stream
    /// mid-pipeline (the same reason d_stage exists for admission bursts).
    mixtok: std::collections::HashMap<usize, CudaSlice<u32>>,
    /// OVERLAP stream: long prefills replay their captured graph here while the decode graph
    /// free-runs on the main stream. Decoders keep full step cadence through a prefill instead
    /// of one token per chunk pass — and the prefill keeps its graph and plans (no eager cost).
    stream2: Arc<CudaStream>,
    /// Prefill graphs captured ON stream2 (a CudaGraph replays on its capture stream).
    pfslot_graphs2: std::collections::HashMap<usize, cudarc::driver::CudaGraph>,
    /// Shadow scratch swapped in during stream2 work: the prefill tail writes xf16/h/xs and the
    /// Lt workspace, all of which the concurrent decode step also touches. Disjoint buffers or
    /// the two streams race on them.
    xf16b: CudaSlice<u16>,
    hb2: CudaSlice<f32>,
    xsb: CudaSlice<f32>,
    ltws2: CudaSlice<u8>,
    blas: CudaBlas,
    /// cuBLASLt handle — FP8 matmul is Lt-only
    lt: lt_sys::cublasLtHandle_t,
    ltws: CudaSlice<u8>,
    /// f16 scratch for the dequantised weight of the matrix currently in flight, and for the f16
    /// activation operand. Sized to the largest matrix in the model.
    wf16: CudaSlice<u16>,
    xf16: CudaSlice<u16>,
    yf32: CudaSlice<f32>,
    f_dq: CudaFunction,
    f_tobf16: CudaFunction,
    f_tof16: CudaFunction,
}

/// Prefill working set for up to `tmax` tokens processed together.
struct Prefill {
    tmax: usize,
    x: CudaSlice<f32>,
    h: CudaSlice<f32>,
    qkv: CudaSlice<f32>,
    qkv16: CudaSlice<f16>,
    vt16: CudaSlice<f16>,
    score16: CudaSlice<f16>,
    rowsum: CudaSlice<f32>,
    attn: CudaSlice<f32>,
    g: CudaSlice<f32>,
    u: CudaSlice<f32>,
    hb16: CudaSlice<u16>,
    attn16: CudaSlice<u16>,
    /// attention output as bf16, straight from the reduce
    abuf16: CudaSlice<u16>,
    gu16: CudaSlice<u16>,
    /// merged gate|up result, f32, for the q4 arm
    gu32: CudaSlice<f32>,
    g16: CudaSlice<u16>,
    logits: CudaSlice<f32>,
    tok: CudaSlice<u32>,
}

/// Buffers for M sequences decoding in lockstep — static batching, which is what steady-state
/// concurrency looks like for a throughput measurement (arrivals/departures, i.e. true continuous
/// batching, are a scheduler on top of exactly these kernels).
pub struct Batch {
    /// Current width the kernels run at. Shrinks to the active set so a server with 64 slots
    /// does not do 64 sequences' work for one request.
    pub m: usize,
    /// Width the buffers were allocated for; `m` may be anything up to this.
    mmax: usize,
    max_seq: usize,
    x: CudaSlice<f32>,
    h: CudaSlice<f32>,
    qkv: CudaSlice<f32>,
    abuf: CudaSlice<f32>,
    g: CudaSlice<f32>,
    u: CudaSlice<f32>,
    hb16: CudaSlice<u16>,
    attn16: CudaSlice<u16>,
    /// attention output as bf16, straight from the reduce
    abuf16: CudaSlice<u16>,
    gu16: CudaSlice<u16>,
    /// merged gate|up result, f32, for the q4 arm
    gu32: CudaSlice<f32>,
    g16: CudaSlice<u16>,
    logits: CudaSlice<f32>,
    xs: CudaSlice<f32>,
    kc: Vec<CudaSlice<u16>>,
    vc: Vec<CudaSlice<u16>>,
    pm: CudaSlice<f32>,
    pl: CudaSlice<f32>,
    pacc: CudaSlice<f32>,
    tok: CudaSlice<u32>,
    out: CudaSlice<u32>,
    /// Per-slot position — continuous batching means slots sit at different context lengths.
    pos: CudaSlice<i32>,
    /// Per-slot occupancy; a zero freezes that slot while the rest of the batch steps.
    active: CudaSlice<u32>,
    nsplit: usize,
    /// Gated-DeltaNet per-slot state, indexed by LAYER (empty slot for full-attention layers):
    /// `dn_s` is [m, nv, dk, dv] and `dn_ring` [m, conv_dim, kernel]. This is the whole reason a
    /// linear-attention model is memory-hungry per slot — ~2.1 MB per slot per linear layer.
    dn_s: Vec<CudaSlice<f32>>,
    dn_ring: Vec<CudaSlice<f32>>,
    /// DeltaNet per-step scratch (shared across layers).
    dn_mixed: CudaSlice<f32>,
    dn_z: CudaSlice<f32>,
    dn_b: CudaSlice<f32>,
    dn_a: CudaSlice<f32>,
    dn_core: CudaSlice<f32>,
    dn_core16: CudaSlice<u16>,
    /// Host mirror of occupancy, so admission does not need a device readback.
    pub occupied: Vec<bool>,
}

fn ecfg(len: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (len.div_ceil(256) as u32, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    }
}

impl CudaQwen3 {
    /// Quantize a host f32 matrix to `q4_0` and upload. Scales are stored f16 (the format the
    /// GEMV binds — 18 B per 32 weights instead of 20).
    fn up_q4(stream: &Arc<CudaStream>, w: &[f32], rows: usize, cols: usize) -> Result<Q4Dev> {
        let (sc, qz) = quantize_q4_0(w, rows, cols);
        let sc16: Vec<f16> = sc.iter().map(|&s| f16::from_f32(s)).collect();
        Ok(Q4Dev {
            scales: stream.memcpy_stod(&sc16)?,
            quants: stream.memcpy_stod(&qz)?,
            quants_p: None,
            rows,
            nblk: cols / 32,
        })
    }

    pub fn load(dir: &Path, max_seq: usize) -> Result<Self> {
        let cfg_json: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json")).context("config.json")?)?;
        let mut cfg = Cfg::from_json(&cfg_json)?;
        // Refuse unknown architectures loudly. Every family below is "RMSNorm + SwiGLU + NeoX RoPE"
        // optionally with q/k-norm and MoE — except Gemma4, which gets its own forward shape
        // (sandwich norms, GeGLU, per-layer-type geometry, sliding windows, layer_scalar) behind
        // cfg.sandwich. Gemma 2/3 ((1+w) norms) remain deliberately absent.
        const SUPPORTED: [&str; 15] = [
            "Qwen3ForCausalLM", "Qwen3MoeForCausalLM", "Qwen2ForCausalLM",
            "LlamaForCausalLM", "MistralForCausalLM",
            // multimodal wrappers whose TEXT stack is plain RMSNorm+SwiGLU+RoPE; benchmarked
            // text-only, with the vision tower simply not loaded
            "Mistral3ForConditionalGeneration", "LlamaForConditionalGeneration",
            "Qwen2_5_VLForConditionalGeneration",
            "Gemma4ForConditionalGeneration",
            // Gated-DeltaNet hybrids (Qwen3.5/3.6, Ornith): 30 of 40 layers are the recurrence,
            // the rest full attention with partial rope + an output gate, over a 256-expert MoE
            // with a shared expert. Every one of those is implemented; the recurrence is gated
            // against the CPU reference at 7e-7 (output) and 3e-8 (state).
            "Qwen3_5MoeForCausalLM", "Qwen3_5MoeForConditionalGeneration",
            "Qwen3_6MoeForCausalLM", "Qwen3_6MoeForConditionalGeneration",
            "OrnithForCausalLM", "OrnithForConditionalGeneration",
        ];
        // Qwen3.5/3.6 (and Ornith, same architecture) are deliberately NOT here. `layer_types`
        // shows 48 of 64 layers are `linear_attention`: gated DeltaNet with a conv1d of kernel 4
        // and recurrent state, plus attn_output_gate and partial_rotary_factor 0.25 on the 16 full
        // attention layers. That is a second attention mechanism, not a config flag, so the loader
        // refuses rather than running three quarters of the model as something it is not.
        let tcfg = cfg_json.get("text_config").unwrap_or(&cfg_json);
        if let Some(lt) = tcfg.get("layer_types").and_then(|x| x.as_array()) {
            let exotic: Vec<&str> = lt
                .iter()
                .filter_map(|x| x.as_str())
                .filter(|k| *k != "full_attention")
                // sliding_attention is implemented for the sandwich (Gemma4) forward pass
                .filter(|k| !(cfg.sandwich && *k == "sliding_attention"))
                // linear_attention is the gated-DeltaNet arm: kernels gated against the CPU
                // reference (output 7e-7, state 3e-8), weights loaded per linear layer, state
                // carried per slot. It is no longer "a second attention mechanism we refuse".
                .filter(|k| !(!cfg.layer_linear.is_empty() && *k == "linear_attention"))
                .collect();
            if !exotic.is_empty() {
                anyhow::bail!(
                    "checkpoint interleaves {} non-full-attention layers ({:?}); this engine \
                     implements full attention only",
                    exotic.len(), exotic.first()
                );
            }
        }
        if let Some(arch) = cfg_json.get("architectures").and_then(|a| a.get(0)).and_then(|a| a.as_str()) {
            anyhow::ensure!(SUPPORTED.contains(&arch),
                "architecture {arch:?} not implemented by the CUDA engine (supported: {SUPPORTED:?})");
        }
        if !cfg.sandwich
            && cfg_json.get("sliding_window").and_then(|x| x.as_u64()).is_some_and(|w| w > 0)
            && cfg_json.get("use_sliding_window").and_then(|x| x.as_bool()).unwrap_or(true)
        {
            anyhow::bail!("model uses sliding-window attention, which this engine does not implement");
        }

        let ctx = CudaContext::new(0).context("cuda device 0")?;
        // cudarc tracks every slice usage with events to auto-synchronise ACROSS streams. During graph
        // capture those waits reference uncaptured work and abort the capture with
        // STREAM_CAPTURE_ISOLATION. This engine runs everything on ONE stream, so program order
        // already guarantees ordering and the events are pure overhead. Must be disabled BEFORE
        // any allocation: tracking is attached to a slice at creation.
        unsafe { ctx.disable_event_tracking() };
        // A NON-default stream: CUDA forbids graph capture on the legacy default stream
        // (CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED), and the captured graph is the whole point.
        let stream = ctx.new_stream().context("create stream")?;
        let st = crate::weights::LazySt::open(dir).context("open safetensors")?;
        let t = |n: &str| -> Result<Vec<f32>> { st.tensor_f32(n) };

        // Multimodal wrappers put the text stack under `model.language_model.`; text-only models
        // use `model.`. Probe instead of assuming — a wrong guess is a confusing missing-tensor error.
        let pfx = if st.tensor_f32("model.layers.0.input_layernorm.weight").is_ok() {
            "model"
        } else if st.tensor_f32("model.language_model.layers.0.input_layernorm.weight").is_ok() {
            "model.language_model"
        } else if st.tensor_f32("language_model.model.layers.0.input_layernorm.weight").is_ok() {
            "language_model.model"
        } else {
            anyhow::bail!("cannot locate the decoder layers in this checkpoint");
        };
        let embed_h = t(&format!("{pfx}.embed_tokens.weight"))?;
        let embed = stream.memcpy_stod(&embed_h)?;

        // Qwen3 has q_norm/k_norm; Llama and Mistral do not. Presence of the tensor decides.
        cfg.qk_norm = st.tensor_f32(&format!("{pfx}.layers.0.self_attn.q_norm.weight")).is_ok();
        let nh = cfg.n_heads;
        let ones = vec![1f32; cfg.head_dim];
        let mut layers = Vec::with_capacity(cfg.layers);
        // ---- gated DeltaNet weights, one set per linear-attention layer ----
        // Loaded BEFORE the attention layers so a missing tensor fails fast, and as bf16
        // mirrors (see DnLayer) — the recurrence itself is f32 and reads its state, not these.
        let mut dn: Vec<Option<DnLayer>> = Vec::with_capacity(cfg.layers);
        if !cfg.layer_linear.is_empty() {
            let (nk, nv, dk, dv, kn) =
                (cfg.dn_nk, cfg.dn_nv, cfg.dn_dk, cfg.dn_dv, cfg.dn_conv_kernel);
            let conv_dim = 2 * nk * dk + nv * dv;
            let bf = |v: &[f32]| -> Vec<u16> {
                v.iter().map(|x| (x.to_bits() >> 16) as u16).collect()
            };
            for i in 0..cfg.layers {
                if !cfg.layer_linear[i] { dn.push(None); continue; }
                let q = format!("{pfx}.layers.{i}.linear_attn");
                let qkv_h = t(&format!("{q}.in_proj_qkv.weight"))?;
                anyhow::ensure!(qkv_h.len() == conv_dim * cfg.hidden,
                    "layer {i} in_proj_qkv is {} elements, expected conv_dim*hidden = {}",
                    qkv_h.len(), conv_dim * cfg.hidden);
                let z_h = t(&format!("{q}.in_proj_z.weight"))?;
                let b_h = t(&format!("{q}.in_proj_b.weight"))?;
                let a_h = t(&format!("{q}.in_proj_a.weight"))?;
                let o_h = t(&format!("{q}.out_proj.weight"))?;
                // conv1d is [conv_dim, 1, kernel] in the checkpoint; the middle axis is a
                // depthwise singleton, so the flat buffer is already [conv_dim, kernel].
                let cw_h = t(&format!("{q}.conv1d.weight"))?;
                anyhow::ensure!(cw_h.len() == conv_dim * kn,
                    "layer {i} conv1d is {} elements, expected conv_dim*kernel = {}",
                    cw_h.len(), conv_dim * kn);
                dn.push(Some(DnLayer {
                    qkv: stream.memcpy_stod(&bf(&qkv_h))?,
                    z: stream.memcpy_stod(&bf(&z_h))?,
                    bp: stream.memcpy_stod(&bf(&b_h))?,
                    ap: stream.memcpy_stod(&bf(&a_h))?,
                    out: stream.memcpy_stod(&bf(&o_h))?,
                    conv_w: stream.memcpy_stod(&cw_h)?,
                    a_log: stream.memcpy_stod(&t(&format!("{q}.A_log"))?)?,
                    dt_bias: stream.memcpy_stod(&t(&format!("{q}.dt_bias"))?)?,
                    norm_w: stream.memcpy_stod(&t(&format!("{q}.norm.weight"))?)?,
                }));
            }
            eprintln!("gated DeltaNet: {} linear layers loaded (nk {nk} nv {nv} dk {dk} dv {dv} \
                       conv {kn}, conv_dim {conv_dim})",
                      cfg.layer_linear.iter().filter(|x| **x).count());
        }
        for i in 0..cfg.layers {
            let p = format!("{pfx}.layers.{i}");
            // Per-layer geometry (Gemma4 sliding vs global); uniform models get cfg's values.
            let (hd, nkv, window, rope2) = cfg.lgeom(i);
            let v_from_k = cfg.layer_v_from_k.get(i).copied().unwrap_or(false);
            // Gemma names the pre-FFN norm `pre_feedforward_layernorm`; on Qwen-family models
            // `post_attention_layernorm` IS the pre-FFN norm (there are only two).
            let ffn_norm_name = if cfg.sandwich {
                format!("{p}.pre_feedforward_layernorm.weight")
            } else {
                format!("{p}.post_attention_layernorm.weight")
            };
            // A linear-attention layer has no self_attn block at all -- its q/k/v/o live under
            // linear_attn and are already loaded into DnLayer. Give the attention fields the same
            // 32-float dummy the MoE path gives the dense FFN fields, so LayerDev stays uniform
            // and the step dispatch (which sends these layers to deltanet_step) never reads them.
            let is_linear = cfg.layer_linear.get(i).copied().unwrap_or(false);
            layers.push(LayerDev {
                attn_norm: stream.memcpy_stod(&t(&format!("{p}.input_layernorm.weight"))?)?,
                ffn_norm: stream.memcpy_stod(&t(&ffn_norm_name)?)?,
                post_attn_norm: if cfg.sandwich {
                    Some(stream.memcpy_stod(&t(&format!("{p}.post_attention_layernorm.weight"))?)?)
                } else { None },
                post_ffw_norm: if cfg.sandwich {
                    Some(stream.memcpy_stod(&t(&format!("{p}.post_feedforward_layernorm.weight"))?)?)
                } else { None },
                ls: if cfg.sandwich {
                    *t(&format!("{p}.layer_scalar"))?.first().context("layer_scalar empty")?
                } else { 1.0 },
                hd, nkv, window, rope2,
                qkv: if is_linear { Self::up_q4(&stream, &[0f32; 32], 1, 32)? } else {
                    let mut c = t(&format!("{p}.self_attn.q_proj.weight"))?;
                    // ATTENTION OUTPUT GATE (Qwen3.5/3.6): q_proj emits 2*nh*hd — the first
                    // half is q, the second a gate applied to the attention result. Split it
                    // off here and append it AFTER v, so the fused buffer is [q|k|v|gate]:
                    // rope and attention launch with per = nh + 2*nkv and never touch the gate
                    // rows, which keeps both kernels unchanged.
                    let gate_rows = if c.len() == 2 * nh * hd * cfg.hidden {
                        anyhow::ensure!(cfg.attn_out_gate,
                            "layer {i} q_proj is 2x wide but attn_output_gate is not set");
                        Some(c.split_off(nh * hd * cfg.hidden))
                    } else { None };
                    let k = t(&format!("{p}.self_attn.k_proj.weight"))?;
                    c.extend_from_slice(&k);
                    if v_from_k {
                        // No v_proj on Gemma4 global layers — V is the raw k_proj output. Fusing
                        // q|k|K (the k rows AGAIN) keeps the [q|k|v] buffer layout identical for
                        // every layer at the cost of duplicated k rows in the GEMM (~11% of this
                        // layer's qkv flops); v_derive then norms the copy in place.
                        c.extend_from_slice(&k);
                    } else {
                        c.extend_from_slice(&t(&format!("{p}.self_attn.v_proj.weight"))?);
                    }
                    let mut qkv_rows = (nh + 2 * nkv) * hd;
                    if let Some(g) = gate_rows { c.extend_from_slice(&g); qkv_rows += nh * hd; }
                    Self::up_q4(&stream, &c, qkv_rows, cfg.hidden)?
                },
                o: if is_linear { Self::up_q4(&stream, &[0f32; 32], 1, 32)? } else { Self::up_q4(
                    &stream,
                    &t(&format!("{p}.self_attn.o_proj.weight"))?,
                    cfg.hidden,
                    nh * hd,
                )? },
                q_norm: if cfg.qk_norm && !is_linear {
                    stream.memcpy_stod(&t(&format!("{p}.self_attn.q_norm.weight"))?)?
                } else { stream.memcpy_stod(&ones)? },
                k_norm: if cfg.qk_norm && !is_linear {
                    stream.memcpy_stod(&t(&format!("{p}.self_attn.k_norm.weight"))?)?
                } else { stream.memcpy_stod(&ones)? },
                gate: if cfg.num_experts > 0 { Self::up_q4(&stream, &[0f32; 32], 1, 32)? } else { Self::up_q4(
                    &stream,
                    &t(&format!("{p}.mlp.gate_proj.weight"))?,
                    cfg.intermediate,
                    cfg.hidden,
                )? },
                up: if cfg.num_experts > 0 { Self::up_q4(&stream, &[0f32; 32], 1, 32)? } else { Self::up_q4(
                    &stream,
                    &t(&format!("{p}.mlp.up_proj.weight"))?,
                    cfg.intermediate,
                    cfg.hidden,
                )? },
                down: if cfg.num_experts > 0 { Self::up_q4(&stream, &[0f32; 32], 1, 32)? } else { Self::up_q4(
                    &stream,
                    &t(&format!("{p}.mlp.down_proj.weight"))?,
                    cfg.hidden,
                    cfg.intermediate,
                )? },
                router: None, w1: None, w2: None,
                sw_gu: None, sw_down: None, sw_gate: None,
                f16_qkv: None, f16_o: None, f16_gate: None, f16_gu: None, f16_up: None, f16_down: None,
            });
        }
        // MoE experts, if present. Concatenate per-expert q4 blocks so one GEMV can address any
        // expert by row offset (mirrors the wgpu engine's layout).
        if cfg.num_experts > 0 {
            let (e_n, mi) = (cfg.num_experts, cfg.moe_inter);
            for i in 0..cfg.layers {
                let p = format!("{pfx}.layers.{i}.mlp");
                // gate|up interleaved per expert: expert e owns rows [e*2mi, e*2mi+mi) = gate and
                // the next mi = up, so ONE GEMV produces both (q4_0 is row-wise, so this is exact).
                let mut s1 = Vec::new(); let mut q1 = Vec::new();
                let mut s2 = Vec::new(); let mut q2 = Vec::new();
                // FUSED EXPERT TENSORS (Qwen3.5/3.6): one [E, 2*mi, hidden] blob instead of
                // 2*E named matrices. Flattened it is ALREADY the row order the per-expert loop
                // below builds, so it quantises in one call — and that matters here: 256 experts
                // x 40 layers is 20k tensor lookups on the split layout.
                if let Ok(gu) = t(&format!("{p}.experts.gate_up_proj")) {
                    anyhow::ensure!(gu.len() == e_n * 2 * mi * cfg.hidden,
                        "fused gate_up_proj is {} elements, expected E*2*mi*hidden = {}",
                        gu.len(), e_n * 2 * mi * cfg.hidden);
                    let (a1, b1) = quantize_q4_0(&gu, e_n * 2 * mi, cfg.hidden);
                    s1.extend(a1.iter().map(|&v| f16::from_f32(v))); q1.extend(b1);
                    drop(gu);
                    let dn_t = t(&format!("{p}.experts.down_proj"))?;
                    anyhow::ensure!(dn_t.len() == e_n * cfg.hidden * mi,
                        "fused down_proj is {} elements, expected E*hidden*mi = {}",
                        dn_t.len(), e_n * cfg.hidden * mi);
                    let (a2, b2) = quantize_q4_0(&dn_t, e_n * cfg.hidden, mi);
                    s2.extend(a2.iter().map(|&v| f16::from_f32(v))); q2.extend(b2);
                    layers[i].router = Some(stream.memcpy_stod(&t(&format!("{p}.gate.weight"))?)?);
                    layers[i].w1 = Some(Q4Dev { scales: stream.memcpy_stod(&s1)?, quants: stream.memcpy_stod(&q1)?, quants_p: None, rows: e_n * 2 * mi, nblk: cfg.hidden / 32 });
                    layers[i].w2 = Some(Q4Dev { scales: stream.memcpy_stod(&s2)?, quants: stream.memcpy_stod(&q2)?, quants_p: None, rows: e_n * cfg.hidden, nblk: mi / 32 });
                    if cfg.shared_expert_inter > 0 {
                        let si = cfg.shared_expert_inter;
                        let mut gu = t(&format!("{p}.shared_expert.gate_proj.weight"))?;
                        gu.extend_from_slice(&t(&format!("{p}.shared_expert.up_proj.weight"))?);
                        layers[i].sw_gu = Some(Self::up_q4(&stream, &gu, 2 * si, cfg.hidden)?);
                        layers[i].sw_down = Some(Self::up_q4(&stream,
                            &t(&format!("{p}.shared_expert.down_proj.weight"))?, cfg.hidden, si)?);
                        layers[i].sw_gate = Some(stream.memcpy_stod(
                            &t(&format!("{p}.shared_expert_gate.weight"))?)?);
                    }
                    if i % 8 == 0 { eprintln!("  moe layer {i}/{} quantized (fused)", cfg.layers); }
                    continue;
                }
                for e in 0..e_n {
                    let (a, b) = quantize_q4_0(&t(&format!("{p}.experts.{e}.gate_proj.weight"))?, mi, cfg.hidden);
                    s1.extend(a.iter().map(|&v| f16::from_f32(v))); q1.extend(b);
                    let (a, b) = quantize_q4_0(&t(&format!("{p}.experts.{e}.up_proj.weight"))?, mi, cfg.hidden);
                    s1.extend(a.iter().map(|&v| f16::from_f32(v))); q1.extend(b);
                    let (a, b) = quantize_q4_0(&t(&format!("{p}.experts.{e}.down_proj.weight"))?, cfg.hidden, mi);
                    s2.extend(a.iter().map(|&v| f16::from_f32(v))); q2.extend(b);
                }
                layers[i].router = Some(stream.memcpy_stod(&t(&format!("{p}.gate.weight"))?)?);
                layers[i].w1 = Some(Q4Dev { scales: stream.memcpy_stod(&s1)?, quants: stream.memcpy_stod(&q1)?, quants_p: None, rows: e_n * 2 * mi, nblk: cfg.hidden / 32 });
                layers[i].w2 = Some(Q4Dev { scales: stream.memcpy_stod(&s2)?, quants: stream.memcpy_stod(&q2)?, quants_p: None, rows: e_n * cfg.hidden, nblk: mi / 32 });
                if i % 8 == 0 { eprintln!("  moe layer {i}/{} quantized", cfg.layers); }
            }
        }

        let final_norm = stream.memcpy_stod(&t(&format!("{pfx}.norm.weight"))?)?;
        let lm_head = if cfg.tied {
            Self::up_q4(&stream, &embed_h, cfg.vocab, cfg.hidden)?
        } else {
            // Three placements seen in the wild: top level (`lm_head.weight`), inside the decoder
            // (`{pfx}.lm_head.weight`), and BESIDE it — Mistral 3 has the decoder at
            // `language_model.model.layers.*` but the head at `language_model.lm_head.weight`.
            let w = match t("lm_head.weight") {
                Ok(w) => w,
                Err(_) => match t(&format!("{pfx}.lm_head.weight")) {
                    Ok(w) => w,
                    Err(_) => {
                        let stem = pfx.rsplit_once('.').map(|(a, _)| a).unwrap_or(pfx);
                        t(&format!("{stem}.lm_head.weight"))
                            .with_context(|| format!("no lm_head near {pfx}"))?
                    }
                },
            };
            Self::up_q4(&stream, &w, cfg.vocab, cfg.hidden)?
        };

        // Compile for the ACTUAL device, not a floor. compute_70 was the floor that keeps the
        // V100S working, but mma.sync.m16n8k16 — the whole basis of fa_prefill — only exists from
        // sm_80, and asking for it at compute_70 fails PTX JIT rather than degrading.
        let (ccmaj, ccmin) = {
            use cudarc::driver::sys::CUdevice_attribute as A;
            (ctx.attribute(A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)?,
             ctx.attribute(A::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)?)
        };
        let archs: &'static str = match (ccmaj, ccmin) {
            (7, 0) => "compute_70",
            (7, _) => "compute_75",
            (8, 0) => "compute_80",
            (8, 6) => "compute_86",
            (8, _) => "compute_89",
            // 90a, not 90: wgmma (fa_prefill_hop) is an ARCH-ACCELERATED instruction and ptxas
            // rejects every one of them under a plain .target sm_90. The `a` variants are not
            // forward-compatible, which is fine here — this arm only ever runs on a cc 9.x device.
            (9, _) => "compute_90a",
            (m, _) if m >= 10 => "compute_100",
            _ => "compute_70",   // anything older than Volta cannot run this engine anyway
        };
        let ptx = compile_ptx_with_opts(
            KERNELS,
            CompileOptions { arch: Some(archs), ..Default::default() },
        )
        .map_err(|e| anyhow::anyhow!("nvrtc ({archs}): {e:?}"))?;
        let m = ctx.load_module(ptx).context("load ptx")?;
        let f = |n: &str| m.load_function(n).context(n.to_string());

        // maxes for the shared scratch; per-layer slabs use each layer's own kvdim.
        // LAZY: the single-stream KV slabs are dead weight on the server (the seq path reads
        // batch slot 0 since the prompt-blindness fix) — ~2 GB on Gemma-31B that helped OOM the
        // boot. grow_seq_kv() sizes them on first single-stream use (the bench path).
        let (hd, nkv) = (cfg.head_dim, cfg.n_kv);
        let mut kc = Vec::with_capacity(cfg.layers);
        let mut vc = Vec::with_capacity(cfg.layers);
        for _ in 0..cfg.layers {
            kc.push(stream.alloc_zeros::<u16>(1)?);
            vc.push(stream.alloc_zeros::<u16>(1)?);
        }

        // widest matrix decides the f16 scratch: MoE gate|up if present, else the LM head
        let widest = if cfg.num_experts > 0 {
            (cfg.num_experts * 2 * cfg.moe_inter) * cfg.hidden
        } else {
            cfg.vocab * cfg.hidden
        };
        let blas = CudaBlas::new(stream.clone()).context("cublas")?;
        let mut lt: lt_sys::cublasLtHandle_t = std::ptr::null_mut();
        unsafe {
            let st = lt_sys::cublasLtCreate(&mut lt);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "cublasLtCreate: {st:?}");
        }
        let ltws = stream.alloc_zeros::<u8>(64 << 20)?;

        Ok(Self {
            blas,
            lt,
            ltws,
            stream2: ctx.new_stream()?,
            pfslot_graphs2: std::collections::HashMap::new(),
            xf16b: stream.alloc_zeros::<u16>(8192)?,
            hb2: stream.alloc_zeros::<f32>(cfg.hidden)?,
            xsb: stream.alloc_zeros::<f32>(2 * 8192)?,
            ltws2: stream.alloc_zeros::<u8>(64 << 20)?,
            // LAZY (vocab*hidden u16 = 2.8 GB on Gemma-31B); gemm_tc grows it on first use —
            // the batched server never calls gemm_tc, so it never pays.
            wf16: stream.alloc_zeros::<u16>(1)?,
            xf16: stream.alloc_zeros::<u16>(8192 * 1)?,
            yf32: stream.alloc_zeros::<f32>(64 * widest / cfg.hidden.max(1))?,
            f_dq: f("dequant_q4_bf16")?,
            f_tobf16: f("f32_to_bf16")?,
            f_tof16: f("f32_to_f16")?,
            x: stream.alloc_zeros::<f32>(cfg.hidden)?,
            h: stream.alloc_zeros::<f32>(cfg.hidden)?,
            qkvbuf: stream.alloc_zeros::<f32>(cfg.max_qkv_width())?,
            abuf: stream.alloc_zeros::<f32>(nh * hd)?,
            gbuf: stream.alloc_zeros::<f32>(cfg.intermediate)?,
            ubuf: stream.alloc_zeros::<f32>(cfg.intermediate)?,
            pm: stream.alloc_zeros::<f32>(nh * max_seq.div_ceil(64))?,
            pl: stream.alloc_zeros::<f32>(nh * max_seq.div_ceil(64))?,
            pacc: stream.alloc_zeros::<f32>(nh * max_seq.div_ceil(64) * hd)?,
            nsplit: max_seq.div_ceil(64),
            logits: stream.alloc_zeros::<f32>(cfg.vocab)?,
            argm: stream.alloc_zeros::<u32>(1)?,
            // 2 floats (max, 1/max) PER ROW: the GEMV uses row 0, but the f16 GEMM path scales
            // per row and prefill runs M = prompt length, so this must cover the widest batch.
            xs: stream.alloc_zeros::<f32>(2 * 8192)?,
            rlog: stream.alloc_zeros::<f32>(cfg.num_experts.max(1))?,
            inv_freq: stream.memcpy_stod(&cfg.inv_freq)?,
            inv_freq2: if cfg.inv_freq2.is_empty() {
                stream.alloc_zeros::<f32>(1)?
            } else {
                stream.memcpy_stod(&cfg.inv_freq2)?
            },
            qkni: cfg.qk_norm as i32,
            eidx: stream.alloc_zeros::<u32>(cfg.topk.max(1))?,
            ewt: stream.alloc_zeros::<f32>(cfg.topk.max(1))?,
            exs: stream.alloc_zeros::<f32>(2 * cfg.topk.max(1))?,
            eg: stream.alloc_zeros::<f32>(cfg.topk.max(1) * 2 * cfg.moe_inter.max(1))?,
            eu: stream.alloc_zeros::<f32>(cfg.topk.max(1) * cfg.moe_inter.max(1))?,
            edn: stream.alloc_zeros::<f32>(cfg.topk.max(1) * cfg.hidden)?,
            d_token: stream.alloc_zeros::<u32>(1)?,
            d_firsts: stream.alloc_zeros::<u32>(256)?,
            // Sized ONCE for the worst-case admission burst and never reallocated. The grow-on-
            // demand version was the only post-boot device reallocation in the engine, and its
            // realloc correlated exactly with the prefill-graph replays writing garbage over
            // p.tok (clean when bursts stayed under the initial capacity, corrupt after a growth
            // realloc between a graph's capture and its replay). Whatever the precise allocator/
            // graph interaction, removing the realloc removes the trigger.
            d_stage: stream.alloc_zeros::<u32>(64 * 4096)?,
            f_slot_ctl: f("slot_ctl")?,
            f_set_i32: f("set_i32")?,
            d_pos: stream.alloc_zeros::<i32>(1)?,
            d_kvoff: stream.alloc_zeros::<i32>(1)?,
            d_kvoffs: stream.alloc_zeros::<i32>(64)?,
            wave_x: stream.alloc_zeros::<f32>(1)?,
            wave_logits: stream.alloc_zeros::<f32>(1)?,
            wave_first: stream.alloc_zeros::<u32>(1)?,
            wave_h: stream.alloc_zeros::<f32>(1)?,
            wave_b16: stream.alloc_zeros::<u16>(1)?,
            f_attn_wave: { let ff = f("attn_wave")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 84000)
                    .context("opt in to >48 KB dynamic shared for attn_wave")?; } ff },
            d_out: stream.alloc_zeros::<u32>(max_seq)?,
            mixtok: std::collections::HashMap::new(),
            graph: None,
            pfgraph: None,
            pfslot_graphs: std::collections::HashMap::new(),
            bgraphs: std::collections::HashMap::new(),
            f_gemv: f("gemv_q4")?,
            f_rms: f("rmsnorm")?,
            f_rms_add: f("rmsnorm_add")?,
            f_qkrope: f("qk_norm_rope")?,
            f_attn_part: f("attn_partial")?,
            f_attn_red: f("attn_reduce")?,
            f_silu: f("silu_mul")?,
            f_add: f("add_inplace")?,
            f_embed: f("embed_lookup")?,
            f_argmax: f("argmax_f32")?,
            f_absmax: f("absmax")?,
            f_gemv_f32: f("gemv_f32")?,
            f_topk: f("moe_topk")?,
            f_topk_b: f("moe_topk_b")?,
            f_combine_b: f("moe_combine_b")?,
            f_bfr: f("bf16_round_f32")?,
            f_moe_count: f("moe_count")?,
            f_moe_scan: f("moe_scan")?,
            f_moe_scatter: f("moe_scatter")?,
            f_gemv_moe_g: f("gemv_q4_moe_g")?,
            mcnt: stream.alloc_zeros::<i32>(1)?,
            moff: stream.alloc_zeros::<i32>(1)?,
            mcur: stream.alloc_zeros::<i32>(1)?,
            mplist: stream.alloc_zeros::<u32>(1)?,
            router16: Vec::new(),
            moe_rows: 0,
            mrlog: stream.alloc_zeros::<f32>(1)?,
            meidx: stream.alloc_zeros::<u32>(1)?,
            mewt: stream.alloc_zeros::<f32>(1)?,
            meg: stream.alloc_zeros::<f32>(1)?,
            meu: stream.alloc_zeros::<f32>(1)?,
            medn: stream.alloc_zeros::<f32>(1)?,
            mexs: stream.alloc_zeros::<f32>(1)?,
            f_absmax_slots: f("absmax_slots")?,
            f_gemv_moe: f("gemv_q4_moe")?,
            f_combine: f("moe_combine")?,
            f_silu_moe: f("silu_mul_moe")?,
            f_kv_write: f("kv_write")?,
            f_advance: f("advance")?,
            f_gemv_b: f("gemv_q4_batch")?,
            f_embed_b: f("embed_batch")?,
            f_rms_add_b: f("rmsnorm_add_b")?,
            f_rms_bf: f("rmsnorm_add_bf")?,
            f_rope_b: f("qk_rope_b")?,
            f_rope_b_f8: f("qk_rope_b_f8")?,
            f_rope_p_f8: f("rope_prefill_f8")?,
            f_attn_part_b_f8: f("attn_partial_b_f8")?,
            // e4m3 KV needs the uniform hd=128 fast paths; Gemma (sandwich, per-layer hd) and
            // any non-128 model keep bf16 slabs regardless of the env.
            fp8kv: std::env::var("FP8KV").as_deref() == Ok("1")
                && cfg.head_dim == 128 && !cfg.sandwich && cfg.layer_hd.is_empty(),
            // Swept at m=8/2k: 1024 -> 564.9 tok/s, 2048 -> 578.0 (+2.3%), 4096/8192 slightly
            // below 2048. Short contexts (<=4 chunks) never reach these Y values.
            attnb_tgt: std::env::var("ATTNB_TGT").ok().and_then(|v| v.parse().ok()).unwrap_or(2048),
            attnw: std::env::var("ATTNW").as_deref() == Ok("1"),
            f_attn_part_w: f("attn_partial_w")?,
            // DEFAULT ON (MOET=0 disables). The kernel was fast and WRONG for one release:
            // it inherits qg2_body's inner loop, which decodes the fragment-PERMUTED q4
            // packing, while expert matrices load in the ORIGINAL packing — same speed, wrong
            // numbers, and a 16-token gate too short to see it. init_batch now permutes the
            // expert weights once (and does it there, not in the step, because the first
            // batched step runs inside graph capture and an allocation there aborts it).
            // Re-gated at 24 tokens with cycle detection: 22 distinct of 24, coherent,
            // deterministic, on both a short and a 2 000-token prompt.
            moet: std::env::var("MOET").as_deref() != Ok("0") && cfg.num_experts > 0,
            moet_pf: std::env::var("MOET_PF").as_deref() != Ok("0"),
            f_attn_out_gate: f("attn_out_gate_b")?,
            f_shared_expert: f("shared_expert_add")?,
            f_add_inplace: f("add_inplace")?,
            f_dn_conv: f("dn_conv")?,
            f_dn_step: { let ff = f("dn_step")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 131072)
                    .context("opt in to >48 KB dynamic shared for dn_step")?; } ff },
            f_moe_t: f("gemm_q4_moe_t")?,
            f_moe_t64: f("gemm_q4_moe_t64")?,
            f_moe_th2: f("gemm_q4_moe_th2")?,
            f_moe_t64h2: f("gemm_q4_moe_t64h2")?,
            f_moe_t64h4: f("gemm_q4_moe_t64h4")?,
            f_moe_th4: f("gemm_q4_moe_th4")?,
            f_moe_wl: f("moe_worklist")?,
            f_moe_sort: f("moe_sort_fused")?,
            mwork: stream.alloc_zeros::<u32>(1)?,
            mwn: stream.alloc_zeros::<i32>(1)?,
            attnt: std::env::var("ATTNT").as_deref() == Ok("1"),
            attnt_auto: std::env::var("ATTNT").is_err(),
            attnt_ctx: std::env::var("ATTNT_CTX").ok().and_then(|v| v.parse().ok()).unwrap_or(512),
            ctx_live: 0,
            f_attn_part_t: { let ff = f("attn_partial_t")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 143616)
                    .context("opt in to >48 KB dynamic shared for attn_partial_t")?; } ff },
            f_attn_part_b: f("attn_partial_b")?,
            f_attn_part_g: f("attn_partial_g")?,
            f_kv_gather: f("kv_gather")?,
            f_attn_strip: { let ff = f("attn_strip")?; if ccmaj >= 8 {
                    ff.set_attribute(
                        cudarc::driver::sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                        101376)
                    .context("opt in to >48 KB dynamic shared for attn_strip")?; } ff },
            f_attn_part_f: { let ff = f("attn_partial_f")?; if ccmaj >= 8 {
                    ff.set_attribute(
                        cudarc::driver::sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                        100000)
                    .context("opt in to >48 KB dynamic shared for attn_partial_f")?; } ff },
            f_attn_red_b: f("attn_reduce_b")?,
            f_attn_red_b16: f("attn_reduce_b16")?,
            f_argmax_b: f("argmax_b")?,
            f_advance_b: f("advance_b")?,
            f_embed_p: f("embed_prefill")?,
            f_rope_p: f("rope_prefill")?,
            f_vderive: f("v_derive")?,
            f_sandwich: f("sandwich_bf")?,
            f_q4rep_wg: f("q4_repack_wg")?,
            f_gwksum: f("gw_ksum")?,
            f_rms_hf: f("rmsnorm_add_hf")?,
            f_attn_red_h16: f("attn_reduce_h16")?,
            f_silu_hf: f("silu_mul_hf")?,
            ml_arm: std::env::var("MARLIN").as_deref() == Ok("1"),
            f_q4ml: { let ff = f("gemm_q4_ml")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 65536)
                    .context("opt in to >48 KB dynamic shared for gemm_q4_ml")?; } ff },
            f_q4ml64: { let ff = f("gemm_q4_ml64")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 49152)
                    .context("opt in for gemm_q4_ml64")?; } ff },
            f_q4ml32: { let ff = f("gemm_q4_ml32")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 45056)
                    .context("opt in for gemm_q4_ml32")?; } ff },
            gwpart: None,
            mlpart: None,
            f_q4wg64: { let ff = f("gemm_q4_wg64")?; if ccmaj >= 9 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 54272)
                    .context("opt in to >48 KB dynamic shared for gemm_q4_wg64")?; } ff },
            f_q4wg: { let ff = f("gemm_q4_wg")?; if ccmaj >= 9 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                // 2 x double-buffered (A[128][64] + B[128][64]) bf16 + align pad = ~64.3 KB
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 75776)
                    .context("opt in to >48 KB dynamic shared for gemm_q4_wg")?; } ff },
            f_softmax_c: f("softmax_causal")?,
            f_scale_out: f("scale_out_rows")?,
            f_silu_bf: f("silu_mul_bf")?,
            f_silu_gu: f("silu_gu_bf")?,
            f_fa: { let ff = f("fa_prefill")?; if ccmaj >= 8 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, FA_SHARED as i32)
                    .context("opt in to >48 KB dynamic shared for fa_prefill")?; } ff },
            f_fa_hop: { let ff = f("fa_prefill_hop")?; if ccmaj == 9 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, FH_SHARED as i32)
                    .context("opt in to >48 KB dynamic shared for fa_prefill_hop")?; } ff },
            f_fa_ws: { let ff = f("fa_prefill_ws")?; if ccmaj == 9 {
                use cudarc::driver::sys::CUfunction_attribute as FA;
                ff.set_attribute(FA::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, FW_SHARED as i32)
                    .context("opt in to >48 KB dynamic shared for fa_prefill_ws")?; } ff },
            f_vt: f("vt_pack")?,
            f_q4gemm: f("gemm_q4_tc")?,
            f_q4v2: [f("gemm_q4_v2_rg1")?, f("gemm_q4_v2_rg2")?,
                     f("gemm_q4_v2_rg4n2")?, f("gemm_q4_v2_rg4")?],
            f_q4red: f("qg2_reduce")?,
            f_q4rep: f("q4_repack")?,
            f_tofp8: f("bf16_to_fp8")?,
            f_absmax_bf: f("absmax_bf")?,
            f_rms_bf_ax: f("rmsnorm_add_bf_amax")?,
            f_silu_bf_ax: f("silu_mul_bf_amax")?,
            f_attn_red_b16_ax: f("attn_reduce_b16_amax")?,
            f_amax2s: f("amax_to_scale")?,
            fp8_amax: stream.alloc_zeros::<f32>(1)?,
            fp8_algo_cache: std::collections::HashMap::new(),
            fp8dw_dead: false,
            bf16_plans: std::collections::HashMap::new(),
            ctx_hint: 0,
            fp8_dsc: {
                // seed amax=64: rms-normed activations rarely exceed ~30, so step 1 clips only
                // extreme outliers while the rotate kernel installs real amax from step 2 on
                let n = cfg.layers * 4;
                let mut v = Vec::with_capacity(n * 3);
                for _ in 0..n { v.extend_from_slice(&[64.0f32, 64.0 / 448.0, 0.0]); }
                stream.memcpy_stod(&v)?
            },
            f_rms_fp8: f("rmsnorm_add_fp8")?,
            f_silu_fp8: f("silu_mul_fp8")?,
            f_attn_red_fp8: f("attn_reduce_fp8")?,
            f_sc_rot: f("fp8_scale_rotate")?,
            f_amax_f_part: f("absmax_f_part")?,
            f_tofp8_f: f("f32_to_fp8")?,
            f_mkscale: f("mk_fp8_scale")?,
            f_amax_part: f("absmax_bf_part")?,
            f_amax_fin: f("absmax_finish")?,
            fp8_part: stream.alloc_zeros::<f32>(64)?,
            fp8_pplans: Vec::new(),
            fp8_pplans_qkv: Vec::new(),
            fp8_pplans_o: Vec::new(),
            fp8_pplans_down: Vec::new(),
            fp8_pplan_t: 0,
            fp8_px: stream.alloc_zeros::<u8>(1)?,
            fp8_psx: Vec::new(),
            fp8_psmax: Vec::new(),
            fp8_gu: Vec::new(),
            fp8_plans: Vec::new(),
            fp8_qkv: Vec::new(),
            fp8_o: Vec::new(),
            fp8_down: Vec::new(),
            fp8_plans_qkv: Vec::new(),
            fp8_plans_o: Vec::new(),
            fp8_plans_down: Vec::new(),
            fp8_plan_m: 0,
            fp8_sxmax: stream.alloc_zeros::<f32>(1)?,
            fp8_plan: None,
            fp8_x: stream.alloc_zeros::<u8>(1)?,
            fp8_sx: stream.alloc_zeros::<f32>(1)?,
            q4p: Vec::new(),
            q4gu: Vec::new(),
            dummy_f32: stream.alloc_zeros::<f32>(1)?,
            dummy_f32b: stream.alloc_zeros::<f32>(1)?,
            dummy_u16: stream.alloc_zeros::<u16>(1)?,
            dummy_u16b: stream.alloc_zeros::<u16>(1)?,
            q4_on: true,
            q4free: false,
            q4part: None,
            fa: std::env::var("FA").as_deref() != Ok("0") && ccmaj >= 8,
            // Hopper's wgmma path. HOP=0 falls back to the mma.sync kernel for A/B measurement.
            hop: std::env::var("HOP").as_deref() != Ok("0") && ccmaj == 9,
            // Default on Hopper. HOP=1 selects the non-specialised wgmma kernel, HOP=0 mma.sync.
            ws: !matches!(std::env::var("HOP").as_deref(), Ok("0") | Ok("1")) && ccmaj == 9,
            fw_ticket: stream.alloc_zeros::<u32>(1)?,
            sm_count: ctx.attribute(
                cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)?,
            pf: None,
            batch: None,
            _ctx: ctx,
            stream,
            cfg,
            embed,
            layers,
            dn,
            final_norm,
            lm_head,
            f16_lm: None,
            kc,
            vc,
            max_seq,
        })
    }

    /// Fill `xs` with (max|x|, 1/max|x|) — must run on an activation before its GEMVs.
    /// Takes only the fields it needs so it composes with an outstanding `&self.layers[i]` borrow.
    fn absmax(
        stream: &Arc<CudaStream>,
        f: &CudaFunction,
        x: &CudaSlice<f32>,
        n: usize,
        xs: &mut CudaSlice<f32>,
    ) -> Result<()> {
        let ni = n as i32;
        unsafe {
            stream
                .launch_builder(f)
                .arg(x)
                .arg(xs)
                .arg(&ni)
                // One block, so threads-per-block is the only parallelism there is.
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    fn gemv(
        stream: &Arc<CudaStream>,
        f: &CudaFunction,
        w: &Q4Dev,
        x: &CudaSlice<f32>,
        xs: &CudaSlice<f32>,
        y: &mut CudaSlice<f32>,
        accumulate: bool,
    ) -> Result<()> {
        let cfg = LaunchConfig {
            grid_dim: (w.rows.div_ceil(32) as u32, 1, 1),
            block_dim: (1024, 1, 1),   // 32 warps = 32 rows sharing one staged activation tile
            shared_mem_bytes: 0,
        };
        let (nb, rows) = (w.nblk as i32, w.rows as i32);
        let accf = if accumulate { 1i32 } else { 0i32 };
        unsafe {
            stream
                .launch_builder(f)
                .arg(&w.scales)
                .arg(&w.quants)
                .arg(x)
                .arg(xs)
                .arg(y)
                .arg(&nb)
                .arg(&rows)
                .arg(&accf)
                .launch(cfg)?;
        }
        Ok(())
    }

    /// x += h, then RMSNorm into `out` (+ scale). One launch instead of add_inplace + rmsnorm.
    #[allow(clippy::too_many_arguments)]
    fn rms_add(
        stream: &Arc<CudaStream>,
        f: &CudaFunction,
        x: &mut CudaSlice<f32>,
        h: &CudaSlice<f32>,
        w: &CudaSlice<f32>,
        out: &mut CudaSlice<f32>,
        xs: &mut CudaSlice<f32>,
        n: usize,
        eps: f32,
    ) -> Result<()> {
        let ni = n as i32;
        unsafe {
            stream
                .launch_builder(f)
                .arg(x).arg(h).arg(w).arg(out).arg(xs).arg(&ni).arg(&eps)
                // One block per row, so threads-per-block is the only parallelism there is.
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    fn rms(
        stream: &Arc<CudaStream>,
        f: &CudaFunction,
        x: &CudaSlice<f32>,
        w: &CudaSlice<f32>,
        out: &mut CudaSlice<f32>,
        xs: &mut CudaSlice<f32>,
        n: usize,
        eps: f32,
    ) -> Result<()> {
        let c = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let ni = n as i32;
        unsafe {
            stream
                .launch_builder(f)
                .arg(x)
                .arg(w)
                .arg(out)
                .arg(xs)
                .arg(&ni)
                .arg(&eps)
                .launch(c)?;
        }
        Ok(())
    }

    /// One decode step. Takes NO host arguments: the token id, the position and the KV write
    /// offset all live in device memory, and nothing here allocates — both are required for the
    /// step to be capturable as a CUDA graph.
    pub fn step(&mut self, _token: u32, _pos: usize) -> Result<u32> {
        // The single-stream kernels have no sliding-window / scale-1 / v_norm arms; the server
        // gates SEQ1 off via seq_capable(), and this backstops the bench paths.
        anyhow::ensure!(!self.cfg.sandwich, "sequential step not implemented for sandwich models");
        if self.batch.is_none() {
            self.grow_seq_kv()?;
        }
        let (hd, nh, nkv, hidden) =
            (self.cfg.head_dim, self.cfg.n_heads, self.cfg.n_kv, self.cfg.hidden);
        let (kvdim, eps, hi) = (nkv * hd, self.cfg.eps, self.cfg.hidden as i32);
        let (nhi, nkvi, hdi) = (nh as i32, nkv as i32, hd as i32);
        

        let esc = self.cfg.embed_scale;
        unsafe {
            self.stream
                .launch_builder(&self.f_embed)
                .arg(&self.embed)
                .arg(&mut self.x)
                .arg(&mut self.h)
                .arg(&self.d_token)
                .arg(&hi)
                .arg(&esc)
                .launch(ecfg(hidden))?;
        }

        for li in 0..self.cfg.layers {
            // ---- attention ----
            Self::rms_add(&self.stream, &self.f_rms_add, &mut self.x, &self.h,
                          &self.layers[li].attn_norm, &mut self.abuf, &mut self.xs, hidden, eps)?;
            std::mem::swap(&mut self.h, &mut self.abuf);
            Self::gemv(&self.stream, &self.f_gemv, &self.layers[li].qkv, &self.h, &self.xs, &mut self.qkvbuf, false)?;

            // SEQ KV LIVES IN BATCH SLOT 0 once a batch exists: the server prefills b.kc[li]
            // (slot 0 = offset 0, same [max_seq, kvdim] layout as the single-stream slabs).
            // Attending over the never-filled self.kc was the server seq path's PROMPT
            // BLINDNESS — fluent output seeded only by the prefill's first token, which passed
            // every eyeball test and no HF gate (they all ran the batch arm or the bench, whose
            // non-batch prefill DOES fill self.kc).
            let (kcp, vcp, kv_stride): (*mut CudaSlice<u16>, *mut CudaSlice<u16>, i32) =
                match self.batch.as_mut() {
                    Some(b) => (&mut b.kc[li] as *mut _, &mut b.vc[li] as *mut _, b.max_seq as i32),
                    None => (&mut self.kc[li] as *mut _, &mut self.vc[li] as *mut _,
                             self.max_seq as i32),
                };
            unsafe {
                self.stream
                    .launch_builder(&self.f_qkrope)
                    .arg(&mut self.qkvbuf)
                    .arg(&self.rlog)
                    .arg(&self.layers[li].q_norm)
                    .arg(&self.layers[li].k_norm)
                    .arg(&mut *kcp).arg(&mut *vcp)
                    .arg(&nhi).arg(&nkvi).arg(&hdi)
                    .arg(&self.d_pos).arg(&self.inv_freq).arg(&self.qkni).arg(&eps).arg(&kv_stride)
                    .launch(LaunchConfig { grid_dim: ((nh + 2 * nkv) as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 })?;
            }
            let nsi = self.nsplit as i32;
            unsafe {
                self.stream
                    .launch_builder(&self.f_attn_part)
                    .arg(&self.qkvbuf).arg(&*kcp).arg(&*vcp)
                    .arg(&mut self.pm).arg(&mut self.pl).arg(&mut self.pacc)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&self.d_pos).arg(&nsi).arg(&kv_stride)
                    .launch(LaunchConfig {
                        grid_dim: (nh as u32, self.nsplit as u32, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: ((hd + 256) * 4) as u32,
                    })?;
                self.stream
                    .launch_builder(&self.f_attn_red)
                    .arg(&self.pm).arg(&self.pl).arg(&self.pacc).arg(&mut self.abuf)
                    .arg(&hdi).arg(&nsi)
                    .launch(LaunchConfig { grid_dim: (nh as u32, 1, 1), block_dim: (128, 1, 1), shared_mem_bytes: 0 })?;
            }
            // n_heads*head_dim, NOT hidden. They are equal for Qwen3 (16x128=2048) and Llama
            // (32x128=4096) so this never showed, but Mistral 3 has 32x128=4096 against hidden 5120:
            // the scale was computed over 1024 elements past the end of the attention output, which
            // loads fine and then emits a repeating token.
            Self::absmax(&self.stream, &self.f_absmax, &self.abuf, nh * hd, &mut self.xs)?;
            Self::gemv(&self.stream, &self.f_gemv, &self.layers[li].o, &self.abuf, &self.xs, &mut self.h, false)?;

            // ---- MLP / MoE ---- (residual add folded into this norm)
            {
                let (x, h) = (&mut self.x, &self.h);
                let _ = (x, h);
            }
            Self::rms_add(&self.stream, &self.f_rms_add, &mut self.x, &self.h,
                          &self.layers[li].ffn_norm, &mut self.abuf, &mut self.xs, hidden, eps)?;
            std::mem::swap(&mut self.h, &mut self.abuf);
            if self.cfg.num_experts > 0 {
                let (k, mi, e_n) = (self.cfg.topk, self.cfg.moe_inter, self.cfg.num_experts);
                let (ei, ki) = (e_n as i32, k as i32);
                let normi = if self.cfg.norm_topk { 1i32 } else { 0i32 };
                unsafe {
                    self.stream.launch_builder(&self.f_gemv_f32)
                        .arg(self.layers[li].router.as_ref().unwrap())
                        .arg(&self.h).arg(&mut self.rlog).arg(&ei).arg(&hi)
                        .launch(LaunchConfig { grid_dim: (e_n as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_topk)
                        .arg(&self.rlog).arg(&mut self.eidx).arg(&mut self.ewt)
                        .arg(&ei).arg(&ki).arg(&normi)
                        .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 })?;
                }
                if li == 0 && std::env::var("MOE_DUMP").is_ok() {
                    self.stream.synchronize()?;
                    let lg: Vec<f32> = self.stream.memcpy_dtov(&self.rlog.slice(0..e_n))?;
                    let ids: Vec<u32> = self.stream.memcpy_dtov(&self.eidx.slice(0..k))?;
                    let wt: Vec<f32> = self.stream.memcpy_dtov(&self.ewt.slice(0..k))?;
                    let mut top: Vec<(usize, f32)> = lg.iter().cloned().enumerate().collect();
                    top.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
                    eprintln!("[moeS l0] logits top4 {:?} sel {:?} wt {:?}", &top[..4], ids, wt);
                }
                let (st0, mii, shared1, shared0, one_i) = (0i32, mi as i32, 1i32, 0i32, 1i32);
                let nb1 = (hidden / 32) as i32;
                let mi2 = (2 * mi) as i32;
                let gcfg = LaunchConfig { grid_dim: ((2 * mi).div_ceil(32) as u32, k as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
                unsafe {
                    self.stream.launch_builder(&self.f_gemv_moe)
                        .arg(&self.layers[li].w1.as_ref().unwrap().scales)
                        .arg(&self.layers[li].w1.as_ref().unwrap().quants)
                        .arg(&self.h).arg(&self.xs).arg(&mut self.eg)
                        .arg(&nb1).arg(&mi2).arg(&self.eidx).arg(&st0).arg(&shared1).arg(&one_i).launch(gcfg)?;
                    self.stream.launch_builder(&self.f_silu_moe)
                        .arg(&self.eg).arg(&mut self.eu).arg(&mut self.exs).arg(&mii)
                        .launch(LaunchConfig { grid_dim: (k as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                }
                let nb2 = (mi / 32) as i32;
                let dcfg = LaunchConfig { grid_dim: (hidden.div_ceil(32) as u32, k as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
                unsafe {
                    self.stream.launch_builder(&self.f_gemv_moe)
                        .arg(&self.layers[li].w2.as_ref().unwrap().scales)
                        .arg(&self.layers[li].w2.as_ref().unwrap().quants)
                        .arg(&self.eu).arg(&self.exs).arg(&mut self.edn)
                        .arg(&nb2).arg(&hi).arg(&self.eidx).arg(&mii).arg(&shared0).arg(&one_i).launch(dcfg)?;
                    self.stream.launch_builder(&self.f_combine)
                        .arg(&self.edn).arg(&self.ewt).arg(&mut self.h).arg(&hi).arg(&ki)
                        .launch(ecfg(hidden))?;
                }
            } else {
                let inter = self.cfg.intermediate;
                Self::gemv(&self.stream, &self.f_gemv, &self.layers[li].gate, &self.h, &self.xs, &mut self.gbuf, false)?;
                Self::gemv(&self.stream, &self.f_gemv, &self.layers[li].up, &self.h, &self.xs, &mut self.ubuf, false)?;
                let (ni, act) = (inter as i32, self.cfg.gelu as i32);
                unsafe {
                    self.stream.launch_builder(&self.f_silu)
                        .arg(&mut self.gbuf).arg(&self.ubuf).arg(&ni).arg(&act).launch(ecfg(inter))?;
                }
                Self::absmax(&self.stream, &self.f_absmax, &self.gbuf, inter, &mut self.xs)?;
                Self::gemv(&self.stream, &self.f_gemv, &self.layers[li].down, &self.gbuf, &self.xs, &mut self.h, false)?;
            }
        }

        Self::rms_add(&self.stream, &self.f_rms_add, &mut self.x, &self.h, &self.final_norm,
                      &mut self.abuf, &mut self.xs, hidden, eps)?;
        std::mem::swap(&mut self.h, &mut self.abuf);
        Self::gemv(&self.stream, &self.f_gemv, &self.lm_head, &self.h, &self.xs, &mut self.logits, false)?;
        let (v, mo) = (self.cfg.vocab as i32, self.max_seq as i32);
        unsafe {
            self.stream.launch_builder(&self.f_argmax)
                .arg(&self.logits).arg(&mut self.d_token).arg(&v)
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
            self.stream.launch_builder(&self.f_advance)
                .arg(&mut self.d_pos).arg(&self.d_token).arg(&mut self.d_out).arg(&mo)
                .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 })?;
        }
        Ok(0)
    }

    /// Seed the device-side token/position (host -> device once, not per step).
    pub fn seed(&mut self, token: u32, pos: usize) -> Result<()> {
        self.stream.memcpy_htod(&[token], &mut self.d_token)?;
        self.stream.memcpy_htod(&[pos as i32], &mut self.d_pos)?;
        self.sync()
    }

    /// Capture ONE decode step as a CUDA graph. Every per-step input (token id, position, KV write
    /// offset) already lives in device memory and every grid size is position-independent, so the
    /// captured graph is valid for every subsequent step. This is the fix for the thing that lost
    /// us the 1.7B head-to-head: ~540 kernel launches/token at ~8us of launch overhead each. vLLM
    /// wins there for exactly this reason (CUDA graphs); replaying one graph collapses it.
    pub fn capture(&mut self) -> Result<()> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        self.sync()?;
        self.stream
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
        // record (does not execute)
        let r = self.step(0, 0);
        let g = self.stream.end_capture(
            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );
        // the step's own error is the informative one; end_capture just reports INVALIDATED
        r.context("during graph capture")?;
        let g = g?.context("stream capture produced no graph")?;
        g.upload()?;
        self.graph = Some(g);
        self.sync()
    }

    /// Capture the BATCHED step as a graph too — ~280 tiny launches per step (norms, rope,
    /// attention, silu, absmax across 28 layers) are not amortised by batching, and at M=8 they are
    /// roughly a third of the step.
    /// Narrow (or widen) the batch to `m` sequences. Slots >= m are then ignored entirely, so the
    /// caller must keep the live set packed into [0, m) — `free_slot` returns the lowest free index,
    /// which does exactly that. Buffers and KV slabs are unaffected: they are indexed by slot.
    pub fn set_width(&mut self, m: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(m >= 1 && m <= b.mmax, "width {m} outside 1..={}", b.mmax);
        b.m = m;
        Ok(())
    }

    pub fn width(&self) -> usize {
        self.batch.as_ref().map(|b| b.m).unwrap_or(0)
    }

    /// Is a graph already captured for the current width?
    /// Scheduler-provided max live context; part of the decode-graph key via the ATTNT policy.
    pub fn set_ctx_live(&mut self, c: usize) {
        self.ctx_live = c;
    }

    /// The tensor-attention policy this step (and therefore which graph key applies).
    fn attnt_now(&self) -> bool {
        (self.attnt || (self.attnt_auto && self.ctx_live >= self.attnt_ctx))
            && self.cfg.head_dim == 128 && !self.fp8kv && self.cfg.layer_hd.is_empty()
            && self.cfg.n_heads % self.cfg.n_kv == 0 && (self.cfg.n_heads / self.cfg.n_kv) <= 8
    }

    pub fn has_batch_graph(&self) -> bool {
        self.bgraphs.contains_key(&(self.width(), self.attnt_now()))
    }

    pub fn capture_batch(&mut self) -> Result<()> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        // plan building tunes Lt algos with real timed matmuls — do it BEFORE capture
        let w = self.width();
        self.fp8_decode_plans(w)?;
        if self.fp8_qkv.is_empty() {
            self.bf16_decode_plans(w)?;
        }
        self.sync()?;
        self.stream
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
        let r = self.step_batch();
        let g = self.stream.end_capture(
            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );
        r.context("during batched graph capture")?;
        let g = g?.context("no graph from batched capture")?;
        g.upload()?;
        let w = self.width();
        let arm = self.attnt_now();
        self.bgraphs.insert((w, arm), g);
        self.sync()
    }

    pub fn run_graph_batch(&self, n: usize) -> Result<()> {
        let g = self
            .bgraphs
            .get(&(self.width(), self.attnt_now()))
            .context("capture_batch() first for this width")?;
        for _ in 0..n {
            g.launch()?;
        }
        Ok(())
    }

    /// Replay the captured step `n` times with NO host synchronisation in the loop.
    pub fn run_graph(&self, n: usize) -> Result<()> {
        let g = self.graph.as_ref().context("capture() first")?;
        for _ in 0..n {
            g.launch()?;
        }
        Ok(())
    }

    /// Read the tokens the device recorded at absolute positions `[from, to)`.
    pub fn tokens(&self, from: usize, to: usize) -> Result<Vec<u32>> {
        Ok(self.stream.memcpy_dtov(&self.d_out.slice(from..to))?)
    }

    /// Allocate the M-sequence batch state. `bmax` is the per-sequence context budget.
    pub fn init_batch(&mut self, m: usize, bmax: usize) -> Result<()> {
        // The M<=8 cap was a shared-memory bound of the batched GEMV; the batched path now runs
        // on cuBLAS, so M is limited only by scratch and KV memory.
        anyhow::ensure!((1..=64).contains(&m), "batched path supports M in 1..=64");
        // MoE routers as bf16 copies: HF's gate is a bf16 Linear, and expert selection must
        // reproduce ITS rounding (bf16 GEMM + bf16-rounded logits), not out-precision it.
        if self.cfg.num_experts > 0 && self.router16.is_empty() {
            let n = (self.cfg.num_experts * self.cfg.hidden) as i32;
            let mut v = Vec::with_capacity(self.cfg.layers);
            for li in 0..self.cfg.layers {
                let r = self.layers[li].router.as_ref().context("MoE model without router")?;
                let mut r16 = self.stream.alloc_zeros::<u16>(n as usize)?;
                unsafe {
                    self.stream.launch_builder(&self.f_tobf16)
                        .arg(r).arg(&mut r16).arg(&n).launch(ecfg(n as usize))?;
                }
                v.push(r16);
            }
            self.router16 = v;
            self.sync()?;
        }
        if self.cfg.num_experts > 0 {
            self.moe_grow(m)?;
        }
        let c = &self.cfg;
        let (hd, nh, nkv) = (c.head_dim, c.n_heads, c.n_kv);
        let nsplit = bmax.div_ceil(64);
        let st = &self.stream;
        let mut kc = Vec::with_capacity(c.layers);
        let mut vc = Vec::with_capacity(c.layers);
        for li in 0..c.layers {
            // bf16: this is the dominant traffic of a decode step at scale, ~2.2 GB per step at
            // M=64 against 0.79 GB of weights, and vLLM's cache is bf16 too. Per-layer kvdim:
            // Gemma4's sliding and global layers cache different widths.
            // FP8KV=1: the slab holds ONE e4m3 byte per element, packed into half-length u16
            // buffers — the _f8 kernel twins index it as bytes; every element-offset formula
            // carries over unchanged.
            let (lhd, lnkv, _, _) = c.lgeom(li);
            let elems = m * bmax * lnkv * lhd;
            let words = if self.fp8kv { elems.div_ceil(2) } else { elems };
            kc.push(st.alloc_zeros::<u16>(words)?);
            vc.push(st.alloc_zeros::<u16>(words)?);
        }
        // Permuted q4 copies for the q4 decode arm (Q4DEC=1): same bytes as the q4 weights, and
        // 3.55x smaller than the bf16 mirrors they are meant to replace.
        // MEMORY-GATED: the arm costs ~1.45x the q4 bytes (q4p + merged q4gu) ON TOP of the bf16
        // mirrors built below. On Gemma-31B that is ~24 GB over a ~64 GB mirror+embed floor —
        // the arm was added after the Q4FREE fit work and silently pushed the 31B load OOM
        // (engine thread died before "model loaded"). Project the bill against free VRAM and
        // skip the arm when it cannot fit: the mirror arm serves every width, just slower at
        // m<=16, and a booting server beats a fast one that never loads.
        let arm_fits = {
            let (mut arm, mut mir, mut orig) = (0usize, 0usize, 0usize);
            for li in 0..self.cfg.layers {
                let l = &self.layers[li];
                for w in [&l.qkv, &l.o, &l.gate, &l.up, &l.down] {
                    arm += w.rows * w.nblk * 16;
                    mir += w.rows * w.nblk * 64;
                    orig += w.rows * w.nblk * 16;
                }
                arm += (l.gate.rows + l.up.rows) * l.gate.nblk * 18; // merged q4gu: qp + f16 scales
            }
            mir += self.lm_head.rows * self.lm_head.nblk * 64;
            let freed = if std::env::var("Q4FREE").as_deref() == Ok("1") { orig } else { 0 };
            let need = (arm + mir + (6usize << 30)).saturating_sub(freed);
            match cudarc::driver::result::mem_get_info() {
                Ok((free, _)) => {
                    if free < need {
                        eprintln!("q4 decode arm SKIPPED: needs ~{:.1} GiB (arm {:.1} + mirrors {:.1} + margin - freed {:.1}), {:.1} free — mirror arm serves all widths",
                            need as f64 / 1e9, arm as f64 / 1e9, mir as f64 / 1e9, freed as f64 / 1e9, free as f64 / 1e9);
                    }
                    free >= need
                }
                Err(_) => true,
            }
        };
        if arm_fits && std::env::var("Q4DEC").as_deref() != Ok("0") && self.cfg.num_experts == 0 && self.q4p.is_empty() {
            let mut all = Vec::with_capacity(self.cfg.layers);
            for li in 0..self.cfg.layers {
                let mut set: Vec<CudaSlice<u32>> = Vec::with_capacity(5);
                for k in 0..5 {
                    let l = &self.layers[li];
                    let w = match k { 0 => &l.qkv, 1 => &l.o, 2 => &l.gate, 3 => &l.up, _ => &l.down };
                    let (nbk, rws) = (w.nblk, w.rows);
                    let mut out = self.stream.alloc_zeros::<u32>(rws * nbk * 4)?;
                    let nbi = (rws * nbk) as i64;
                    unsafe {
                        self.stream.launch_builder(&self.f_q4rep)
                            .arg(&w.quants).arg(&mut out).arg(&nbi).launch(ecfg(rws * nbk))?;
                    }
                    set.push(out);
                }
                let a: [CudaSlice<u32>; 5] = set.try_into().map_err(|_| anyhow::anyhow!("5"))?;
                all.push(a);
            }
            self.q4p = all;
            // merged gate|up: the q4 layout is row-major by output row, so concatenating the two
            // matrices' rows IS the merged matrix — no repacking beyond the usual permute.
            let mut gus = Vec::with_capacity(self.cfg.layers);
            for li in 0..self.cfg.layers {
                let (g, u) = (&self.layers[li].gate, &self.layers[li].up);
                anyhow::ensure!(g.nblk == u.nblk, "gate/up block mismatch");
                let (nbk, rows) = (g.nblk, g.rows + u.rows);
                let mut q = self.stream.alloc_zeros::<u32>(rows * nbk * 4)?;
                let mut sc = self.stream.alloc_zeros::<f16>(rows * nbk)?;
                let gq = g.rows * nbk * 4;
                self.stream.memcpy_dtod(&g.quants, &mut q.slice_mut(0..gq))?;
                self.stream.memcpy_dtod(&u.quants, &mut q.slice_mut(gq..rows * nbk * 4))?;
                let gs = g.rows * nbk;
                self.stream.memcpy_dtod(&g.scales, &mut sc.slice_mut(0..gs))?;
                self.stream.memcpy_dtod(&u.scales, &mut sc.slice_mut(gs..rows * nbk))?;
                let mut qp = self.stream.alloc_zeros::<u32>(rows * nbk * 4)?;
                let nbi = (rows * nbk) as i64;
                unsafe {
                    self.stream.launch_builder(&self.f_q4rep)
                        .arg(&q).arg(&mut qp).arg(&nbi).launch(ecfg(rows * nbk))?;
                }
                gus.push((qp, sc));
            }
            self.q4gu = gus;
            let widest = ((nh + 2 * nkv) * hd).max(self.cfg.intermediate).max(self.cfg.hidden);
            self.q4part = Some(self.stream.alloc_zeros::<f32>(8 * m * widest)?);
            self.sync()?;
        }

        // f16 mirrors for the tensor-core batched path (dense arm)
        {
            // `ws` may be several q4 matrices to lay down back-to-back in one buffer (gate|up).
            let mut mk = |ws: &[&Q4Dev]| -> Result<CudaSlice<u16>> {
                let n: usize = ws.iter().map(|w| w.rows * w.nblk * 32).sum();
                let mut out = self.stream.alloc_zeros::<u16>(n)?;
                let mut off = 0i64;
                for w in ws {
                    let (nbi, ri) = (w.nblk as i32, w.rows as i32);
                    let blocks = (w.rows * w.nblk).div_ceil(256);
                    unsafe {
                        self.stream.launch_builder(&self.f_dq)
                            .arg(&w.scales).arg(&w.quants).arg(&mut out).arg(&nbi).arg(&ri).arg(&off)
                            .launch(LaunchConfig { grid_dim: (blocks as u32,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    }
                    off += (w.rows * w.nblk * 32) as i64;
                }
                Ok(out)
            };
            // Q4FREE=1: drop each layer's ORIGINAL q4 quants AS SOON AS its mirrors exist.
            // Freeing after the whole build (the first version) left a PEAK of originals+all
            // mirrors co-resident — 85 GB on Gemma-31B (q4 16.5 + mirrors 58.6 + f32 embed 5.6
            // + tied lm q4 ×2), an OOM the interleaved free turns into ~72 GB. Scales stay (the
            // q4 arm reads them). Opt-in because it kills every path that multiplies originals:
            // step()/seq (SEQ1 — seq_capable gates it), the M=1 gemv arm (self.q4free gates it),
            // gemm_tc/gemm_q4 verify paths, and the MoE arm.
            let q4free = std::env::var("Q4FREE").as_deref() == Ok("1");
            let memlog = std::env::var("MEMLOG").is_ok();
            // MIRRORFREE=1 (requires FP8=1): the bf16 mirrors exist only long enough to derive
            // each layer's e4m3 copies, then DROP — every path that would multiply them (prefill,
            // wave, decode) has an fp8 arm. Peak co-residency is ONE layer's mirrors instead of
            // all of them: a 24B goes from 47 GB mirrors + 24 GB fp8 (OOM next to anything) to
            // ~24 GB fp8 total, and the 31B's 61.6 GB mirror wall disappears. FA-path models
            // only for now — the score-path o-GEMM (Gemma) still multiplies mirrors.
            // Dense models only. Sandwich/score-path (Gemma) is covered since the score-path
            // o-GEMM gained its f32-input fp8 arm and sandwich decode runs the two-stage chains.
            let mirror_free = std::env::var("FP8").is_ok()
                && std::env::var("MIRRORFREE").as_deref() == Ok("1")
                && self.cfg.num_experts == 0;
            let (mut vq8, mut vo8, mut vgu8, mut vd8) =
                (Vec::new(), Vec::new(), Vec::new(), Vec::new());
            for li in 0..self.cfg.layers {
                if memlog && li % 8 == 0 {
                    self.stream.synchronize()?;
                    if let Ok((free, _)) = cudarc::driver::result::mem_get_info() {
                        eprintln!("[mem] mirror build layer {li}: {:.1} GiB free", free as f64 / 1e9);
                    }
                }
                let l = &self.layers[li];
                let (a, b2, gu, e) =
                    (mk(&[&l.qkv])?, mk(&[&l.o])?, mk(&[&l.gate, &l.up])?, mk(&[&l.down])?);
                // the dequant kernels above must finish before their source is dropped
                if q4free || mirror_free { self.stream.synchronize()?; }
                if mirror_free {
                    let st = self.stream.clone();
                    vq8.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8,
                        &self.f_mkscale, &a, a.len())?);
                    vo8.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8,
                        &self.f_mkscale, &b2, b2.len())?);
                    vgu8.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8,
                        &self.f_mkscale, &gu, gu.len())?);
                    vd8.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8,
                        &self.f_mkscale, &e, e.len())?);
                    // conversions must finish before the mirrors drop at end of scope
                    self.stream.synchronize()?;
                } else {
                    let l = &mut self.layers[li];
                    l.f16_qkv = Some(a); l.f16_o = Some(b2); l.f16_gu = Some(gu); l.f16_down = Some(e);
                }
                if q4free {
                    let l = &mut self.layers[li];
                    let st2 = self.stream.clone();
                    for w in [&mut l.qkv, &mut l.o, &mut l.gate, &mut l.up, &mut l.down] {
                        w.quants = st2.alloc_zeros::<u32>(1)?;
                    }
                }
            }
            if mirror_free {
                self.fp8_qkv = vq8;
                self.fp8_o = vo8;
                self.fp8_gu = vgu8;
                self.fp8_down = vd8;
                let m_slots = self.batch.as_ref().map(|bt| bt.mmax).unwrap_or(64);
                let mut wid = self.cfg.hidden.max(self.cfg.intermediate);
                for li in 0..self.cfg.layers { wid = wid.max(self.layers[li].o.nblk * 32); }
                self.fp8_x = self.stream.alloc_zeros::<u8>(m_slots * wid)?;
                self.fp8_sx = self.stream.alloc_zeros::<f32>(1)?;
                self.fp8_sxmax = self.stream.alloc_zeros::<f32>(1)?;
                self.fp8_plan_m = 0;
                self.fp8_pplan_t = 0;
            }
            let lm = mk(&[&self.lm_head])?;
            self.f16_lm = Some(lm);
            if q4free {
                // the tied/untied LM head's originals are equally dead (gemv/seq gated off)
                self.stream.synchronize()?;
                let st2 = self.stream.clone();
                self.lm_head.quants = st2.alloc_zeros::<u32>(1)?;
                self.q4free = true;
            }
            self.sync()?;
        }

        // FP8 e4m3 mirrors of gate|up (FP8=1). vLLM has a precision tier we lack — with
        // --quantization fp8 it does 9775/16745 at M=32/64 against our 8580/15134 — and the kernel
        // headroom is measured at 1.71x bf16, so this is competing in THEIR tier, not cross-pairing.
        if std::env::var("FP8").is_ok() && self.cfg.num_experts == 0 && self.fp8_gu.is_empty() {
            let mut v = Vec::with_capacity(self.cfg.layers);
            for li in 0..self.cfg.layers {
                let n = (2 * self.cfg.intermediate * self.cfg.hidden) as i64;
                let src = self.layers[li].f16_gu.as_ref().context("f16 mirrors first")?;
                let mut sc = self.stream.alloc_zeros::<f32>(1)?;
                let mut q = self.stream.alloc_zeros::<u8>(n as usize)?;
                unsafe {
                    self.stream.launch_builder(&self.f_absmax_bf).arg(src).arg(&mut sc).arg(&n)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8).arg(src).arg(&mut q).arg(&sc).arg(&n)
                        .launch(ecfg((n as usize).div_ceil(2)))?;
                }
                let mut swd = self.stream.alloc_zeros::<f32>(1)?;
                unsafe {
                    self.stream.launch_builder(&self.f_mkscale).arg(&sc).arg(&mut swd)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
                }
                v.push((q, swd));
            }
            self.fp8_gu = v;
            // The step's OTHER three weight matrices, same scheme — with all four in e4m3 the
            // decode step reads half the bytes the bf16 mirrors cost.
            let (mut vq, mut vo, mut vd) = (Vec::new(), Vec::new(), Vec::new());
            for li in 0..self.cfg.layers {
                let st = self.stream.clone();
                let l = &self.layers[li];
                vq.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8, &self.f_mkscale,
                    l.f16_qkv.as_ref().context("f16 mirrors first")?, l.qkv.rows * l.qkv.nblk * 32)?);
                vo.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8, &self.f_mkscale,
                    l.f16_o.as_ref().context("f16 mirrors first")?, l.o.rows * l.o.nblk * 32)?);
                vd.push(Self::fp8_mirror(&st, &self.f_absmax_bf, &self.f_tofp8, &self.f_mkscale,
                    l.f16_down.as_ref().context("f16 mirrors first")?, l.down.rows * l.down.nblk * 32)?);
            }
            self.fp8_qkv = vq;
            self.fp8_o = vo;
            self.fp8_down = vd;
            // sized for the widest per-GEMM activation (down sees m*inter; o sees m*nh*lhd)
            let mut wid = self.cfg.hidden.max(self.cfg.intermediate);
            for li in 0..self.cfg.layers { wid = wid.max(self.layers[li].o.nblk * 32); }
            self.fp8_x = self.stream.alloc_zeros::<u8>(m * wid)?;
            // Device-computed per-step activation scale (absmax two-stage) — no host readback.
            self.fp8_sx = self.stream.alloc_zeros::<f32>(1)?;
            self.fp8_sxmax = self.stream.alloc_zeros::<f32>(1)?;
            // Plans are built by the step's pre-loop rebuild (fp8_plan_m == 0 forces it) so the
            // four plan sets have ONE build site and stay in lockstep with the live width.
            self.fp8_plan_m = 0;
            self.sync()?;
        }

        // Marlin split-K partials, sized for the widest GEMM this model will run at this batch
        // width (max over qkv / o / gate|up / down row counts) x the largest ksplit the
        // dispatcher can pick. Allocated BEFORE any graph capture so it never moves.
        if self.ml_arm {
            let widest_rows = ((nh + 2 * nkv) * hd).max(c.hidden).max(c.intermediate);
            let need = 8 * m * widest_rows;
            if self.mlpart.as_ref().map(|p| p.len()).unwrap_or(0) < need {
                self.mlpart = Some(st.alloc_zeros::<f32>(need)?);
            }
        }
        // ---- gated-DeltaNet state, one pair of buffers per linear layer ----
        let (dnk, dnv, ddk, ddv, dkn) =
            (c.dn_nk, c.dn_nv, c.dn_dk, c.dn_dv, c.dn_conv_kernel);
        let dn_conv_dim = if c.layer_linear.is_empty() { 0 } else { 2 * dnk * ddk + dnv * ddv };
        let mut dn_s = Vec::new();
        let mut dn_ring = Vec::new();
        for li in 0..c.layers {
            let linear = c.layer_linear.get(li).copied().unwrap_or(false);
            // one-element placeholders on full-attention layers: indexing stays by layer id,
            // which keeps the forward pass free of a second index map
            let (ns, nr) = if linear { (m * dnv * ddk * ddv, m * dn_conv_dim * dkn) } else { (1, 1) };
            dn_s.push(st.alloc_zeros::<f32>(ns)?);
            dn_ring.push(st.alloc_zeros::<f32>(nr)?);
        }
        let dn_mixed = st.alloc_zeros::<f32>((m * dn_conv_dim).max(1))?;
        let dn_z = st.alloc_zeros::<f32>((m * dnv * ddv).max(1))?;
        let dn_b = st.alloc_zeros::<f32>((m * dnv).max(1))?;
        let dn_a = st.alloc_zeros::<f32>((m * dnv).max(1))?;
        let dn_core = st.alloc_zeros::<f32>((m * dnv * ddv).max(1))?;
        let dn_core16 = st.alloc_zeros::<u16>((m * dnv * ddv).max(1))?;
        if dn_conv_dim > 0 {
            let per_slot = (dnv * ddk * ddv + dn_conv_dim * dkn) * 4;
            eprintln!("DeltaNet state: {:.1} MB/slot/layer x {} linear layers x {m} slots = {:.1} GB",
                per_slot as f64 / 1e6, c.layer_linear.iter().filter(|x| **x).count(),
                per_slot as f64 * c.layer_linear.iter().filter(|x| **x).count() as f64 * m as f64 / 1e9);
        }
        self.batch = Some(Batch {
            m, mmax: m, max_seq: bmax, nsplit,
            x: st.alloc_zeros::<f32>(m * c.hidden)?,
            h: st.alloc_zeros::<f32>(m * c.hidden)?,
            qkv: st.alloc_zeros::<f32>(m * (nh + 2 * nkv) * hd)?,
            // max(nh*hd, hidden), NOT nh*hd: the decode tail swaps h <-> abuf every step, so
            // the two buffers ALTERNATE identities by step parity. On models where nh*hd <
            // hidden (Mistral: 4096 vs 5120) an nh*hd-sized abuf meant every odd-parity step
            // treated a 32k-element buffer as the 40k-element residual — a silent OOB write
            // the mixed step's bounds-checked slice was the first to catch.
            abuf: st.alloc_zeros::<f32>(m * (nh * hd).max(c.hidden))?,
            g: st.alloc_zeros::<f32>(m * c.intermediate)?,
            u: st.alloc_zeros::<f32>(m * c.intermediate)?,
            hb16: st.alloc_zeros::<u16>(m * c.hidden)?,
            attn16: st.alloc_zeros::<u16>(m * c.n_heads * c.head_dim)?,
            abuf16: st.alloc_zeros::<u16>(m * c.n_heads * c.head_dim)?,
            gu16: st.alloc_zeros::<u16>(m * 2 * c.intermediate)?,
            gu32: st.alloc_zeros::<f32>(m * 2 * c.intermediate)?,
            g16: st.alloc_zeros::<u16>(m * c.intermediate)?,
            logits: st.alloc_zeros::<f32>(m * c.vocab)?,
            xs: st.alloc_zeros::<f32>(2 * m)?,
            kc, vc,
            pm: st.alloc_zeros::<f32>(m * nh * nsplit)?,
            pl: st.alloc_zeros::<f32>(m * nh * nsplit)?,
            pacc: st.alloc_zeros::<f32>(m * nh * nsplit * hd)?,
            tok: st.alloc_zeros::<u32>(m)?,
            out: st.alloc_zeros::<u32>(m * bmax)?,
            pos: st.alloc_zeros::<i32>(m)?,
            active: st.alloc_zeros::<u32>(m)?,
            dn_s, dn_ring, dn_mixed, dn_z, dn_b, dn_a, dn_core, dn_core16,
            occupied: vec![false; m],
        });
        // gemm_bf16 stages activations through xf16, whose boot size is one row; every batched
        // LM-head call needs m * hidden. The wave path grew it lazily, which left pure-long
        // workloads (no waves) writing past an 8192-element buffer.
        if self.xf16.len() < m * c.hidden {
            self.xf16 = self.stream.alloc_zeros::<u16>(64 * c.hidden)?;
            self.xf16b = self.stream.alloc_zeros::<u16>(64 * c.hidden)?;
        }
        // THE MoE TENSOR ARM READS PERMUTED WEIGHTS. gemm_q4_moe_t inherits qg2_body's inner
        // loop, whose nibble extraction assumes the fragment-friendly permutation q4_repack
        // produces — but expert matrices load with quants_p: None, the ORIGINAL q4_0 packing
        // the per-pair GEMV consumes. Decoding one layout with the other's index map is
        // exactly as fast and completely wrong: that is how this shipped as "+113%, ids
        // identical" past a 16-token gate.
        //
        // It builds HERE, in init_batch, and not in the step: the first batched step runs
        // INSIDE capture_batch's graph capture, and an allocation during capture aborts it —
        // which is precisely what the first attempt did (engine dead, zero tokens).
        if self.moet {
            let n = self.cfg.layers;
            for li in 0..n {
                let mut w1 = self.layers[li].w1.take().context("MoE w1")?;
                self.ensure_permuted(&mut w1)?;
                self.layers[li].w1 = Some(w1);
                let mut w2 = self.layers[li].w2.take().context("MoE w2")?;
                self.ensure_permuted(&mut w2)?;
                self.layers[li].w2 = Some(w2);
            }
            self.sync()?;
            eprintln!("MoE tensor arm: expert weights permuted for the mma fragment layout");
        }
        Ok(())
    }

    pub fn seed_batch(&mut self, tokens: &[u32], pos: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(tokens.len() == b.m, "expected {} tokens", b.m);
        let m = b.m;
        self.stream.memcpy_htod(tokens, &mut b.tok)?;
        self.stream.memcpy_htod(&vec![pos as i32; m], &mut b.pos)?;
        self.stream.memcpy_htod(&vec![1u32; m], &mut b.active)?;
        b.occupied = vec![true; m];
        self.sync()
    }

    /// Free a slot: it stops advancing while the rest of the batch continues.
    pub fn release_slot(&mut self, slot: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(slot < b.mmax, "slot out of range");
        self.stream.memcpy_htod(&[0u32], &mut b.active.slice_mut(slot..slot + 1))?;
        b.occupied[slot] = false;
        Ok(())
    }

    /// First free slot, if any.
    pub fn free_slot(&self) -> Option<usize> {
        self.batch.as_ref()?.occupied.iter().position(|o| !o)
    }

    /// Place a sequence in `slot` at `pos` with `token` as its next input, and mark it active.
    pub fn occupy_slot(&mut self, slot: usize, token: u32, pos: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(slot < b.mmax, "slot out of range");
        self.stream.memcpy_htod(&[token], &mut b.tok.slice_mut(slot..slot + 1))?;
        self.stream.memcpy_htod(&[pos as i32], &mut b.pos.slice_mut(slot..slot + 1))?;
        self.stream.memcpy_htod(&[1u32], &mut b.active.slice_mut(slot..slot + 1))?;
        b.occupied[slot] = true;
        Ok(())
    }

    /// Patch a slot's next-input token (the batched admission path occupies with a placeholder
    /// before the prefill burst is drained, then writes the real sampled token here).
    pub fn set_slot_token(&mut self, slot: usize, token: u32) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(slot < b.mmax, "slot out of range");
        self.stream.memcpy_htod(&[token], &mut b.tok.slice_mut(slot..slot + 1))?;
        Ok(())
    }

    /// Tokens sampled this step, one per slot (host readback — one small copy per step).
    /// The token the single-stream path (incl. `prefill`) most recently sampled.
    pub fn last_token(&self) -> Result<u32> {
        Ok(self.stream.memcpy_dtov(&self.d_token)?[0])
    }

    pub fn batch_last_tokens(&self) -> Result<Vec<u32>> {
        let b = self.batch.as_ref().context("init_batch first")?;
        Ok(self.stream.memcpy_dtov(&b.tok)?)
    }

    /// One lockstep decode step for all M sequences. Dense path.
    pub fn step_batch(&mut self) -> Result<()> {
        let mut b = self.batch.take().context("init_batch first")?;
        let r = self.step_batch_inner(&mut b);
        self.batch = Some(b);
        r
    }

    fn step_batch_inner(&mut self, b: &mut Batch) -> Result<()> {
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let eps = c.eps;
        let (m, bmax, nsplit) = (b.m, b.max_seq, b.nsplit);
        // partial rotary width (0 = full rope); Qwen3.5/3.6 rotate only head_dim * 0.25
        let roti = c.rot_dim as i32;
        // DECODE_PROF=1: per-stage sync timing across the step (the decode analogue of
        // PREFILL_PROF). Only meaningful on an EAGER step — the server bypasses the decode
        // graph under this env for exactly that reason.
        let prof = std::env::var("DECODE_PROF").is_ok();
        let mut tp = [0f64; 10];
        let mut tk = std::time::Instant::now();
        macro_rules! mark { ($i:expr) => { if prof {
            self.stream.synchronize()?;
            let n2 = std::time::Instant::now();
            tp[$i] += n2.duration_since(tk).as_secs_f64();
            tk = n2;
        } } }
        // TWO ARMS, chosen by batch width.
        //
        // Batched decode is memory-bound, and the weights dominate: the q4 GEMV streams ~0.79 GB
        // per step (4.5 bits/param) against the bf16 GEMM's ~2.82 GB, a 3.55x advantage that holds
        // for ANY M because both read every weight once. What kills the GEMV at large M is its
        // inner loop over the M activations per weight block, which turns it ALU-bound. So: q4
        // GEMV while it is still bandwidth-bound, cuBLAS once there is enough arithmetic to hide
        // the wider weights. GEMV_MAXM matches the kernel's MAXM register bound.
        //
        // (The earlier note here said cuBLAS was 5x slower, measured when every weight was
        // dequantised per step. Resident bf16 mirrors removed that cost, so the comparison is now
        // purely weight bytes against ALU work.)
        // MEASURED crossover (Qwen3-1.7B, H100 PCIe, aggregate tok/s, GEMV vs GEMM):
        //   M=1  325.8 vs 285.0   GEMV wins
        //   M=4  843.5 vs 1121.9  GEMM wins
        //   M=8 1134.7 vs 2215.1  GEMM wins decisively
        // So the q4 GEMV's 3.55x weight-byte advantage only survives at M=1; by M=4 its per-M
        // inner loop has made it ALU-bound and cuBLAS is ahead.
        // Q4 DECODE ARM: multiply the q4 weights directly instead of their bf16 mirrors. Decode is
        // memory-bound on weights — 0.79 GB per step against 2.82 GB — and the standalone kernel is
        // at cuBLAS parity, so back-to-back inside the captured graph (no per-launch cost) is where
        // that 3.55x should finally show. Behind Q4DEC so the default path is untouched.
        // NOT ON BY DEFAULT — requires Q4DEC=1, which is also what builds the permuted copies.
        //
        // SETTLED by VERIFY_ARM=1, and only because the divergence was chased instead of assumed.
        // Against the single-stream q4 reference over 24 greedy tokens, the bf16 arm matched on all
        // 24 while this arm differed at exactly one (20412 vs 107614) — which looked like a defect.
        // The top-2 logit margin at that step decides it:
        //   bf16 arm: top1 107614, logit 19.86710, margin over top2 0.00181
        //   q4   arm: top1  20412, logit 19.90377, margin over top2 0.00196
        // A 0.0018 margin on a logit of 19.9 is 0.009 % — a genuine tie, far under bf16 rounding
        // noise — and both arms re-converge on the very next token. Not a defect, so this is ON by
        // default (Q4DEC=0 disables). The cost is one permuted q4 copy of the five projections,
        // which is the same bytes as the q4 weights and 3.55x smaller than the bf16 mirrors.
        //
        // MEASURED crossover, end-to-end aggregate tok/s (Qwen3-1.7B, H100):
        //   M=8  2224 -> 2353 (+5.8%)   M=16 4328 -> 4483 (+3.6%)
        //   M=32 7737 -> 6623 (-14%)    M=64 13765 -> 7785 (-43%)
        // which tracks the standalone GEMM exactly: parity at M<=16, 1.9x behind at 64. So the q4
        // arm takes the narrow widths and the bf16 mirrors keep the wide ones.
        // M=32 re-probed after the FFN became one GEMM: 7025 vs 7739 for the mirrors. Still loses,
        // so the crossover really is 16 — the GEMM's 1.25x deficit at 32 is more than the fusion saves.
        const Q4_MAXM: usize = 16;
        // MARLIN ARM: raw-q4 GEMM with register-resident dequant. It reads 0.79 GB/step where
        // the bf16 mirrors read 2.82 GB, and the isolated-GEMM deficit to cuBLAS (1.9x) is
        // measured against a 90 TF number the FULL step never reaches — the whole mirror step
        // runs at ~42 TF effective, which is what makes this worth wiring end to end.
        let ml_arm = self.ml_arm && c.num_experts == 0 && m <= 128 && self.layers[0].qkv.quants.len() > 1;
        let q4_arm = !ml_arm && self.q4_on && !self.q4p.is_empty() && m <= Q4_MAXM && c.num_experts == 0;
        const GEMV_MAXM: usize = 1;
        // The gemv arm reads ORIGINAL q4 quants — dead after Q4FREE. The sandwich (Gemma4) flow
        // never emits the f32-normed input + absmax this arm consumes, so it is off there too.
        let gemv_arm = !ml_arm && m <= GEMV_MAXM && std::env::var("BGEMV").as_deref() != Ok("0")
            && !self.q4free && !c.sandwich && c.num_experts == 0;
        // fp8 decode chains use producer-side absmax (the _amax kernel variants) — the 3-launch
        // absmax chain per GEMM was ~3 ms of serialized tiny-kernel GPU time at M=64
        let fp8dec = !self.fp8_qkv.is_empty() && !ml_arm && !gemv_arm && !c.sandwich;
        // Delayed-scale direct-write producers: OPT-IN (FP8DW=1). They work at narrow widths but
        // some Lt algo/plan combinations return NOT_SUPPORTED *inside CUDA graph capture* at
        // 64 slots, which aborts the decode-graph capture and kills the engine thread. The
        // amax-fold chains are the shipped default until that is understood; the kernels and
        // the rotate machinery stay in-tree for the follow-up.
        let fp8dw = fp8dec && std::env::var("FP8DW").as_deref() == Ok("1") && !self.fp8dw_dead;
        // Sandwich (Gemma) decode also runs the fp8 GEMMs, but its norms are f_sandwich — no
        // producer accumulates the absmax, so these chains keep the two-stage absmax kernels.
        let fp8sw = !self.fp8_qkv.is_empty() && !ml_arm && !gemv_arm && c.sandwich;
        let shb = (m * 64 * 17 * 4) as u32;
        let (mi, hi, nhi, nkvi, hdi, nsi, bmi) =
            (m as i32, hidden as i32, nh as i32, nkv as i32, hd as i32, nsplit as i32, bmax as i32);
        let per = (nh + 2 * nkv) * hd;
        // 1024 threads for the per-row reductions (rmsnorm/sandwich): 256 left each block
        // walking 2048 elements with a barrier per step — 749 us/step at M=64, measured.
        let one = LaunchConfig { grid_dim: (m as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 };
        let one256 = LaunchConfig { grid_dim: (m as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };

        // FP8 plan sets are width-specific; rebuild all four together BEFORE the layer loop —
        // qkv consumes its plan at layer 0. Idempotent no-op once built for this width, and
        // capture_batch() pre-calls it so the ALGO TUNING SYNCS never run inside a capture.
        if !self.fp8_gu.is_empty() && self.fp8_plan_m != m {
            self.fp8_decode_plans(m)?;
        }
        // plan building may have just auto-disabled FP8DW — the flag must agree with the plans
        // THIS step runs (dw skips the amax/tofp8 chains, which the standard plans depend on)
        let fp8dw = fp8dw && !self.fp8dw_dead;
        let esc = c.embed_scale;
        unsafe {
            self.stream.launch_builder(&self.f_embed_b)
                .arg(&self.embed).arg(&mut b.x).arg(&mut b.h).arg(&b.tok).arg(&hi).arg(&mi).arg(&esc)
                .launch(ecfg(m * hidden))?;
        }
        mark!(0);
        for li in 0..c.layers {
            // Per-layer attention geometry (Gemma4); uniform models see cfg's values unchanged.
            let (lhd, lnkv, lwin, lrope2) = c.lgeom(li);
            let (nkvi, hdi) = (lnkv as i32, lhd as i32);
            let wini = lwin as i32;
            let ascale = if c.attn_scale_one { 1.0f32 } else { 1.0 / (lhd as f32).sqrt() };
            // ---- GATED DELTANET LAYER (Qwen3.5/3.6, Ornith) ----
            // The token mixer is a recurrence, not attention: norm -> four projections ->
            // depthwise causal conv -> delta rule against the persistent state -> gated
            // RMSNorm -> out_proj. It writes b.h exactly as the attention branch does, so the
            // FFN below and the residual bookkeeping are untouched.
            if c.layer_linear.get(li).copied().unwrap_or(false) {
                let (nk, nv, dk, dv, kn) = (c.dn_nk, c.dn_nv, c.dn_dk, c.dn_dv, c.dn_conv_kernel);
                let conv_dim = 2 * nk * dk + nv * dv;
                unsafe {
                    self.stream.launch_builder(&self.f_rms_bf)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].attn_norm)
                        .arg(&mut b.hb16).arg(&hi).arg(&eps).launch(one)?;
                }
                {
                    let d = self.dn[li].as_ref().context("DeltaNet weights for a linear layer")?;
                    Self::gemm_pre(&self.stream, &self.blas, &d.qkv, &b.hb16, &mut b.dn_mixed,
                                   m, conv_dim, hidden)?;
                    Self::gemm_pre(&self.stream, &self.blas, &d.z, &b.hb16, &mut b.dn_z,
                                   m, nv * dv, hidden)?;
                    Self::gemm_pre(&self.stream, &self.blas, &d.bp, &b.hb16, &mut b.dn_b,
                                   m, nv, hidden)?;
                    Self::gemm_pre(&self.stream, &self.blas, &d.ap, &b.hb16, &mut b.dn_a,
                                   m, nv, hidden)?;
                }
                {
                    let (cdi, kni, mi2) = (conv_dim as i32, kn as i32, m as i32);
                    let (nki, nvi2, dki, dvi) = (nk as i32, nv as i32, dk as i32, dv as i32);
                    let shb2 = ((dk * dv + 2 * dk + 32) * 4) as u32;
                    let d = self.dn[li].as_ref().unwrap();
                    unsafe {
                        self.stream.launch_builder(&self.f_dn_conv)
                            .arg(&mut b.dn_mixed).arg(&mut b.dn_ring[li]).arg(&d.conv_w)
                            .arg(&cdi).arg(&kni).arg(&mi2)
                            .launch(ecfg(m * conv_dim))?;
                        self.stream.launch_builder(&self.f_dn_step)
                            .arg(&mut b.dn_s[li]).arg(&b.dn_mixed).arg(&b.dn_z)
                            .arg(&b.dn_b).arg(&b.dn_a)
                            .arg(&d.a_log).arg(&d.dt_bias).arg(&d.norm_w).arg(&mut b.dn_core)
                            .arg(&nki).arg(&nvi2).arg(&dki).arg(&dvi).arg(&eps)
                            .launch(LaunchConfig { grid_dim: (nv as u32, m as u32, 1),
                                                   block_dim: (dv as u32, 1, 1),
                                                   shared_mem_bytes: shb2 })?;
                        let nci = (m * nv * dv) as i32;
                        self.stream.launch_builder(&self.f_tobf16)
                            .arg(&b.dn_core).arg(&mut b.dn_core16).arg(&nci)
                            .launch(ecfg(m * nv * dv))?;
                    }
                    Self::gemm_pre(&self.stream, &self.blas, &d.out, &b.dn_core16, &mut b.h,
                                   m, hidden, nv * dv)?;
                }
                mark!(3);
                // fall through to the shared FFN/MoE block below
            } else {
            unsafe {
                if c.sandwich {
                    // Layer top: fold the PREVIOUS layer's post-FFN norm + residual + layer_scalar
                    // together with this layer's input norm (h = 0 at li == 0, so the dummy postw
                    // contributes rms(0)·w = 0 and ls of the previous layer applies correctly).
                    let (postw, pls) = if li == 0 {
                        (&self.layers[0].attn_norm, 1.0f32)
                    } else {
                        (self.layers[li - 1].post_ffw_norm.as_ref().unwrap(), self.layers[li - 1].ls)
                    };
                    self.stream.launch_builder(&self.f_sandwich)
                        .arg(&mut b.x).arg(&b.h).arg(postw).arg(&self.layers[li].attn_norm)
                        .arg(&mut b.hb16).arg(&hi).arg(&eps).arg(&pls).launch(one)?;
                } else if gemv_arm {
                    self.stream.launch_builder(&self.f_rms_add_b)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].attn_norm)
                        .arg(&mut b.abuf).arg(&mut b.xs).arg(&hi).arg(&eps).launch(one256)?;
                } else if fp8dw {
                    let mut slot = self.fp8_dsc.slice_mut((li * 4 + 0) * 3..(li * 4 + 0) * 3 + 3);
                    self.stream.launch_builder(&self.f_rms_fp8)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].attn_norm).arg(&mut self.fp8_x)
                        .arg(&hi).arg(&eps).arg(&mut slot).launch(one)?;
                } else if fp8dec {
                    self.stream.launch_builder(&self.f_rms_bf_ax)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].attn_norm).arg(&mut b.hb16)
                        .arg(&hi).arg(&eps).arg(&mut self.fp8_amax).launch(one)?;
                } else {
                    self.stream.launch_builder(if ml_arm { &self.f_rms_hf } else { &self.f_rms_bf })
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].attn_norm).arg(&mut b.hb16)
                        .arg(&hi).arg(&eps).launch(one)?;
                }
            }
            if ml_arm {
                Self::gemm_q4_ml_raw(&self.stream, &self.f_q4ml, &self.f_gwksum,
                    self.mlpart.as_mut().unwrap(), &self.layers[li].qkv, &b.hb16, &mut b.qkv,
                    m, (nh + 2 * lnkv) * lhd, hidden)?;
            } else if gemv_arm {
                std::mem::swap(&mut b.h, &mut b.abuf);
                Self::gemv_b(&self.stream, &self.f_gemv_b, &self.layers[li].qkv, &b.h, &b.xs, &mut b.qkv, m, hidden, shb)?;
            } else if q4_arm {
                Self::gemm_q4s(&self.stream, &self.f_q4v2, &self.f_q4red,
                    self.q4part.as_mut().unwrap(), &self.layers[li].qkv.scales, &self.q4p[li][0],
                    self.layers[li].qkv.nblk, &b.hb16, &mut b.qkv, m, (nh + 2 * lnkv) * lhd, hidden)?;
            } else if fp8dec || fp8sw {
                let nx = (m * hidden) as i64;
                if !fp8dw {
                    unsafe {
                        if fp8dec {
                            // producer (rms_bf_ax) already accumulated the absmax
                            self.stream.launch_builder(&self.f_amax2s)
                                .arg(&mut self.fp8_amax).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
                        } else {
                            let g = 64i32;
                            self.stream.launch_builder(&self.f_amax_part)
                                .arg(&b.hb16).arg(&mut self.fp8_part).arg(&nx)
                                .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                            self.stream.launch_builder(&self.f_amax_fin)
                                .arg(&self.fp8_part).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx).arg(&g)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        }
                        self.stream.launch_builder(&self.f_tofp8)
                            .arg(&b.hb16).arg(&mut self.fp8_x).arg(&self.fp8_sxmax).arg(&nx)
                            .launch(ecfg((m * hidden).div_ceil(2)))?;
                    }
                }
                let (lt, ltws) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt, &self.fp8_plans_qkv[li], ltws, 1 << 20,
                    &self.fp8_qkv[li].0, &self.fp8_x, &mut b.qkv, 1.0)?;
            } else {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (m, self.layers[li].qkv.rows, hidden),
                    self.layers[li].f16_qkv.as_ref().unwrap(), &b.hb16, &mut b.qkv,
                    m, (nh + 2 * lnkv) * lhd, hidden)?;
            }
            mark!(1);
            unsafe {
                if c.v_norm {
                    self.stream.launch_builder(&self.f_vderive)
                        .arg(&mut b.qkv).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&eps).arg(&mi)
                        .launch(LaunchConfig { grid_dim: ((m * lnkv) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                }
                let invf = if lrope2 { &self.inv_freq2 } else { &self.inv_freq };
                self.stream.launch_builder(if self.fp8kv { &self.f_rope_b_f8 } else { &self.f_rope_b })
                    .arg(&mut b.qkv).arg(&mut b.kc[li]).arg(&mut b.vc[li])
                    .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(invf).arg(&self.qkni).arg(&eps).arg(&bmi)
                    .arg(&b.active).arg(&roti)
                    .launch(LaunchConfig { grid_dim: ((m * (nh + 2 * lnkv)) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                mark!(2);
                // GROUPED ARM: one block per (slot, kv head) reads each K/V row ONCE for the
                // whole GQA group. Llama 32q/8kv was reading every K/V row four times — the
                // largest non-GEMM term of the step at ~790 GB/s of a ~2 TB/s HBM.
                // MEASURED NEGATIVE, default OFF (ATTNG=<hpb> to try again). Sharing K/V rows
                // across query heads INSIDE a block (hpb=2 and hpb=4 both) cost 20% on Llama
                // M=64 (5503 -> 4385 / 4393) even though it removes 2-4x of the K/V traffic:
                // the base kernel's 4-rows-per-warp, 2-loads-in-flight structure hides latency
                // better than the grouped version's per-head shuffle reductions and smem score
                // round-trip, and the grid shrinks by hpb. The right way to bank the traffic
                // saving is the flash-decoding structure — stage the K/V CHUNK in shared memory
                // once, then let one warp per query head consume it from smem — not in-register
                // sharing. Output is bit-identical either way (ids matched).
                // Context-aware arm choice, measured: at ~150 tokens KV is ~7% of the step's
                // bytes and the flash kernel's staging costs more than the 2-4x KV traffic it
                // saves (M=8 short: 797 vs 820); at 2k it is ~35% and flash wins (M=1 2k:
                // 173.8 vs 168.6). ATTNF=1/0 forces it on/off for A/B.
                let attnf_env = std::env::var("ATTNF").ok();
                let flash_ok = lhd == 128 && nh % lnkv == 0 && (nh / lnkv) >= 2 && (nh / lnkv) <= 8;
                // MEASURED policy: flash wins only at M=1 with long context (2k: 173.8 vs
                // 168.6). At M=8/2k it LOSES (Llama 348 vs 355) and at short context it loses
                // (797 vs 820) — the KV traffic it saves is a small share of the step unless
                // the batch is narrow AND the context is long.
                let flash = !self.fp8kv && flash_ok && match attnf_env.as_deref() {
                    Some("1") => true,
                    Some("0") => false,
                    _ => m == 1 && self.ctx_hint >= 512,
                };
                let hpb: usize = std::env::var("ATTNG").ok().and_then(|v| v.parse().ok()).unwrap_or(0);
                let grouped = !self.fp8kv && hpb >= 2 && lhd == 128 && nh % lnkv == 0
                    && (nh / lnkv) % hpb == 0 && nh % hpb == 0;
                if flash {
                    let gq = nh / lnkv;
                    // vred is [wpg][gq][hd] and wpg = nwarp/gq, so it is nwarp*hd floats for
                    // ANY gq. Sizing it 2*gq*hd assumed wpg==2 and overflowed on gq=2 models
                    // (Qwen: wpg=4) — the server returned zero tokens.
                    let nwarp = 8usize;
                    let fsh = ((2 * 64 * lhd * 2) + (gq * lhd + gq * 64 + nwarp * lhd) * 4) as u32;
                    self.stream.launch_builder(&self.f_attn_part_f)
                        .arg(&b.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                        .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                        .arg(&ascale).arg(&wini)
                        .launch(LaunchConfig { grid_dim: ((m * lnkv) as u32,
                                                          (2048usize.div_ceil(m * lnkv)).clamp(4, nsplit.max(4)) as u32,
                                                          1),
                                               block_dim: (256,1,1), shared_mem_bytes: fsh })?;
                } else if grouped {
                    let hpbi = hpb as i32;
                    let gsh = ((hpb * lhd + hpb * 64 + 8 * hpb * lhd) * 4) as u32;
                    self.stream.launch_builder(&self.f_attn_part_g)
                        .arg(&b.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                        .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                        .arg(&ascale).arg(&wini).arg(&hpbi)
                        .launch(LaunchConfig { grid_dim: ((m * nh / hpb) as u32,
                                                          (1024usize.div_ceil(m * nh / hpb)).clamp(4, nsplit.max(4)) as u32,
                                                          1),
                                               block_dim: (256,1,1), shared_mem_bytes: gsh })?;
                } else if self.attnt_now() && lhd == 128 {
                    // TENSOR arm: blocks are (slot, kv head); each of the 4 warps owns whole
                    // 64-key chunks. Partials contract unchanged.
                    let yt = (self.attnb_tgt.div_ceil((m * lnkv).max(1)))
                        .clamp(1, nsplit.div_ceil(4).max(1)) as u32;
                    self.stream.launch_builder(&self.f_attn_part_t)
                        .arg(&b.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                        .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                        .arg(&ascale).arg(&wini)
                        .launch(LaunchConfig { grid_dim: ((m * lnkv) as u32, yt, 1),
                                               block_dim: (128, 1, 1),
                                               shared_mem_bytes: 143616 })?;
                } else if self.attnw && lhd == 128 && !self.fp8kv {
                    // WARP-PRIVATE chunks: nwarp chunks per block-y, no block barriers in the
                    // loop. Same pm/pl/pacc contract; the reduce is unchanged.
                    let yw = (self.attnb_tgt.div_ceil(m * nh))
                        .clamp(1, nsplit.div_ceil(8).max(1)) as u32;
                    self.stream.launch_builder(&self.f_attn_part_w)
                        .arg(&b.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                        .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                        .arg(&ascale).arg(&wini)
                        .launch(LaunchConfig { grid_dim: ((m * nh) as u32, yw, 1),
                                               block_dim: (256, 1, 1),
                                               shared_mem_bytes: ((lhd + 8 * 64) * 4) as u32 })?;
                } else {
                self.stream.launch_builder(if self.fp8kv { &self.f_attn_part_b_f8 } else { &self.f_attn_part_b })
                    .arg(&b.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                    .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                    .arg(&ascale).arg(&wini)
                    // The kernel strides chunks by gridDim.y; nsplit stays the PARTIALS STRIDE
                    // only. Y is sized so dispatch neither scales with max_seq (the M=64 storm)
                    // nor starves the GPU of blocks (the M=1 long-context collapse).
                    // ATTNB_TGT: the block target behind Y. DECODE_PROF at m=8/2k measured this
                    // kernel LATENCY-bound at 0.68 TB/s — each block walking 8 chunks serially
                    // with two loads in flight; more Y = more loads in flight machine-wide.
                    // Short contexts (<=4 chunks) are structurally unaffected.
                    .launch(LaunchConfig { grid_dim: ((m * nh) as u32,
                                                      (self.attnb_tgt.div_ceil(m * nh)).clamp(4, nsplit.max(4)) as u32,
                                                      1 ),
                                           block_dim: (256,1,1), shared_mem_bytes: ((lhd + 256) * 4) as u32 })?;
                }
                mark!(3);
                let rcfg = LaunchConfig { grid_dim: ((m * nh) as u32, 1, 1),
                                          block_dim: (128, 1, 1), shared_mem_bytes: 0 };
                if gemv_arm {
                    // only the q4 GEMV path consumes f32 attention output and the absmax scales
                    self.stream.launch_builder(&self.f_attn_red_b)
                        .arg(&b.pm).arg(&b.pl).arg(&b.pacc).arg(&mut b.abuf).arg(&hdi).arg(&nsi)
                        .arg(&nhi).arg(&b.pos).arg(&wini).launch(rcfg)?;
                    // same n_heads*head_dim vs hidden confusion as the single-stream path
                    let ahi = (nh * lhd) as i32;
                    self.stream.launch_builder(&self.f_absmax_slots)
                        .arg(&b.abuf).arg(&mut b.xs).arg(&ahi).arg(&ahi).launch(one256)?;
                } else if fp8dw {
                    let mut slot = self.fp8_dsc.slice_mut((li * 4 + 1) * 3..(li * 4 + 1) * 3 + 3);
                    self.stream.launch_builder(&self.f_attn_red_fp8)
                        .arg(&b.pm).arg(&b.pl).arg(&b.pacc).arg(&mut self.fp8_x).arg(&hdi).arg(&nsi)
                        .arg(&nhi).arg(&b.pos).arg(&wini).arg(&mut slot).launch(rcfg)?;
                } else if fp8dec {
                    self.stream.launch_builder(&self.f_attn_red_b16_ax)
                        .arg(&b.pm).arg(&b.pl).arg(&b.pacc).arg(&mut b.abuf16).arg(&hdi).arg(&nsi)
                        .arg(&nhi).arg(&b.pos).arg(&wini).arg(&mut self.fp8_amax).launch(rcfg)?;
                } else {
                    // both GEMM arms want bf16, so reduce straight into it and skip the absmax,
                    // whose output nothing on those paths reads
                    self.stream.launch_builder(if ml_arm { &self.f_attn_red_h16 }
                                               else { &self.f_attn_red_b16 })
                        .arg(&b.pm).arg(&b.pl).arg(&b.pacc).arg(&mut b.abuf16).arg(&hdi).arg(&nsi)
                        .arg(&nhi).arg(&b.pos).arg(&wini).launch(rcfg)?;
                }
            }
            // ATTN_DUMP2: fingerprint layer-0 attention partials + reduced output so the scalar
            // and tensor arms can be diffed numerically (pm/pl localize the S side, pacc the PV
            // side). Eager runs only.
            if li == 0 && prof && std::env::var("ATTN_DUMP2").is_ok() {
                self.stream.synchronize()?;
                let pmv: Vec<f32> = self.stream.memcpy_dtov(&b.pm.slice(0..4 * nsplit))?;
                let plv: Vec<f32> = self.stream.memcpy_dtov(&b.pl.slice(0..4 * nsplit))?;
                let pav: Vec<f32> = self.stream.memcpy_dtov(&b.pacc.slice(0..8))?;
                for gh in 0..4 {
                    let sps: Vec<String> = (0..4.min(nsplit))
                        .map(|s| format!("{:.4}/{:.3}", pmv[gh * nsplit + s], plv[gh * nsplit + s]))
                        .collect();
                    eprintln!("[attn2 gh={gh}] pm/pl sp0-3: {}", sps.join(" "));
                }
                eprintln!("[attn2] pacc[gh0 sp0 d0-7]: {:?}",
                    pav.iter().map(|x| (x * 1e3).round() / 1e3).collect::<Vec<_>>());
                self.dump_attn("dec-l0", &b.abuf16, &[0, 1], nh * lhd)?;
                // output gate (Qwen3.5/3.6): the gate rows sit past q|k|v in the fused buffer
                if c.attn_out_gate {
                    let per_hd = (nh + 2 * lnkv + nh) * lhd;
                    let (goff, nhhd) = (((nh + 2 * lnkv) * lhd) as i32, (nh * lhd) as i32);
                    let (phd, mi3) = (per_hd as i32, m as i32);
                    unsafe {
                    self.stream.launch_builder(&self.f_attn_out_gate)
                        .arg(&b.qkv).arg(&mut b.abuf16).arg(&goff).arg(&nhhd).arg(&phd).arg(&mi3)
                        .launch(ecfg(m * nh * lhd))?;
                    }
                }
            }
            mark!(4);
            if ml_arm {
                Self::gemm_q4_ml_raw(&self.stream, &self.f_q4ml, &self.f_gwksum,
                    self.mlpart.as_mut().unwrap(), &self.layers[li].o, &b.abuf16, &mut b.h,
                    m, hidden, nh * lhd)?;
            } else if gemv_arm {
                Self::gemv_b(&self.stream, &self.f_gemv_b, &self.layers[li].o, &b.abuf, &b.xs, &mut b.h, m, nh * lhd, shb)?;
            } else if q4_arm {
                Self::gemm_q4s(&self.stream, &self.f_q4v2, &self.f_q4red,
                    self.q4part.as_mut().unwrap(), &self.layers[li].o.scales, &self.q4p[li][1],
                    self.layers[li].o.nblk, &b.abuf16, &mut b.h, m, hidden, nh * lhd)?;
            } else if fp8dec || fp8sw {
                let nx = (m * nh * lhd) as i64;
                if !fp8dw {
                    unsafe {
                        if fp8dec {
                            // producer (attn_red_b16_ax) already accumulated the absmax
                            self.stream.launch_builder(&self.f_amax2s)
                                .arg(&mut self.fp8_amax).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
                        } else {
                            let g = 64i32;
                            self.stream.launch_builder(&self.f_amax_part)
                                .arg(&b.abuf16).arg(&mut self.fp8_part).arg(&nx)
                                .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                            self.stream.launch_builder(&self.f_amax_fin)
                                .arg(&self.fp8_part).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx).arg(&g)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        }
                        self.stream.launch_builder(&self.f_tofp8)
                            .arg(&b.abuf16).arg(&mut self.fp8_x).arg(&self.fp8_sxmax).arg(&nx)
                            .launch(ecfg((m * nh * lhd).div_ceil(2)))?;
                    }
                }
                let (lt, ltws) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt, &self.fp8_plans_o[li], ltws, 1 << 20,
                    &self.fp8_o[li].0, &self.fp8_x, &mut b.h, 1.0)?;
            } else {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (m, hidden, self.layers[li].o.nblk * 32),
                    self.layers[li].f16_o.as_ref().unwrap(), &b.abuf16, &mut b.h,
                    m, hidden, nh * lhd)?;
            }
            mark!(5);
            }   // end of the full-attention branch; both arms have written b.h
            unsafe {
                if c.sandwich {
                    // x = x + rms(attn_h)·post_attn_norm (ls = 1 mid-layer), then the FFN input
                    // norm into hb16 — one launch, mirroring the pre-norm models' fusion.
                    let lone = 1.0f32;
                    self.stream.launch_builder(&self.f_sandwich)
                        .arg(&mut b.x).arg(&b.h)
                        .arg(self.layers[li].post_attn_norm.as_ref().unwrap())
                        .arg(&self.layers[li].ffn_norm)
                        .arg(&mut b.hb16).arg(&hi).arg(&eps).arg(&lone).launch(one)?;
                } else if gemv_arm || c.num_experts > 0 {
                    // MoE shares this flavour: the expert GEMVs consume f32 rows + per-row absmax
                    self.stream.launch_builder(&self.f_rms_add_b)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].ffn_norm)
                        .arg(&mut b.abuf).arg(&mut b.xs).arg(&hi).arg(&eps).launch(one256)?;
                } else if fp8dw {
                    let mut slot = self.fp8_dsc.slice_mut((li * 4 + 2) * 3..(li * 4 + 2) * 3 + 3);
                    self.stream.launch_builder(&self.f_rms_fp8)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].ffn_norm).arg(&mut self.fp8_x)
                        .arg(&hi).arg(&eps).arg(&mut slot).launch(one)?;
                } else if fp8dec {
                    self.stream.launch_builder(&self.f_rms_bf_ax)
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].ffn_norm).arg(&mut b.hb16)
                        .arg(&hi).arg(&eps).arg(&mut self.fp8_amax).launch(one)?;
                } else {
                    self.stream.launch_builder(if ml_arm { &self.f_rms_hf } else { &self.f_rms_bf })
                        .arg(&mut b.x).arg(&b.h).arg(&self.layers[li].ffn_norm).arg(&mut b.hb16)
                        .arg(&hi).arg(&eps).launch(one)?;
                }
            }
            let (n8, ii, ni) = ((m * inter / 8) as i32, inter as i32, (m * inter) as i32);
            let act = c.gelu as i32;
            if c.num_experts > 0 {
                // BATCHED MoE FFN: bf16 router GEMM -> per-row top-k -> per-(row,expert)-pair
                // grouped GEMVs (gemv_q4_moe with blockIdx.y = pair; x_div folds k pairs onto
                // their row's input + scale) -> fused silu+absmax -> down -> weighted combine.
                // Phase 1 reads each pair's expert weights independently (~4x traffic vs a
                // dedup pass at m=64); correctness first, expert-major dedup is the next lever.
                std::mem::swap(&mut b.h, &mut b.abuf);
                // ---- SHARED EXPERT (Qwen3.5/3.6) ----
                // A dense SwiGLU every token runs ALONGSIDE its routed experts, added under a
                // sigmoid gate of the layer input. It runs HERE, before the routed path, because
                // b.h still holds that input — the routed combine overwrites it. The gated
                // result parks in b.abuf (free after the swap) and is added back afterwards.
                let shared = self.layers[li].sw_gu.is_some();
                if shared {
                    let si = c.shared_expert_inter;
                    let (sii, sni) = (si as i32, (m * si) as i32);
                    let act2 = c.gelu as i32;
                    let gu_w = self.layers[li].sw_gu.as_ref().unwrap().clone_ref();
                    let dn_w = self.layers[li].sw_down.as_ref().unwrap().clone_ref();
                    unsafe {
                        self.stream.launch_builder(&self.f_absmax_slots)
                            .arg(&b.h).arg(&mut b.xs).arg(&hi).arg(&hi).launch(one256)?;
                    }
                    Self::gemv_b(&self.stream, &self.f_gemv_b, &gu_w, &b.h, &b.xs,
                                 &mut b.gu32, m, hidden, shb)?;
                    unsafe {
                        self.stream.launch_builder(&self.f_silu_gu)
                            .arg(&b.gu32).arg(&mut self.xf16).arg(&sii).arg(&sni).arg(&act2)
                            .launch(ecfg(m * si))?;
                        self.stream.launch_builder(&self.f_absmax_slots)
                            .arg(&b.gu32).arg(&mut b.xs).arg(&sii).arg(&sii).launch(one256)?;
                    }
                    Self::gemv_b(&self.stream, &self.f_gemv_b, &dn_w, &b.gu32, &b.xs,
                                 &mut b.abuf, m, si, shb)?;
                    let gt_w = self.layers[li].sw_gate.as_ref().unwrap();
                    unsafe {
                        self.stream.launch_builder(&self.f_shared_expert)
                            .arg(&mut b.abuf).arg(&b.h).arg(gt_w).arg(&hi).arg(&mi)
                            .launch(LaunchConfig { grid_dim: (m as u32, 1, 1),
                                                   block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    }
                }
                let (e_n, k, mi) = (c.num_experts, c.topk, c.moe_inter);
                let pairs = m * k;
                // Tiles are 64 rows tall everywhere, and the kernel skips the sub-tiles routing
                // did not fill. That guard is what makes one height safe for both regimes.
                // Height buys arithmetic intensity -- a quantised weight quad feeds MT mma
                // instead of one -- and prefill has enough pairs per expert to use it (~125 at
                // T=2000, worth +17% on 8x2k). Decode does not: streams are admitted staggered,
                // so they sit at different positions, route differently, and leave only a few
                // pairs per expert. Without the guard those tiles are ~94% empty and the wasted
                // mma cost 16% at M=64. Width still follows the block count, which is why narrow
                // batches take the 64-column entry. tsh must match the height the kernel
                // consumes: the sort emits tiles of 16*MT.
                // DECODE keeps 16-row tiles. Height buys arithmetic intensity -- one
                // quantised weight quad feeds MT mma instead of one -- but only if the rows
                // exist, and at decode they do not: streams are admitted staggered, so they
                // sit at different positions, route to different experts, and leave a
                // handful of pairs per expert. 64-row tiles there are ~94% empty and the
                // wasted mma cost 16% at M=64. Skipping the empty sub-tiles is worse still
                // (-5 to -12%): a data-dependent break inside the unrolled mma loop defeats
                // the unroll and spills the fragment array. Prefill is the opposite regime
                // and takes the tall tile at its own site.
                let tsh: i32 = 4;
                let nih = (m * hidden) as i32;
                unsafe {
                    self.stream.launch_builder(&self.f_tobf16)
                        .arg(&b.h).arg(&mut b.hb16).arg(&nih).launch(ecfg(m * hidden))?;
                }
                Self::gemm_pre(&self.stream, &self.blas, &self.router16[li], &b.hb16,
                    &mut self.mrlog, m, e_n, hidden)?;

                let (ei, ki) = (e_n as i32, k as i32);
                let normi = if c.norm_topk { 1i32 } else { 0i32 };
                let (nb1, mi2, mii, hsi) = ((hidden / 32) as i32, (2 * mi) as i32, mi as i32, hidden as i32);
                let (xs1, xdivk, one_d, sh0) = (hidden as i32, k as i32, 1i32, 0i32);
                unsafe {
                    self.stream.launch_builder(&self.f_topk_b)
                        .arg(&self.mrlog).arg(&mut self.meidx).arg(&mut self.mewt)
                        .arg(&ei).arg(&ki).arg(&normi)
                        .launch(LaunchConfig { grid_dim: (m as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                }
                if li == 0 && std::env::var("MOE_DUMP").is_ok() {
                    self.stream.synchronize()?;
                    let lg: Vec<f32> = self.stream.memcpy_dtov(&self.mrlog.slice(0..e_n))?;
                    let ids: Vec<u32> = self.stream.memcpy_dtov(&self.meidx.slice(0..k))?;
                    let wt: Vec<f32> = self.stream.memcpy_dtov(&self.mewt.slice(0..k))?;
                    let mut top: Vec<(usize, f32)> = lg.iter().cloned().enumerate().collect();
                    top.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
                    eprintln!("[moeB l0 r0] logits top4 {:?} sel {:?} wt {:?}", &top[..4], ids, wt);
                }
                unsafe {
                    // measured crossover: grouped wins at m>=32 (few pairs per expert below
                    // that, and the empty-expert blocks + serial pair-chains cost more than the
                    // weight amortization saves — m=8 was 10% WORSE grouped)
                    let grouped = self.moet || (std::env::var("MOEG").as_deref() != Ok("0") && pairs >= 256);
                    let pi32 = pairs as i32;
                    if grouped {
                        // expert-major: counting-sort pairs by expert, then each weight read
                        // serves up to PBG pairs (the per-pair kernel cost 4x+ the distinct-
                        // expert traffic floor and lost every M>=8 cell)
                        let cap = self.mwork.len() as i32;
                        let ssh = (3 * e_n * 4) as u32;
                        self.stream.launch_builder(&self.f_moe_sort)
                            .arg(&self.meidx).arg(&mut self.mcnt).arg(&mut self.moff)
                            .arg(&mut self.mplist).arg(&mut self.mwork).arg(&mut self.mwn)
                            .arg(&pi32).arg(&ei).arg(&cap).arg(&tsh)
                            .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1),
                                                   shared_mem_bytes: ssh })?;
                    }
                    // Pick the tile width from the work available. See gemm_q4_moe_t64.
                    let (f_t, bn) = if pairs <= 128 { (&self.f_moe_t64, 64usize) }
                                    else { (&self.f_moe_t, 128usize) };
                    let w1 = self.layers[li].w1.as_ref().context("MoE w1")?;
                    if self.moet {
                        // TENSOR arm: one block per (128-column tile, expert); each walks that
                        // expert's pair list in 16-pair tiles, gathering rows into shared.
                        self.stream.launch_builder(f_t)
                            .arg(&w1.scales).arg(w1.quants_p.as_ref().unwrap())
                            .arg(&b.h).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&hsi)
                            .arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&xs1).arg(&xdivk).arg(&self.mwork).arg(&self.mwn)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(bn) as u32, 132, 1),
                                                   block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    } else if grouped {
                        self.stream.launch_builder(&self.f_gemv_moe_g)
                            .arg(&w1.scales).arg(&w1.quants)
                            .arg(&b.h).arg(&b.xs).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&xs1).arg(&sh0).arg(&xdivk)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(32) as u32, e_n as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    } else {
                        self.stream.launch_builder(&self.f_gemv_moe)
                            .arg(&w1.scales).arg(&w1.quants)
                            .arg(&b.h).arg(&b.xs).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&self.meidx).arg(&xs1).arg(&sh0).arg(&xdivk)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(32) as u32, pairs as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    }
                    self.stream.launch_builder(&self.f_silu_moe)
                        .arg(&self.meg).arg(&mut self.meu).arg(&mut self.mexs).arg(&mii)
                        .launch(LaunchConfig { grid_dim: (pairs as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    let w2 = self.layers[li].w2.as_ref().context("MoE w2")?;
                    let nb2 = (mi / 32) as i32;
                    if self.moet && self.moet_pf {
                        self.stream.launch_builder(f_t)
                            .arg(&w2.scales).arg(w2.quants_p.as_ref().unwrap())
                            .arg(&self.meu).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&mii)
                            .arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&mii).arg(&one_d).arg(&self.mwork).arg(&self.mwn)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(bn) as u32, 132, 1),
                                                   block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    } else if grouped {
                        self.stream.launch_builder(&self.f_gemv_moe_g)
                            .arg(&w2.scales).arg(&w2.quants)
                            .arg(&self.meu).arg(&self.mexs).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&mii).arg(&sh0).arg(&one_d)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(32) as u32, e_n as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    } else {
                        self.stream.launch_builder(&self.f_gemv_moe)
                            .arg(&w2.scales).arg(&w2.quants)
                            .arg(&self.meu).arg(&self.mexs).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&self.meidx).arg(&mii).arg(&sh0).arg(&one_d)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(32) as u32, pairs as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    }
                    let rowsi = m as i32;
                    self.stream.launch_builder(&self.f_combine_b)
                        .arg(&self.medn).arg(&self.mewt).arg(&mut b.h).arg(&hsi).arg(&ki).arg(&rowsi)
                        .launch(ecfg(m * hidden))?;
                    if shared {
                        let nhh = (m * hidden) as i32;
                        self.stream.launch_builder(&self.f_add_inplace)
                            .arg(&mut b.h).arg(&b.abuf).arg(&nhh).launch(ecfg(m * hidden))?;
                    }
                }
            } else if ml_arm {
                Self::gemm_q4_ml_raw(&self.stream, &self.f_q4ml, &self.f_gwksum,
                    self.mlpart.as_mut().unwrap(), &self.layers[li].gate, &b.hb16, &mut b.g,
                    m, inter, hidden)?;
                Self::gemm_q4_ml_raw(&self.stream, &self.f_q4ml, &self.f_gwksum,
                    self.mlpart.as_mut().unwrap(), &self.layers[li].up, &b.hb16, &mut b.u,
                    m, inter, hidden)?;
                unsafe {
                    self.stream.launch_builder(&self.f_silu)
                        .arg(&mut b.g).arg(&b.u).arg(&ni).arg(&act).launch(ecfg(m * inter))?;
                    self.stream.launch_builder(&self.f_tof16)
                        .arg(&b.g).arg(&mut b.g16).arg(&ni).launch(ecfg(m * inter))?;
                }
                Self::gemm_q4_ml_raw(&self.stream, &self.f_q4ml, &self.f_gwksum,
                    self.mlpart.as_mut().unwrap(), &self.layers[li].down, &b.g16, &mut b.h,
                    m, hidden, inter)?;
            } else if gemv_arm {
                std::mem::swap(&mut b.h, &mut b.abuf);
                Self::gemv_b(&self.stream, &self.f_gemv_b, &self.layers[li].gate, &b.h, &b.xs, &mut b.g, m, hidden, shb)?;
                Self::gemv_b(&self.stream, &self.f_gemv_b, &self.layers[li].up, &b.h, &b.xs, &mut b.u, m, hidden, shb)?;
                unsafe {
                    self.stream.launch_builder(&self.f_silu)
                        .arg(&mut b.g).arg(&b.u).arg(&ni).arg(&act).launch(ecfg(m * inter))?;
                    self.stream.launch_builder(&self.f_absmax_slots)
                        .arg(&b.g).arg(&mut b.xs).arg(&ii).arg(&ii).launch(one256)?;
                }
                Self::gemv_b(&self.stream, &self.f_gemv_b, &self.layers[li].down, &b.g, &b.xs, &mut b.h, m, inter, shb)?;
            } else if q4_arm {
                // ONE GEMM for gate|up (they share the input), then a fused silu that also emits the
                // bf16 the down projection wants: 3 kernels -> 2.
                Self::gemm_q4s(&self.stream, &self.f_q4v2, &self.f_q4red,
                    self.q4part.as_mut().unwrap(), &self.q4gu[li].1, &self.q4gu[li].0,
                    self.layers[li].gate.nblk, &b.hb16, &mut b.gu32, m, 2 * inter, hidden)?;
                unsafe {
                    self.stream.launch_builder(&self.f_silu_gu)
                        .arg(&b.gu32).arg(&mut self.xf16).arg(&ii).arg(&ni).arg(&act).launch(ecfg(m * inter))?;
                }
                Self::gemm_q4s(&self.stream, &self.f_q4v2, &self.f_q4red,
                    self.q4part.as_mut().unwrap(), &self.layers[li].down.scales, &self.q4p[li][4],
                    self.layers[li].down.nblk, &self.xf16, &mut b.h, m, hidden, inter)?;
            } else if fp8dec || fp8sw {
                // FP8 FFN: quantise the activation per step (per-tensor scale, device scale
                // pointers on the Lt descriptor, alpha = 1), one planned Lt matmul per GEMM.
                // fp8dw: producers already wrote e4m3 with the delayed scale — no chains at all.
                // fp8dec reads the producer-accumulated absmax; fp8sw (sandwich) runs two-stage.
                let nx = (m * hidden) as i64;
                if !fp8dw {
                    unsafe {
                        if fp8dec {
                            self.stream.launch_builder(&self.f_amax2s)
                                .arg(&mut self.fp8_amax).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
                        } else {
                            let g = 64i32;
                            self.stream.launch_builder(&self.f_amax_part)
                                .arg(&b.hb16).arg(&mut self.fp8_part).arg(&nx)
                                .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                            self.stream.launch_builder(&self.f_amax_fin)
                                .arg(&self.fp8_part).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx).arg(&g)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        }
                        self.stream.launch_builder(&self.f_tofp8)
                            .arg(&b.hb16).arg(&mut self.fp8_x).arg(&self.fp8_sxmax).arg(&nx)
                            .launch(ecfg((m * hidden).div_ceil(2)))?;
                    }
                }
                let (lt, ltws) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned(&self.stream, lt, &self.fp8_plans[li], ltws, 1 << 20,
                    &self.fp8_gu[li].0, &self.fp8_x, &mut b.gu16, 1.0)?;
                unsafe {
                    // a missing kernel arg is a HOST segfault: cuLaunchKernel copies params from
                    // the pointer array and reads one past its end
                    if fp8dw {
                        let mut slot = self.fp8_dsc.slice_mut((li * 4 + 3) * 3..(li * 4 + 3) * 3 + 3);
                        self.stream.launch_builder(&self.f_silu_fp8)
                            .arg(&b.gu16).arg(&mut self.fp8_x).arg(&ii).arg(&n8).arg(&act)
                            .arg(&mut slot).launch(ecfg(m * inter / 8))?;
                    } else if fp8dec {
                        self.stream.launch_builder(&self.f_silu_bf_ax)
                            .arg(&b.gu16).arg(&mut b.g16).arg(&ii).arg(&n8).arg(&act)
                            .arg(&mut self.fp8_amax).launch(ecfg(m * inter / 8))?;
                    } else {
                        self.stream.launch_builder(&self.f_silu_bf)
                            .arg(&b.gu16).arg(&mut b.g16).arg(&ii).arg(&n8).arg(&act)
                            .launch(ecfg(m * inter / 8))?;
                    }
                }
                // down in fp8 too — it is a third of the FFN's weight bytes
                let ndx = (m * inter) as i64;
                if !fp8dw {
                    unsafe {
                        if fp8dec {
                            self.stream.launch_builder(&self.f_amax2s)
                                .arg(&mut self.fp8_amax).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
                        } else {
                            let g = 64i32;
                            self.stream.launch_builder(&self.f_amax_part)
                                .arg(&b.g16).arg(&mut self.fp8_part).arg(&ndx)
                                .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                            self.stream.launch_builder(&self.f_amax_fin)
                                .arg(&self.fp8_part).arg(&mut self.fp8_sxmax).arg(&mut self.fp8_sx).arg(&g)
                                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        }
                        self.stream.launch_builder(&self.f_tofp8)
                            .arg(&b.g16).arg(&mut self.fp8_x).arg(&self.fp8_sxmax).arg(&ndx)
                            .launch(ecfg((m * inter).div_ceil(2)))?;
                    }
                }
                Self::fp8_matmul_planned_f32(&self.stream, lt, &self.fp8_plans_down[li], ltws, 1 << 20,
                    &self.fp8_down[li].0, &self.fp8_x, &mut b.h, 1.0)?;
            } else {
                Self::gemm_pre_c16(&self.stream, &self.blas, self.layers[li].f16_gu.as_ref().unwrap(), &b.hb16, &mut b.gu16, m, 2 * inter, hidden)?;
                unsafe {
                    self.stream.launch_builder(&self.f_silu_bf)
                        .arg(&b.gu16).arg(&mut b.g16).arg(&ii).arg(&n8).arg(&act).launch(ecfg(m * inter / 8))?;
                }
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans, (m, hidden, inter),
                    self.layers[li].f16_down.as_ref().unwrap(), &b.g16, &mut b.h,
                    m, hidden, inter)?;
            }
            mark!(6);
        }
        if fp8dw {
            let ns = (c.layers * 4) as i32;
            unsafe {
                self.stream.launch_builder(&self.f_sc_rot)
                    .arg(&mut self.fp8_dsc).arg(&ns)
                    .launch(LaunchConfig { grid_dim: (((c.layers * 4 + 255) / 256) as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
            }
        }
        if c.sandwich {
            // Final stage: the LAST layer's post-FFN norm + residual + layer_scalar folded with
            // the final norm, emitted bf16 straight into the LM-head GEMM.
            let last = c.layers - 1;
            let pls = self.layers[last].ls;
            unsafe {
                self.stream.launch_builder(&self.f_sandwich)
                    .arg(&mut b.x).arg(&b.h)
                    .arg(self.layers[last].post_ffw_norm.as_ref().unwrap())
                    .arg(&self.final_norm)
                    .arg(&mut b.hb16).arg(&hi).arg(&eps).arg(&pls).launch(one)?;
            }
            // final_logit_softcapping is monotonic (30·tanh(z/30)) — greedy argmax unaffected,
            // so serving skips it. Revisit if logprobs are ever exposed.
            Self::gemm_pre(&self.stream, &self.blas, self.f16_lm.as_ref().unwrap(),
                &b.hb16, &mut b.logits, m, self.cfg.vocab, hidden)?;
        } else {
        unsafe {
            self.stream.launch_builder(&self.f_rms_add_b)
                .arg(&mut b.x).arg(&b.h).arg(&self.final_norm).arg(&mut b.abuf).arg(&mut b.xs)
                .arg(&hi).arg(&eps).launch(one256)?;
        }
        std::mem::swap(&mut b.h, &mut b.abuf);
        if gemv_arm {
            Self::gemv_b(&self.stream, &self.f_gemv_b, &self.lm_head, &b.h, &b.xs, &mut b.logits, m, hidden, shb)?;
        } else {
            Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16, self.f16_lm.as_ref().unwrap(), &mut self.xf16, &b.h, &mut b.logits, m, self.cfg.vocab, hidden)?;
        }
        }
        let (v, mo) = (c.vocab as i32, bmax as i32);
        unsafe {
            self.stream.launch_builder(&self.f_argmax_b)
                .arg(&b.logits).arg(&mut b.tok).arg(&v)
                .launch(LaunchConfig { grid_dim: (m as u32,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
            self.stream.launch_builder(&self.f_advance_b)
                .arg(&mut b.pos).arg(&b.tok).arg(&mut b.out).arg(&mo).arg(&mi).arg(&b.active)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
        }
        mark!(7);
        if prof {
            let names = ["embed", "norm+qkv", "rope+quant", "attn_partial", "reduce",
                         "o-gemm", "ffn(all)", "tail(lm+argmax)"];
            let tot: f64 = tp.iter().sum();
            println!("== DECODE PROFILE (m={m}, sum {:.3} ms)", tot * 1e3);
            for (n2, v) in names.iter().zip(tp.iter()) {
                println!("   {n2:<16} {:>7.3} ms  {:>5.1}%", v * 1e3, v / tot * 100.0);
            }
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn gemv_b(stream: &Arc<CudaStream>, f: &CudaFunction, w: &Q4Dev, x: &CudaSlice<f32>,
              xs: &CudaSlice<f32>, y: &mut CudaSlice<f32>, m: usize, cols: usize, shb: u32) -> Result<()> {
        let (nb, rows, mi, ci) = (w.nblk as i32, w.rows as i32, m as i32, cols as i32);
        unsafe {
            stream.launch_builder(f)
                .arg(&w.scales).arg(&w.quants).arg(x).arg(xs).arg(y)
                .arg(&nb).arg(&rows).arg(&mi).arg(&ci)
                .launch(LaunchConfig { grid_dim: (w.rows.div_ceil(64) as u32, 1, 1), block_dim: (1024,1,1), shared_mem_bytes: shb })?;
        }
        Ok(())
    }

    pub fn batch_tokens(&self, seq: usize, from: usize, to: usize) -> Result<Vec<u32>> {
        let b = self.batch.as_ref().context("init_batch first")?;
        let base = seq * b.max_seq;
        Ok(self.stream.memcpy_dtov(&b.out.slice(base + from..base + to))?)
    }

    /// `Y[M, rows] = X[M, cols] * W^T` on TENSOR CORES: dequantise the q4_0 weight into the f16
    /// scratch, convert the activations, then one `cublasGemmEx` (f16 in, fp32 accumulate). Only
    /// worth it for M > 1 — the dequant is fixed cost amortised across the M columns.
    #[allow(clippy::too_many_arguments)]
    fn gemm_tc(
        stream: &Arc<CudaStream>,
        blas: &CudaBlas,
        f_dq: &CudaFunction,
        f_tof16: &CudaFunction,
        wf16: &mut CudaSlice<f16>,
        xf16: &mut CudaSlice<f16>,
        w: &Q4Dev,
        x: &CudaSlice<f32>,
        y: &mut CudaSlice<f32>,
        m: usize,
        cols: usize,
    ) -> Result<()> {
        let (nblk, rows) = (w.nblk, w.rows);
        let (nbi, ri, ni) = (nblk as i32, rows as i32, (m * cols) as i32);
        let blocks = (rows * nblk).div_ceil(256);
        // wf16 is lazily grown (it is vocab*hidden-sized on big models and the batched server
        // never comes through here)
        if wf16.len() < rows * nblk * 32 {
            *wf16 = stream.alloc_zeros::<f16>(rows * nblk * 32)?;
        }
        unsafe {
            stream.launch_builder(f_dq)
                .arg(&w.scales).arg(&w.quants).arg(&mut *wf16).arg(&nbi).arg(&ri)
                .launch(LaunchConfig { grid_dim: (blocks as u32, 1, 1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
            stream.launch_builder(f_tof16)
                .arg(x).arg(&mut *xf16).arg(&ni).launch(ecfg(m * cols))?;
        }
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pw, _a) = wf16.device_ptr(stream);
        let (px, _b) = xf16.device_ptr(stream);
        let (py, _c) = y.device_ptr_mut(stream);
        unsafe {
            blas::gemm_ex(
                *blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                rows as i32, m as i32, cols as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16F, cols as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16F, cols as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, rows as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            ).map_err(|e| anyhow::anyhow!("gemm_ex: {e:?}"))?;
        }
        Ok(())
    }

    /// `Y[M, rows] = X[M, cols] * W^T` with W already f16 resident — one cublasGemmEx on tensor
    /// cores, no dequant in the step.
    /// `Y[M, rows] = X[M, cols] * W^T` all-f32 (CUBLAS_COMPUTE_32F, no TF32): the MoE router.
    /// bf16 routing logits FLIP top-k membership at near-ties — batch diverged from HF greedy at
    /// token 5 with a bf16 router against the f32 seq arm's 14/16 — so the router stays f32.
    fn gemm_f32(
        stream: &Arc<CudaStream>, blas: &CudaBlas,
        w: &CudaSlice<f32>, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pw, _a) = w.device_ptr(stream);
        let (px, _b) = x.device_ptr(stream);
        let (py, _c) = y.device_ptr_mut(stream);
        unsafe {
            blas::gemm_ex(
                *blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                rows as i32, m as i32, cols as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, cols as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, cols as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, rows as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            ).map_err(|e| anyhow::anyhow!("gemm_ex f32: {e:?}"))?;
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn gemm_bf16(
        stream: &Arc<CudaStream>, blas: &CudaBlas, f_tobf16: &CudaFunction,
        w: &CudaSlice<u16>, xbf: &mut CudaSlice<u16>,
        x: &CudaSlice<f32>, y: &mut CudaSlice<f32>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let ni = (m * cols) as i32;
        unsafe {
            // one conversion, no scaling: bf16 carries f32's exponent range
            stream.launch_builder(f_tobf16)
                .arg(x).arg(&mut *xbf).arg(&ni).launch(ecfg(m * cols))?;
        }
        // Scope the raw-pointer guards so `y` is free again for the rescale launch below.
        {
            let (alpha, beta) = (1.0f32, 0.0f32);
            let (pw, _a) = w.device_ptr(stream);
            let (px, _b) = xbf.device_ptr(stream);
            let (py, _c) = y.device_ptr_mut(stream);
            unsafe {
                blas::gemm_ex(
                    *blas.handle(),
                    blas_sys::cublasOperation_t::CUBLAS_OP_T,
                    blas_sys::cublasOperation_t::CUBLAS_OP_N,
                    rows as i32, m as i32, cols as i32,
                    (&alpha) as *const f32 as *const _,
                    pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                    px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                    (&beta) as *const f32 as *const _,
                    py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, rows as i32,
                    blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                    blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
                ).map_err(|e| anyhow::anyhow!("gemm_ex: {e:?}"))?;
            }
        }
        Ok(())
    }

    /// As `gemm_bf16` but the OUTPUT is bf16: for FFN interiors, where an f32 [T, inter] result
    /// would be written, read back and converted again for the next GEMM.
    #[allow(clippy::too_many_arguments)]
    fn gemm_bf16_c16(
        stream: &Arc<CudaStream>, blas: &CudaBlas, f_tobf16: &CudaFunction,
        w: &CudaSlice<u16>, xbf: &mut CudaSlice<u16>,
        x: &CudaSlice<f32>, y: &mut CudaSlice<u16>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let ni = (m * cols) as i32;
        unsafe {
            stream.launch_builder(f_tobf16).arg(x).arg(&mut *xbf).arg(&ni).launch(ecfg(m * cols))?;
        }
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pw, _a) = w.device_ptr(stream);
        let (px, _b) = xbf.device_ptr(stream);
        let (py, _c) = y.device_ptr_mut(stream);
        unsafe {
            blas::gemm_ex(
                *blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                rows as i32, m as i32, cols as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, rows as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            ).map_err(|e| anyhow::anyhow!("gemm_ex c16: {e:?}"))?;
        }
        Ok(())
    }

    /// bf16 in AND out: the FFN's gate|up, whose input is the bf16 rmsnorm and whose output feeds
    /// `silu_mul_bf`. No conversion pass on either side.
    #[allow(clippy::too_many_arguments)]
    fn gemm_pre_c16(
        stream: &Arc<CudaStream>, blas: &CudaBlas,
        w: &CudaSlice<u16>, xbf: &CudaSlice<u16>, y: &mut CudaSlice<u16>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pw, _a) = w.device_ptr(stream);
        let (px, _b) = xbf.device_ptr(stream);
        let (py, _c) = y.device_ptr_mut(stream);
        unsafe {
            blas::gemm_ex(
                *blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                rows as i32, m as i32, cols as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, rows as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            ).map_err(|e| anyhow::anyhow!("gemm_ex pre_c16: {e:?}"))?;
        }
        Ok(())
    }

    /// Input ALREADY bf16 (it came from `silu_mul_bf`), f32 output for the residual stream.
    #[allow(clippy::too_many_arguments)]
    fn gemm_pre(
        stream: &Arc<CudaStream>, blas: &CudaBlas,
        w: &CudaSlice<u16>, xbf: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pw, _a) = w.device_ptr(stream);
        let (px, _b) = xbf.device_ptr(stream);
        let (py, _c) = y.device_ptr_mut(stream);
        unsafe {
            blas::gemm_ex(
                *blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                rows as i32, m as i32, cols as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                px as *const std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_16BF, cols as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void, blas_sys::cudaDataType_t::CUDA_R_32F, rows as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            ).map_err(|e| anyhow::anyhow!("gemm_ex pre: {e:?}"))?;
        }
        Ok(())
    }

    /// Logits of one batch slot after the last step — for deciding whether a divergence between
    /// two arms is a near-tie flipped by rounding or a real defect.
    /// DEBUG: current device decode position.
    pub fn peek_pos(&self) -> Result<i32> {
        Ok(self.stream.memcpy_dtov(&self.d_pos)?[0])
    }

    /// The single-stream path's last logits (for the seq-vs-batch numeric bisect).
    pub fn seq_logits(&self) -> Result<Vec<f32>> {
        Ok(self.stream.memcpy_dtov(&self.logits)?)
    }

    pub fn batch_logits(&self, seq: usize) -> Result<Vec<f32>> {
        let b = self.batch.as_ref().context("init_batch first")?;
        let v = self.cfg.vocab;
        Ok(self.stream.memcpy_dtov(&b.logits.slice(seq * v..(seq + 1) * v))?)
    }

    /// Enable/disable the q4 decode arm at runtime (the permuted weights stay built).
    pub fn set_q4_arm(&mut self, on: bool) { self.q4_on = on; }

    /// Whether the sequential (SEQ1) M=1 decode path can run this model: it multiplies the
    /// ORIGINAL q4 quants (dead after Q4FREE) and has no sandwich/sliding arms (Gemma4).
    pub fn seq_capable(&self) -> bool { !self.q4free && !self.cfg.sandwich && !self.fp8kv }

    /// Whether wave (batched multi-request) prefill can run this model.
    pub fn wave_capable(&self) -> bool { self.cfg.num_experts == 0 }

    /// Move a live slot's whole decode state (KV slab up to `used` positions, token, position,
    /// active flag, token log) into another slot. The server compacts with this so the batch
    /// width tracks the LIVE COUNT: under equal budgets retirement is arrival-ordered, so
    /// without compaction the highest-ever slot pins the width for the entire drain and every
    /// step pays for sequences that already left. ~0.5 ms of dtod against a 2x-wide step.
    /// Caller must ensure the stream is quiescent for `from`/`to` (loop top, post-sync).
    pub fn move_slot(&mut self, from: usize, to: usize, used: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(from < b.mmax && to < b.mmax && from != to, "bad move");
        let n = used.min(b.max_seq);
        let stream = self.stream.clone();
        unsafe {
            let cu = stream.cu_stream();
            // element size: 2 B (bf16) or 1 B (e4m3, FP8KV=1) — slab offsets scale with it.
            // Head-major layout: a slot's block is n_kv planes of max_seq rows, so the move is
            // one copy of the first `n` rows PER PLANE instead of one block copy.
            let esz = if self.fp8kv { 1usize } else { 2 };
            for li in 0..self.cfg.layers {
                let (lhd, lnkv, _, _) = self.cfg.lgeom(li);
                let bytes = n * lhd * esz;
                for buf in [&mut b.kc[li], &mut b.vc[li]] {
                    let (p, _g) = buf.device_ptr_mut(&stream);
                    for kvh in 0..lnkv {
                        cudarc::driver::result::memcpy_dtod_async(
                            p + (((to * lnkv + kvh) * b.max_seq) * lhd * esz) as u64,
                            p + (((from * lnkv + kvh) * b.max_seq) * lhd * esz) as u64,
                            bytes, cu)?;
                    }
                }
            }
            let (pt, _g1) = b.tok.device_ptr_mut(&stream);
            cudarc::driver::result::memcpy_dtod_async(pt + to as u64 * 4, pt + from as u64 * 4, 4, cu)?;
            let (pp, _g2) = b.pos.device_ptr_mut(&stream);
            cudarc::driver::result::memcpy_dtod_async(pp + to as u64 * 4, pp + from as u64 * 4, 4, cu)?;
            let (pa, _g3) = b.active.device_ptr_mut(&stream);
            cudarc::driver::result::memcpy_dtod_async(pa + to as u64 * 4, pa + from as u64 * 4, 4, cu)?;
            let (po, _g4) = b.out.device_ptr_mut(&stream);
            cudarc::driver::result::memcpy_dtod_async(
                po + (to * b.max_seq * 4) as u64, po + (from * b.max_seq * 4) as u64,
                n * 4, cu)?;
        }
        // Host bookkeeping only for the vacated slot: it now sits at/above the compacted width,
        // so no kernel reads its device state (every _b kernel is m-bounded).
        b.occupied[to] = true;
        b.occupied[from] = false;
        Ok(())
    }

    /// Switch the attention implementation at runtime, so the two can be diffed on real weights.
    pub fn set_fa(&mut self, on: bool) { self.fa = on; }

    /// Prefill and return the LM-head logits. This is the end-to-end numeric gate for attention:
    /// flash and score-matrix paths must agree on all 151936 values, not just the argmax.
    pub fn prefill_logits(&mut self, tokens: &[u32]) -> Result<Vec<f32>> {
        self.prefill(tokens)?;
        self.sync()?;
        let p = self.pf.as_ref().unwrap();
        Ok(self.stream.memcpy_dtov(&p.logits)?)
    }

    /// C[m, rows] = X[m, cols] . W^T with W read as q4_0 — no bf16 mirror anywhere.
    /// Same operand and output layout as `gemm_pre`, so the two are directly diffable.
    #[allow(clippy::too_many_arguments)]
    fn gemm_q4(stream: &Arc<CudaStream>, f: &CudaFunction, w: &Q4Dev,
               xbf: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
               m: usize, rows: usize, cols: usize) -> Result<()> {
        let (mi, ri, ci, nb) = (m as i32, rows as i32, cols as i32, w.nblk as i32);
        unsafe {
            stream.launch_builder(f)
                .arg(&w.scales).arg(&w.quants).arg(xbf).arg(y)
                .arg(&mi).arg(&ri).arg(&ci).arg(&nb)
                .launch(LaunchConfig {
                    grid_dim: (m.div_ceil(64) as u32, rows.div_ceil(64) as u32, 1),
                    block_dim: (128, 1, 1),
                    shared_mem_bytes: 0,
                })?;
        }
        Ok(())
    }

    /// Barrier-free q4 GEMM with split-K, then a deterministic reduction.
    /// Same operand and output layout as `gemm_pre`.
    #[allow(clippy::too_many_arguments)]
    /// FP8 (e4m3) matmul through cuBLASLt: D[m x rows] = W[rows x cols]^T . X[m x cols], TN.
    ///
    /// `cublasGemmEx` cannot do FP8 at all (it returns INVALID_VALUE), and cudarc's safe `Matmul`
    /// trait only impls f32/f16/bf16, so this is raw `cublaslt::sys`. The sequence is a direct port
    /// of `bench/probes/fp8_vs_bf16.cu`, which measured **1064 TFLOP/s against bf16's 627** on the
    /// gate|up shape — that 1.70x is the entire reason this path exists.
    ///
    /// `alpha` carries the dequantisation: FP8 stores w/s_w and x/s_x, so the product needs s_w*s_x.
    /// Build the descriptors ONCE for a shape. cuBLASLt's descriptors must outlive the async
    /// matmul; creating and destroying them around each call is a use-after-free (it core-dumps).
    #[allow(clippy::too_many_arguments)]
    #[allow(clippy::too_many_arguments)]
    fn fp8_plan(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t,
        m: usize, rows: usize, cols: usize, wss: usize,
        sa: &CudaSlice<f32>, sbp: u64, d_f32: bool,
        picked: Option<lt_sys::cublasLtMatmulAlgo_t>,
    ) -> Result<Fp8Plan> {
        unsafe {
            let mut op: lt_sys::cublasLtMatmulDesc_t = std::ptr::null_mut();
            let st = lt_sys::cublasLtMatmulDescCreate(
                &mut op, lt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                lt_sys::cudaDataType_t::CUDA_R_32F);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt desc: {st:?}");
            let ta = blas_sys::cublasOperation_t::CUBLAS_OP_T;
            let tb = blas_sys::cublasOperation_t::CUBLAS_OP_N;
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
                &ta as *const _ as *const _, std::mem::size_of_val(&ta));
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
                &tb as *const _ as *const _, std::mem::size_of_val(&tb));
            // A/B scale pointers: cuBLASLt computes D = alpha * (sA*A) * (sB*B), so with A holding
            // w*448/s_w and B holding x*448/s_x, sA = s_w/448 and sB = s_x/448 recover w*x exactly —
            // and because these are DEVICE pointers, the values can change per step without touching
            // the descriptor. That is what makes a fixed activation scale unnecessary.
            let (pa, _ga) = sa.device_ptr(stream);
            let pav = pa as *const std::ffi::c_void;
            let pbv = sbp as *const std::ffi::c_void;
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
                &pav as *const _ as *const _, std::mem::size_of_val(&pav));
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
                &pbv as *const _ as *const _, std::mem::size_of_val(&pbv));

            let (mut la, mut lb, mut ld) = (std::ptr::null_mut(), std::ptr::null_mut(), std::ptr::null_mut());
            let e4 = lt_sys::cudaDataType_t::CUDA_R_8F_E4M3;
            lt_sys::cublasLtMatrixLayoutCreate(&mut la, e4, cols as u64, rows as u64, cols as i64);
            lt_sys::cublasLtMatrixLayoutCreate(&mut lb, e4, cols as u64, m as u64, cols as i64);
            // D in the CONSUMER's dtype: the down projection feeds the f32 residual stream, and
            // taking bf16 there costs a whole extra pass over m*hidden per layer.
            let dt = if d_f32 { lt_sys::cudaDataType_t::CUDA_R_32F }
                     else { lt_sys::cudaDataType_t::CUDA_R_16BF };
            lt_sys::cublasLtMatrixLayoutCreate(&mut ld, dt, rows as u64, m as u64, rows as i64);
            let mut pref: lt_sys::cublasLtMatmulPreference_t = std::ptr::null_mut();
            lt_sys::cublasLtMatmulPreferenceCreate(&mut pref);
            let wsb = wss as u64;
            lt_sys::cublasLtMatmulPreferenceSetAttribute(
                pref, lt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                &wsb as *const _ as *const _, std::mem::size_of_val(&wsb));
            let mut heur: [lt_sys::cublasLtMatmulHeuristicResult_t; 8] = std::mem::zeroed();
            let mut found = 0i32;
            let st = lt_sys::cublasLtMatmulAlgoGetHeuristic(lt, op, la, lb, ld, ld, pref, 8,
                                                            heur.as_mut_ptr(), &mut found);
            lt_sys::cublasLtMatmulPreferenceDestroy(pref);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS && found > 0,
                            "no FP8 algo for [{m} x {cols}] x [{cols} x {rows}]: {st:?}");
            let algo = match picked {
                Some(a) => a,
                None => heur[0].algo,
            };
            Ok(Fp8Plan { op, la, lb, ld, algo, cand: heur[..found as usize].iter().map(|h| h.algo).collect() })
        }
    }

    /// Time this plan's heuristic candidates on the given operands and lock in the fastest.
    /// heuristic[0] is 10-20% off the best kernel for skinny-m fp8 shapes, and every layer of a
    /// width shares the shape — so callers time ONCE per (m, rows, cols) and reuse the winner.
    #[allow(clippy::too_many_arguments)]
    fn fp8_tune_plan(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &mut Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w8: &CudaSlice<u8>, x8: &CudaSlice<u8>, dsz: usize, d_f32: bool,
    ) -> Result<lt_sys::cublasLtMatmulAlgo_t> {
        let mut d32;
        let mut d16;
        let cands = p.cand.clone();
        let mut best = (f64::INFINITY, p.algo);
        if d_f32 { d32 = Some(stream.alloc_zeros::<f32>(dsz)?); d16 = None; }
        else { d32 = None; d16 = Some(stream.alloc_zeros::<u16>(dsz)?); }
        let mut timed: Vec<(f64, lt_sys::cublasLtMatmulAlgo_t)> = Vec::new();
        for a in cands {
            p.algo = a;
            let run = |p: &Fp8Plan, ws: &mut CudaSlice<u8>,
                       d32: &mut Option<CudaSlice<f32>>, d16: &mut Option<CudaSlice<u16>>| {
                if let Some(d) = d32.as_mut() {
                    Self::fp8_matmul_planned_f32(stream, lt, p, ws, wss, w8, x8, d, 1.0)
                } else {
                    Self::fp8_matmul_planned(stream, lt, p, ws, wss, w8, x8, d16.as_mut().unwrap(), 1.0)
                }
            };
            // warm once (also filters algos that fail at run time), then time 3 reps
            if run(p, ws, &mut d32, &mut d16).is_err() {
                continue;
            }
            stream.synchronize()?;
            let t0 = std::time::Instant::now();
            for _ in 0..3 {
                run(p, ws, &mut d32, &mut d16)?;
            }
            stream.synchronize()?;
            let dt = t0.elapsed().as_secs_f64();
            timed.push((dt, a));
            if dt < best.0 {
                best = (dt, a);
            }
        }
        // If NO candidate even runs (delayed-scale descriptors can invalidate every heuristic
        // suggestion), say so instead of returning the UNVERIFIED initial algo — that silent
        // fallback is exactly what let FP8DW ship plans that NOT_SUPPORTED at the first step.
        anyhow::ensure!(!timed.is_empty(), "no Lt algo executes for this shape/descriptor");
        // CAPTURE-SAFETY: some Lt algos run eagerly but return NOT_SUPPORTED inside CUDA graph
        // capture — which ABORTS the decode-graph capture and kills the engine thread. Walk the
        // candidates fastest-first and keep the first that survives a throwaway trial capture.
        timed.sort_by(|a2, b2| a2.0.partial_cmp(&b2.0).unwrap());
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        for (_, a) in &timed {
            p.algo = *a;
            stream.synchronize()?;
            stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
            let r = if let Some(d) = d32.as_mut() {
                Self::fp8_matmul_planned_f32(stream, lt, p, ws, wss, w8, x8, d, 1.0)
            } else {
                Self::fp8_matmul_planned(stream, lt, p, ws, wss, w8, x8, d16.as_mut().unwrap(), 1.0)
            };
            let g = stream.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
            if r.is_ok() && g.is_ok() {
                return Ok(*a);
            }
        }
        p.algo = best.1;
        Ok(best.1)
    }

    /// Quantise one bf16 mirror to an e4m3 copy with its Lt-ready per-tensor scale.
    fn fp8_mirror(
        stream: &Arc<CudaStream>, f_absmax: &CudaFunction, f_tofp8: &CudaFunction,
        f_mkscale: &CudaFunction, src: &CudaSlice<u16>, n: usize,
    ) -> Result<(CudaSlice<u8>, CudaSlice<f32>)> {
        let ni = n as i64;
        let mut sc = stream.alloc_zeros::<f32>(1)?;
        let mut q = stream.alloc_zeros::<u8>(n)?;
        unsafe {
            stream.launch_builder(f_absmax).arg(src).arg(&mut sc).arg(&ni)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
            stream.launch_builder(f_tofp8).arg(src).arg(&mut q).arg(&sc).arg(&ni)
                .launch(ecfg(n.div_ceil(2)))?;
        }
        let mut swd = stream.alloc_zeros::<f32>(1)?;
        unsafe {
            stream.launch_builder(f_mkscale).arg(&sc).arg(&mut swd)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 })?;
        }
        Ok((q, swd))
    }

    /// (Re)build the four per-layer FP8 DECODE plan sets for batch width `m`, timing Lt algo
    /// candidates once per (m, rows, cols) shape and reusing the winner across layers. The
    /// tuning launches real matmuls with syncs — callers must run this OUTSIDE graph capture.
    fn fp8_decode_plans(&mut self, m: usize) -> Result<()> {
        if self.fp8_gu.is_empty() || self.fp8_plan_m == m {
            return Ok(());
        }
        let (hidden, inter, n) = (self.cfg.hidden, self.cfg.intermediate, self.cfg.layers);
        let (dscp, _gd) = self.fp8_dsc.device_ptr(&self.stream);
        let (sxp, _gs) = self.fp8_sx.device_ptr(&self.stream);
        // Delayed-scale (FP8DW) plans can have NO executable Lt algo for a shape — retry the
        // whole build with standard scale pointers rather than shipping half-dw plans (earlier
        // dw plans would read producer slots the degraded step never fills).
        'build: loop {
            let use_dw = std::env::var("FP8DW").as_deref() == Ok("1") && !self.fp8dw_dead;
            let (mut pg, mut pq, mut po, mut pd) =
                (Vec::with_capacity(n), Vec::with_capacity(n), Vec::with_capacity(n), Vec::with_capacity(n));
            for l2 in 0..n {
                let shapes: [(usize, usize, bool, usize); 4] = [
                    (2 * inter, hidden, false, 2),
                    (self.layers[l2].qkv.rows, hidden, true, 0),
                    (hidden, self.layers[l2].o.nblk * 32, true, 1),
                    (hidden, inter, true, 3),
                ];
                for (si, &(rows, cols, d32, site)) in shapes.iter().enumerate() {
                    let (sw, w8) = match si {
                        0 => (&self.fp8_gu[l2].1, &self.fp8_gu[l2].0),
                        1 => (&self.fp8_qkv[l2].1, &self.fp8_qkv[l2].0),
                        2 => (&self.fp8_o[l2].1, &self.fp8_o[l2].0),
                        _ => (&self.fp8_down[l2].1, &self.fp8_down[l2].0),
                    };
                    // delayed-scale: B scale reads the (layer, site) slot's amax/448 slot
                    let sbp = if use_dw { dscp + (((l2 * 4 + site) * 3 + 1) * 4) as u64 } else { sxp };
                    let key = (m, rows, cols, use_dw);
                    let cached = self.fp8_algo_cache.get(&key).copied();
                    let mut plan = Self::fp8_plan(&self.stream, self.lt, m, rows, cols, 1 << 20,
                                                  sw, sbp, d32, cached)?;
                    if cached.is_none() {
                        match Self::fp8_tune_plan(&self.stream, self.lt, &mut plan, &mut self.ltws,
                                                  1 << 20, w8, &self.fp8_x, m * rows, d32) {
                            Ok(a) => { self.fp8_algo_cache.insert(key, a); }
                            Err(e) if use_dw => {
                                eprintln!("FP8DW disabled for this run: {rows}x{cols} at m={m} \
                                           has no executable delayed-scale algo ({e})");
                                self.fp8dw_dead = true;
                                continue 'build;
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    match si { 0 => pg.push(plan), 1 => pq.push(plan), 2 => po.push(plan), _ => pd.push(plan) }
                }
            }
            self.fp8_plans = pg;
            self.fp8_plans_qkv = pq;
            self.fp8_plans_o = po;
            self.fp8_plans_down = pd;
            self.fp8_plan_m = m;
            return Ok(());
        }
    }

    /// Tuned bf16 decode GEMM: uses this width's tuned Lt plan when one exists, else the
    /// cublasGemmEx heuristic path (gemm_pre). Same math, same operands.
    /// NEVER take `&mut self` here and NEVER clone the weight slice to satisfy the borrow
    /// checker: CudaSlice::clone() DEEP COPIES ON DEVICE, so a clone per layer per step copies
    /// the whole model every step (Llama M=64 measured 5315 -> 3542 tok/s before this was
    /// caught). Disjoint fields are passed in explicitly instead.
    #[allow(clippy::too_many_arguments)]
    fn gemm_pre_t(stream: &Arc<CudaStream>, blas: &CudaBlas, lt: lt_sys::cublasLtHandle_t,
                  ltws: &mut CudaSlice<u8>,
                  plans: &std::collections::HashMap<(usize, usize, usize), Fp8Plan>,
                  key: (usize, usize, usize),
                  w: &CudaSlice<u16>, x: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
                  m: usize, rows: usize, cols: usize) -> Result<()> {
        if let Some(p) = plans.get(&key) {
            return Self::bf16_matmul_planned(stream, lt, p, ltws, 1 << 20, w, x, y);
        }
        Self::gemm_pre(stream, blas, w, x, y, m, rows, cols)
    }

    /// Build+tune bf16 decode plans for width `m` (the four projection shapes). Idempotent per
    /// shape; runs OUTSIDE graph capture because tuning launches synced matmuls.
    /// Same as bf16_decode_plans but for a prefill row count (adds the bf16-D gate|up shape).
    fn bf16_prefill_plans(&mut self, t: usize) -> Result<()> {
        if std::env::var("BF16LT").as_deref() == Ok("0") || self.layers[0].f16_qkv.is_none()
            || self.cfg.num_experts > 0 {
            return Ok(());
        }
        let (hidden, inter) = (self.cfg.hidden, self.cfg.intermediate);
        let l0 = &self.layers[0];
        let shapes: [(usize, usize, bool, usize); 4] = [
            (l0.qkv.rows, hidden, true, 0),
            (hidden, l0.o.nblk * 32, true, 1),
            (2 * inter, hidden, false, 2),
            (hidden, inter, true, 3),
        ];
        for &(rows, cols, d32, si) in shapes.iter() {
            let key = (t, rows, cols);
            if self.bf16_plans.contains_key(&key) { continue; }
            let w = match si {
                0 => self.layers[0].f16_qkv.as_ref(),
                1 => self.layers[0].f16_o.as_ref(),
                2 => self.layers[0].f16_gu.as_ref(),
                _ => self.layers[0].f16_down.as_ref(),
            }.context("bf16 mirrors")?;
            let mut plan = Self::bf16_plan(&self.stream, self.lt, t, rows, cols, 1 << 20, d32, None)?;
            let xs = Self::tune_operand(&self.stream, t * cols)?;
            let (lt, ltws) = (self.lt, &mut self.ltws);
            let a = Self::bf16_tune_plan(&self.stream, lt, &mut plan, ltws, 1 << 20, w, &xs, t * rows)?;
            plan.algo = a;
            self.bf16_plans.insert(key, plan);
        }
        Ok(())
    }

    fn bf16_decode_plans(&mut self, m: usize) -> Result<()> {
        // Tuned Lt plans beat cublasGemmEx's per-call heuristic on the decode shapes:
        // Llama-8B bf16 M=64 5401 -> 5501 (+1.9%), M=32 3037 -> 3073 (+1.2%), clean A/B on one
        // binary. (An earlier run of this experiment "showed" -33%; that was a CudaSlice::clone()
        // in the call path deep-copying the weights on device every step, not the plans.)
        if std::env::var("BF16LT").as_deref() == Ok("0") || self.layers[0].f16_qkv.is_none() {
            return Ok(());
        }
        let (hidden, inter) = (self.cfg.hidden, self.cfg.intermediate);
        let l0 = &self.layers[0];
        let shapes: [(usize, usize, bool); 4] = [
            (l0.qkv.rows, hidden, true),
            (hidden, l0.o.nblk * 32, true),
            (2 * inter, hidden, false),
            (hidden, inter, true),
        ];
        for (si, &(rows, cols, d32)) in shapes.iter().enumerate() {
            // MoE models have no dense gate|up/down mirrors — layers[0].f16_gu is a dummy, and
            // tuning a [2*inter x hidden] matmul against it is a huge OOB read that the
            // sanitizer counts by the thousands and that FAULTS at some allocation layouts
            // (this was the entire "batched MoE dies at slots=64" crash: nvjet illegal address
            // from inside capture_batch's plan tuning).
            if self.cfg.num_experts > 0 && si >= 2 { continue; }
            let key = (m, rows, cols);
            if self.bf16_plans.contains_key(&key) { continue; }
            let w = match si {
                0 => self.layers[0].f16_qkv.as_ref(),
                1 => self.layers[0].f16_o.as_ref(),
                2 => self.layers[0].f16_gu.as_ref(),
                _ => self.layers[0].f16_down.as_ref(),
            }.context("bf16 mirrors")?.clone();
            let mut plan = Self::bf16_plan(&self.stream, self.lt, m, rows, cols, 1 << 20, d32, None)?;
            let xs = Self::tune_operand(&self.stream, m * cols)?;
            let a = Self::bf16_tune_plan(&self.stream, self.lt, &mut plan, &mut self.ltws,
                                         1 << 20, &w, &xs, m * rows)?;
            plan.algo = a;
            self.bf16_plans.insert(key, plan);
        }
        Ok(())
    }

    /// ATTN_DUMP=1: fingerprint the given attention-output rows (bf16 [rows, nh*hd]) so the
    /// chunked path and the FA path can be diffed numerically instead of by reading code.
    fn dump_attn(&self, tag: &str, a16: &CudaSlice<u16>, rows: &[usize], width: usize) -> Result<()> {
        for &r in rows {
            let v: Vec<u16> = self.stream.memcpy_dtov(&a16.slice(r * width..(r + 1) * width))?;
            let f: Vec<f32> = v.iter().map(|&b| f32::from_bits((b as u32) << 16)).collect();
            let sum: f32 = f.iter().sum();
            let amax = f.iter().fold(0f32, |m, x| m.max(x.abs()));
            eprintln!("[attn {tag}] row {r}: sum {sum:.5} absmax {amax:.5} first4 {:?}",
                      &f[..4].iter().map(|x| (x * 1e4).round() / 1e4).collect::<Vec<_>>());
        }
        Ok(())
    }

    /// A bf16 cuBLASLt plan: same D = W^T X shape as gemm_pre, but with an explicit algo so the
    /// decode widths can be TUNED. cublasGemmEx picks a heuristic per call; on skinny-m shapes
    /// that heuristic left ~5% on the table for fp8, and the bf16 decode GEMMs are the same
    /// shape family.
    fn bf16_plan(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t,
        m: usize, rows: usize, cols: usize, wss: usize, d_f32: bool,
        picked: Option<lt_sys::cublasLtMatmulAlgo_t>,
    ) -> Result<Fp8Plan> {
        let _ = stream;
        unsafe {
            let mut op: lt_sys::cublasLtMatmulDesc_t = std::ptr::null_mut();
            let st = lt_sys::cublasLtMatmulDescCreate(
                &mut op, lt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                lt_sys::cudaDataType_t::CUDA_R_32F);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt desc bf16: {st:?}");
            let ta = blas_sys::cublasOperation_t::CUBLAS_OP_T;
            let tb = blas_sys::cublasOperation_t::CUBLAS_OP_N;
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
                &ta as *const _ as *const _, std::mem::size_of_val(&ta));
            lt_sys::cublasLtMatmulDescSetAttribute(
                op, lt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
                &tb as *const _ as *const _, std::mem::size_of_val(&tb));
            let (mut la, mut lb, mut ld) = (std::ptr::null_mut(), std::ptr::null_mut(), std::ptr::null_mut());
            let bf = lt_sys::cudaDataType_t::CUDA_R_16BF;
            lt_sys::cublasLtMatrixLayoutCreate(&mut la, bf, cols as u64, rows as u64, cols as i64);
            lt_sys::cublasLtMatrixLayoutCreate(&mut lb, bf, cols as u64, m as u64, cols as i64);
            let dt = if d_f32 { lt_sys::cudaDataType_t::CUDA_R_32F }
                     else { lt_sys::cudaDataType_t::CUDA_R_16BF };
            lt_sys::cublasLtMatrixLayoutCreate(&mut ld, dt, rows as u64, m as u64, rows as i64);
            let mut pref: lt_sys::cublasLtMatmulPreference_t = std::ptr::null_mut();
            lt_sys::cublasLtMatmulPreferenceCreate(&mut pref);
            let wsb = wss as u64;
            lt_sys::cublasLtMatmulPreferenceSetAttribute(
                pref, lt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                &wsb as *const _ as *const _, std::mem::size_of_val(&wsb));
            let mut heur: [lt_sys::cublasLtMatmulHeuristicResult_t; 8] = std::mem::zeroed();
            let mut found = 0i32;
            let st = lt_sys::cublasLtMatmulAlgoGetHeuristic(lt, op, la, lb, ld, ld, pref, 8,
                                                            heur.as_mut_ptr(), &mut found);
            lt_sys::cublasLtMatmulPreferenceDestroy(pref);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS && found > 0,
                            "no bf16 algo for [{m} x {cols}] x [{cols} x {rows}]: {st:?}");
            let algo = picked.unwrap_or(heur[0].algo);
            Ok(Fp8Plan { op, la, lb, ld, algo, cand: heur[..found as usize].iter().map(|h| h.algo).collect() })
        }
    }

    /// Tuning operands must look like REAL activations. Timing kernels on an all-zero buffer
    /// measures a machine that never denormalizes, never saturates and compresses perfectly in
    /// L2 — the candidate ranking it produces is not the ranking the served shapes see.
    fn tune_operand(stream: &Arc<CudaStream>, n: usize) -> Result<CudaSlice<u16>> {
        let mut v = vec![0u16; n];
        let mut x = 0x2545_f491_4f6c_dd1du64;
        for e in v.iter_mut() {
            x ^= x << 13; x ^= x >> 7; x ^= x << 17;
            // bf16 in [-2, 2): exponent 126/127, random mantissa+sign
            let sign = ((x >> 40) & 1) as u16;
            let expo = 126u16 + ((x >> 41) & 1) as u16;
            *e = (sign << 15) | (expo << 7) | ((x >> 3) & 0x7f) as u16;
        }
        Ok(stream.memcpy_stod(&v)?)
    }

    /// Time bf16 plan candidates and keep the fastest that also survives a trial graph capture.
    #[allow(clippy::too_many_arguments)]
    fn bf16_tune_plan(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &mut Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w: &CudaSlice<u16>, x: &CudaSlice<u16>, dsz: usize,
    ) -> Result<lt_sys::cublasLtMatmulAlgo_t> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        let mut d = stream.alloc_zeros::<f32>(dsz)?;
        let cands = p.cand.clone();
        let mut timed: Vec<(f64, lt_sys::cublasLtMatmulAlgo_t)> = Vec::new();
        for a in cands {
            p.algo = a;
            if Self::bf16_matmul_planned(stream, lt, p, ws, wss, w, x, &mut d).is_err() { continue; }
            stream.synchronize()?;
            let t0 = std::time::Instant::now();
            for _ in 0..3 { Self::bf16_matmul_planned(stream, lt, p, ws, wss, w, x, &mut d)?; }
            stream.synchronize()?;
            timed.push((t0.elapsed().as_secs_f64(), a));
        }
        timed.sort_by(|a2, b2| a2.0.partial_cmp(&b2.0).unwrap());
        for (_, a) in &timed {
            p.algo = *a;
            stream.synchronize()?;
            stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
            let r = Self::bf16_matmul_planned(stream, lt, p, ws, wss, w, x, &mut d);
            let g = stream.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
            if r.is_ok() && g.is_ok() { return Ok(*a); }
        }
        Ok(p.algo)
    }

    #[allow(clippy::too_many_arguments)]
    fn bf16_matmul_planned_c16(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w: &CudaSlice<u16>, x: &CudaSlice<u16>, d: &mut CudaSlice<u16>,
    ) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        unsafe {
            let (pw, _a) = w.device_ptr(stream);
            let (px, _b) = x.device_ptr(stream);
            let (pd, _c) = d.device_ptr_mut(stream);
            let (pws, _w) = ws.device_ptr_mut(stream);
            let st = lt_sys::cublasLtMatmul(
                lt, p.op, &alpha as *const f32 as *const _,
                pw as *const _, p.la, px as *const _, p.lb,
                &beta as *const f32 as *const _,
                pd as *const _, p.ld, pd as *mut _, p.ld,
                &p.algo, pws as *mut _, wss, stream.cu_stream() as _);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt bf16 c16: {st:?}");
        }
        Ok(())
    }

    /// gemm_pre_c16 with the tuned plan when one exists (gate|up writes bf16 D).
    #[allow(clippy::too_many_arguments)]
    fn gemm_pre_c16_t(stream: &Arc<CudaStream>, blas: &CudaBlas, lt: lt_sys::cublasLtHandle_t,
                      ltws: &mut CudaSlice<u8>,
                      plans: &std::collections::HashMap<(usize, usize, usize), Fp8Plan>,
                      key: (usize, usize, usize),
                      w: &CudaSlice<u16>, x: &CudaSlice<u16>, y: &mut CudaSlice<u16>,
                      m: usize, rows: usize, cols: usize) -> Result<()> {
        if let Some(p) = plans.get(&key) {
            return Self::bf16_matmul_planned_c16(stream, lt, p, ltws, 1 << 20, w, x, y);
        }
        Self::gemm_pre_c16(stream, blas, w, x, y, m, rows, cols)
    }

    #[allow(clippy::too_many_arguments)]
    fn bf16_matmul_planned(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w: &CudaSlice<u16>, x: &CudaSlice<u16>, d: &mut CudaSlice<f32>,
    ) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        unsafe {
            let (pw, _a) = w.device_ptr(stream);
            let (px, _b) = x.device_ptr(stream);
            let (pd, _c) = d.device_ptr_mut(stream);
            let (pws, _w) = ws.device_ptr_mut(stream);
            let st = lt_sys::cublasLtMatmul(
                lt, p.op, &alpha as *const f32 as *const _,
                pw as *const _, p.la, px as *const _, p.lb,
                &beta as *const f32 as *const _,
                pd as *const _, p.ld, pd as *mut _, p.ld,
                &p.algo, pws as *mut _, wss, stream.cu_stream() as _);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt bf16 matmul: {st:?}");
        }
        Ok(())
    }

    /// Rebuild the four per-layer FP8 PREFILL plan sets for row count `t` (prompt length for the
    /// staged path, wave rows for the wave). One activation-scale buffer per layer, shared by
    /// that layer's four GEMMs — each quant runs stream-ordered right before its matmul.
    fn fp8_prefill_plans(&mut self, t: usize) -> Result<()> {
        if self.fp8_pplan_t == t || self.fp8_gu.is_empty() {
            return Ok(());
        }
        let (hidden, inter, n) = (self.cfg.hidden, self.cfg.intermediate, self.cfg.layers);
        let mut widest = hidden.max(inter);
        for li in 0..n { widest = widest.max(self.layers[li].o.nblk * 32); }
        self.fp8_px = self.stream.alloc_zeros::<u8>(t * widest)?;
        self.fp8_psx = (0..n).map(|_| self.stream.alloc_zeros::<f32>(1))
            .collect::<std::result::Result<Vec<_>, _>>()?;
        self.fp8_psmax = (0..n).map(|_| self.stream.alloc_zeros::<f32>(1))
            .collect::<std::result::Result<Vec<_>, _>>()?;
        let (mut pg, mut pq, mut po, mut pd) =
            (Vec::with_capacity(n), Vec::with_capacity(n), Vec::with_capacity(n), Vec::with_capacity(n));
        for l2 in 0..n {
            let shapes: [(usize, usize, bool); 4] = [
                (2 * inter, hidden, false),
                (self.layers[l2].qkv.rows, hidden, true),
                (hidden, self.layers[l2].o.nblk * 32, true),
                (hidden, inter, true),
            ];
            for (si, &(rows, cols, d32)) in shapes.iter().enumerate() {
                let (sw, w8) = match si {
                    0 => (&self.fp8_gu[l2].1, &self.fp8_gu[l2].0),
                    1 => (&self.fp8_qkv[l2].1, &self.fp8_qkv[l2].0),
                    2 => (&self.fp8_o[l2].1, &self.fp8_o[l2].0),
                    _ => (&self.fp8_down[l2].1, &self.fp8_down[l2].0),
                };
                let key = (t, rows, cols, false);
                let cached = self.fp8_algo_cache.get(&key).copied();
                let (psp, _gp) = self.fp8_psx[l2].device_ptr(&self.stream);
                let mut plan = Self::fp8_plan(&self.stream, self.lt, t, rows, cols, 1 << 20,
                                              sw, psp, d32, cached)?;
                if cached.is_none() {
                    let a = Self::fp8_tune_plan(&self.stream, self.lt, &mut plan, &mut self.ltws,
                                                1 << 20, w8, &self.fp8_px, t * rows, d32)?;
                    self.fp8_algo_cache.insert(key, a);
                }
                match si { 0 => pg.push(plan), 1 => pq.push(plan), 2 => po.push(plan), _ => pd.push(plan) }
            }
        }
        self.fp8_pplans = pg;
        self.fp8_pplans_qkv = pq;
        self.fp8_pplans_o = po;
        self.fp8_pplans_down = pd;
        self.fp8_pplan_t = t;
        Ok(())
    }

    /// As `fp8_matmul_planned` but D is f32 — for the projections feeding the f32 residual
    /// stream (qkv / o / down), where a bf16 D would cost an extra conversion pass.
    #[allow(clippy::too_many_arguments)]
    fn fp8_matmul_planned_f32(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w8: &CudaSlice<u8>, x8: &CudaSlice<u8>, d: &mut CudaSlice<f32>, alpha: f32,
    ) -> Result<()> {
        let beta = 0f32;
        unsafe {
            let (pw, _a) = w8.device_ptr(stream);
            let (px, _b) = x8.device_ptr(stream);
            let (pd, _c) = d.device_ptr_mut(stream);
            let (pws, _w) = ws.device_ptr_mut(stream);
            let st = lt_sys::cublasLtMatmul(
                lt, p.op, &alpha as *const f32 as *const _,
                pw as *const _, p.la, px as *const _, p.lb,
                &beta as *const f32 as *const _,
                pd as *const _, p.ld, pd as *mut _, p.ld,
                &p.algo, pws as *mut _, wss, stream.cu_stream() as _);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt matmul f32: {st:?}");
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn fp8_matmul_planned(
        stream: &Arc<CudaStream>, lt: lt_sys::cublasLtHandle_t, p: &Fp8Plan,
        ws: &mut CudaSlice<u8>, wss: usize,
        w8: &CudaSlice<u8>, x8: &CudaSlice<u8>, d: &mut CudaSlice<u16>, alpha: f32,
    ) -> Result<()> {
        let beta = 0f32;
        unsafe {
            let (pw, _a) = w8.device_ptr(stream);
            let (px, _b) = x8.device_ptr(stream);
            let (pd, _c) = d.device_ptr_mut(stream);
            let (pws, _w) = ws.device_ptr_mut(stream);
            let st = lt_sys::cublasLtMatmul(
                lt, p.op, &alpha as *const f32 as *const _,
                pw as *const _, p.la, px as *const _, p.lb,
                &beta as *const f32 as *const _,
                pd as *const _, p.ld, pd as *mut _, p.ld,
                &p.algo, pws as *mut _, wss, stream.cu_stream() as _);
            anyhow::ensure!(st == lt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS, "lt matmul: {st:?}");
        }
        Ok(())
    }


    /// Launch the q4 GEMM without borrowing all of `self` — the hot path needs to pass one of
    /// self's weights alongside `self`'s function table, and CudaSlice::clone() would DEEP COPY on
    /// device, so cloning per step is not an option.
    #[allow(clippy::too_many_arguments)]
    fn gemm_q4s(
        stream: &Arc<CudaStream>, fs: &[CudaFunction; 4], fred: &CudaFunction,
        part: &mut CudaSlice<f32>, scales: &CudaSlice<f16>, quants_p: &CudaSlice<u32>,
        nblk: usize, xbf: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
        m: usize, rows: usize, cols: usize,
    ) -> Result<()> {
        let (vi, brows, bcols) = if m <= 16 { (0usize, 16usize, 64usize) }
                                 else if m <= 32 { (1, 32, 128) }
                                 else if m <= 64 { (2, 64, 128) }
                                 else { (3, 64, 256) };
        let tiles = m.div_ceil(brows) * rows.div_ceil(bcols);
        let nks = (400 / tiles.max(1)).max(1)
            .min((cols * 9 / (64 * m.max(1))).max(1))
            .clamp(1, 8)
            .min(nblk);
        let (mi, ri, ci, nb, nksi) =
            (m as i32, rows as i32, cols as i32, nblk as i32, nks as i32);
        let grid = (m.div_ceil(brows) as u32, rows.div_ceil(bcols) as u32, nks as u32);
        let thr = 256u32;
        unsafe {
            if nks == 1 {
                stream.launch_builder(&fs[vi])
                    .arg(scales).arg(quants_p).arg(xbf).arg(y)
                    .arg(&mi).arg(&ri).arg(&ci).arg(&nb).arg(&nksi)
                    .launch(LaunchConfig { grid_dim: grid, block_dim: (thr,1,1), shared_mem_bytes: 0 })?;
            } else {
                stream.launch_builder(&fs[vi])
                    .arg(scales).arg(quants_p).arg(xbf).arg(&mut *part)
                    .arg(&mi).arg(&ri).arg(&ci).arg(&nb).arg(&nksi)
                    .launch(LaunchConfig { grid_dim: grid, block_dim: (thr,1,1), shared_mem_bytes: 0 })?;
                let n = (m * rows) as i32;
                stream.launch_builder(fred)
                    .arg(&*part).arg(y).arg(&n).arg(&nksi).launch(ecfg(m * rows))?;
            }
        }
        Ok(())
    }

    /// Build (once) the permuted copy of a weight matrix that `gemm_q4_v2` consumes.
    fn ensure_permuted(&mut self, w: &mut Q4Dev) -> Result<()> {
        if w.quants_p.is_some() {
            return Ok(());
        }
        let nb = (w.rows * w.nblk) as i64;
        let mut out = self.stream.alloc_zeros::<u32>(w.rows * w.nblk * 4)?;
        unsafe {
            self.stream.launch_builder(&self.f_q4rep)
                .arg(&w.quants).arg(&mut out).arg(&nb)
                .launch(ecfg(w.rows * w.nblk))?;
        }
        w.quants_p = Some(out);
        Ok(())
    }

    fn gemm_q4_v2(&mut self, w: &Q4Dev, xbf: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
                  m: usize, rows: usize, cols: usize) -> Result<()> {
        // Split-K depth: SWEPT, not derived. Two costs pull opposite ways — the output tile alone
        // cannot fill the GPU at decode widths, but the partials pass reads nks*m*rows*4 bytes.
        //
        // The earlier formula sized nks to make reduction traffic equal weight traffic, which gave 32
        // at M=8. Measuring says that is 8x too aggressive once the tiles are small: nks=4 takes M=8
        // from 0.019 to 0.015 ms and M=16 from 0.018 to 0.016 — parity with cuBLAS. The tile count is
        // already large at narrow widths (96 column-blocks at BN=64), so the split only needs to top
        // the block count up, and every extra split is pure reduction traffic.
        // So: enough splits for ~400 blocks, never more than the reduction budget allows, cap 8.
        let (vi, brows, bcols) = if m <= 16 { (0usize, 16usize, 64usize) }
                                 else if m <= 32 { (1, 32, 128) }
                                 else if m <= 64 { (2, 64, 128) }
                                 else { (3, 64, 256) };
        let tiles = m.div_ceil(brows) * rows.div_ceil(bcols);
        let nks_par = (400 / tiles.max(1)).max(1);
        let nks_bw = (cols * 9 / (64 * m.max(1))).max(1);
        let nks = std::env::var("QNKS").ok().and_then(|v| v.parse().ok())
            .unwrap_or_else(|| nks_par.min(nks_bw))
            .clamp(1, 8)
            .min(w.nblk);
        let need = nks * m * rows;
        if nks == 1 {
            // no partials at all: the kernel's output layout IS C's
            let (mi, ri, ci, nb, nksi) =
                (m as i32, rows as i32, cols as i32, w.nblk as i32, 1i32);
            unsafe {
                self.stream.launch_builder(&self.f_q4v2[vi])
                    .arg(&w.scales).arg(w.quants_p.as_ref().unwrap()).arg(xbf).arg(y)
                    .arg(&mi).arg(&ri).arg(&ci).arg(&nb).arg(&nksi)
                    .launch(LaunchConfig {
                        grid_dim: (m.div_ceil(brows) as u32, rows.div_ceil(bcols) as u32, 1),
                        block_dim: (256, 1, 1), shared_mem_bytes: 0,
                    })?;
            }
            return Ok(());
        }
        if self.q4part.as_ref().map(|p| p.len()).unwrap_or(0) < need {
            self.q4part = Some(self.stream.alloc_zeros::<f32>(need)?);
        }
        let mut part = self.q4part.take().unwrap();
        let (mi, ri, ci, nb, nksi) =
            (m as i32, rows as i32, cols as i32, w.nblk as i32, nks as i32);
        let r = (|| -> Result<()> {
            unsafe {
                self.stream.launch_builder(&self.f_q4v2[vi])
                    .arg(&w.scales).arg(w.quants_p.as_ref().unwrap()).arg(xbf).arg(&mut part)
                    .arg(&mi).arg(&ri).arg(&ci).arg(&nb).arg(&nksi)
                    .launch(LaunchConfig {
                        grid_dim: (m.div_ceil(brows) as u32, rows.div_ceil(bcols) as u32, nks as u32),
                        block_dim: (256, 1, 1), shared_mem_bytes: 0,
                    })?;
                let n = (m * rows) as i32;
                self.stream.launch_builder(&self.f_q4red)
                    .arg(&part).arg(y).arg(&n).arg(&nksi).launch(ecfg(m * rows))?;
            }
            Ok(())
        })();
        self.q4part = Some(part);
        r
    }

    /// Diff the q4 tensor-core GEMM against cuBLAS on bf16 mirrors, on real weights, and time both.
    /// A q4 GEMM that is fast and subtly wrong would sail past every throughput number we report.
    /// Wide-M q4 GEMM: y[M, rows] = xb[M, cols] · Wᵀ, W q4_0 RAW layout (not permuted).
    pub fn gemm_q4_wg(&mut self, w: &Q4Dev, xb: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
                      m: usize, rows: usize, cols: usize, a_is_f16: bool, repacked: bool,
                      dbg: Option<&mut CudaSlice<u16>>) -> Result<()> {
        let (mi, ri, ci, nbi) = (m as i32, rows as i32, cols as i32, w.nblk as i32);
        let af = a_is_f16 as i32;
        let rp = repacked as i32;
        let null = 0u64;
        // BN=64 while the BN=128 grid would leave SMs idle (measured +21% at M=64, -37% at
        // M=2000 — so it is a WIDTH dispatch, not a replacement).
        let narrow = m <= 128;
        let (bn, smem) = if narrow { (64usize, 54272u32) } else { (128usize, 75776u32) };
        // SPLIT-K when the (rows/BN, M/128) grid leaves the machine idle. Each slice must keep
        // enough K stages to amortise its prologue, hence the >=4-stage floor.
        // MEASURED trigger: split-K pays only where the grid is genuinely starved AND the
        // extra reduce pass is small — M<=32 gained (3->4, 6->7 TF), M=128/256 LOST (52->39,
        // 82->45) because the [ksplit][M][rows] sum outweighed the parallelism.
        let nblocks = rows.div_ceil(bn) * m.div_ceil(128);
        let stages = cols / 64;
        let ksplit = if m > 32 || nblocks >= 128 { 1 }
                     else { (256 / nblocks.max(1)).clamp(1, (stages / 4).max(1)).min(8) };
        if ksplit > 1 {
            let need = ksplit * m * rows;
            if self.gwpart.as_ref().map(|p| p.len()).unwrap_or(0) < need {
                self.gwpart = Some(self.stream.alloc_zeros::<f32>(need)?);
            }
        }
        let ksi = ksplit as i32;
        let f = if narrow { self.f_q4wg64.clone() } else { self.f_q4wg.clone() };
        let mut dummy = self.dummy_f32.clone();
        unsafe {
            let mut lb = self.stream.launch_builder(&f);
            let partp: &mut CudaSlice<f32> = if ksplit > 1 { self.gwpart.as_mut().unwrap() }
                                             else { &mut dummy };
            lb.arg(&w.scales).arg(&w.quants).arg(xb).arg(&mut *y)
                .arg(&mi).arg(&ri).arg(&ci).arg(&nbi).arg(&af).arg(&rp).arg(&ksi).arg(partp);
            match dbg {
                Some(d) => { lb.arg(d); }
                None => { lb.arg(&null); }
            }
            lb.launch(LaunchConfig {
                grid_dim: (rows.div_ceil(bn) as u32, m.div_ceil(128) as u32, ksplit as u32),
                block_dim: (256, 1, 1),
                shared_mem_bytes: smem,
            })?;
        }
        if ksplit > 1 {
            let n = (m * rows) as i32;
            let part = self.gwpart.as_ref().unwrap().clone();
            unsafe {
                self.stream.launch_builder(&self.f_gwksum)
                    .arg(&part).arg(&mut *y).arg(&n).arg(&ksi).launch(ecfg(m * rows))?;
            }
        }
        Ok(())
    }

    /// Marlin-structure q4 GEMM (f16 activations required), associated form: takes the
    /// pieces it needs so a caller can pass DISJOINT fields of self. The &mut self version
    /// forced clone_ref() in the decode loop, and CudaSlice::clone ALLOCATES — inside a
    /// captured graph that both reallocates and orphans the capture (illegal address at M>1
    /// while M=1, which never captures a wide graph, looked fine).
    #[allow(clippy::too_many_arguments)]
    pub fn gemm_q4_ml_raw(stream: &Arc<CudaStream>, f_ml: &CudaFunction, f_ksum: &CudaFunction,
                          part: &mut CudaSlice<f32>, w: &Q4Dev, xb: &CudaSlice<u16>,
                          y: &mut CudaSlice<f32>, m: usize, rows: usize, cols: usize) -> Result<()> {
        let (mi, ri, ci, nbi) = (m as i32, rows as i32, cols as i32, w.nblk as i32);
        let mblk = m.div_ceil(64);
        let nb = rows.div_ceil(128) * mblk;
        let stages = cols / 64;
        let ksplit = if nb >= 228 { 1 } else { (228 / nb.max(1)).clamp(1, (stages / 4).max(1)).min(8) };
        let ksi = ksplit as i32;
        anyhow::ensure!(ksplit == 1 || part.len() >= ksplit * m * rows,
                        "marlin split-K partials too small: need {} have {}", ksplit * m * rows, part.len());
        unsafe {
            stream.launch_builder(f_ml)
                .arg(&w.scales).arg(&w.quants).arg(xb).arg(&mut *y)
                .arg(&mi).arg(&ri).arg(&ci).arg(&nbi).arg(&ksi).arg(&mut *part)
                .launch(LaunchConfig {
                    grid_dim: (rows.div_ceil(128) as u32, mblk as u32, ksplit as u32),
                    block_dim: (128, 1, 1),
                    shared_mem_bytes: 65536,
                })?;
        }
        if ksplit > 1 {
            let n = (m * rows) as i32;
            unsafe {
                stream.launch_builder(f_ksum)
                    .arg(&*part).arg(&mut *y).arg(&n).arg(&ksi).launch(ecfg(m * rows))?;
            }
        }
        Ok(())
    }

    /// Marlin-structure q4 GEMM (f16 activations required).
    pub fn gemm_q4_ml(&mut self, w: &Q4Dev, xb: &CudaSlice<u16>, y: &mut CudaSlice<f32>,
                      m: usize, rows: usize, cols: usize) -> Result<()> {
        let (mi, ri, ci, nbi) = (m as i32, rows as i32, cols as i32, w.nblk as i32);
        // Width dispatch: aim for >= ~2 blocks per SM. BN=32 gives rows/32 blocks (192 on
        // Qwen's gate|up) which is what a decode-width grid needs; the wide tile keeps
        // A-reuse for prefill.
        let mblk = m.div_ceil(64);
        // Wide tile ALWAYS (the 4-ldmatrix : 16-mma ratio is what makes it fast); fill the
        // machine with split-K instead. Target ~2 blocks per SM.
        let nb = rows.div_ceil(128) * mblk;
        let stages = cols / 64;
        let ksplit = if nb >= 228 { 1 } else { (228 / nb.max(1)).clamp(1, (stages / 4).max(1)).min(8) };
        // NEVER grow here: the decode step is a captured CUDA graph, and a realloc between
        // capture and replay leaves the graph writing through a freed pointer (the same
        // keystone failure the prefill graphs hit). init_batch sizes this once; if a caller
        // needs more, that is a bug to surface, not to paper over.
        if ksplit > 1 {
            let need = ksplit * m * rows;
            match self.mlpart.as_ref() {
                Some(p) if p.len() >= need => {}
                _ => {
                    self.mlpart = Some(self.stream.alloc_zeros::<f32>(need)?);
                }
            }
        }
        let ksi = ksplit as i32;
        let mut dummy = self.dummy_f32.clone();
        unsafe {
            let mut lb = self.stream.launch_builder(&self.f_q4ml);
            let pp: &mut CudaSlice<f32> = if ksplit > 1 { self.mlpart.as_mut().unwrap() }
                                          else { &mut dummy };
            lb.arg(&w.scales).arg(&w.quants).arg(xb).arg(&mut *y)
                .arg(&mi).arg(&ri).arg(&ci).arg(&nbi).arg(&ksi).arg(pp)
                .launch(LaunchConfig {
                    grid_dim: (rows.div_ceil(128) as u32, mblk as u32, ksplit as u32),
                    block_dim: (128, 1, 1),
                    shared_mem_bytes: 65536,
                })?;
        }
        if ksplit > 1 {
            let n = (m * rows) as i32;
            let part = self.mlpart.as_ref().unwrap().clone();
            unsafe {
                self.stream.launch_builder(&self.f_gwksum)
                    .arg(&part).arg(&mut *y).arg(&n).arg(&ksi).launch(ecfg(m * rows))?;
            }
        }
        Ok(())
    }

    pub fn verify_q4_gemm(&mut self, t: usize, iters: usize) -> Result<()> {
        anyhow::ensure!(self.layers[0].f16_gu.is_some(), "call init_batch first");
        let (hidden, inter) = (self.cfg.hidden, self.cfg.intermediate);
        if let Ok(one) = std::env::var("Q4ML_ONE") {
            // profiling mode: ONE warm + ONE timed Marlin launch at the requested width
            let want: usize = one.parse().unwrap_or(64);
            if t != want { return Ok(()); }
            let xh: Vec<f32> = (0..t * hidden)
                .map(|i| (((i * 2654435761usize) % 1000) as f32 / 500.0) - 1.0)
                .collect();
            let xf = self.stream.memcpy_stod(&xh)?;
            let mut xh16 = self.stream.alloc_zeros::<u16>(t * hidden)?;
            let ni = (t * hidden) as i32;
            unsafe {
                self.stream.launch_builder(&self.f_tof16)
                    .arg(&xf).arg(&mut xh16).arg(&ni).launch(ecfg(t * hidden))?;
            }
            let mut yw = self.stream.alloc_zeros::<f32>(t * inter)?;
            let gml = self.layers[0].gate.clone_ref();
            self.gemm_q4_ml(&gml, &xh16, &mut yw, t, inter, hidden)?;
            self.sync()?;
            self.gemm_q4_ml(&gml, &xh16, &mut yw, t, inter, hidden)?;
            self.sync()?;
            return Ok(());
        }
        if std::env::var("Q4WG_ONE").is_ok() {
            // profiling mode: ONE warm + ONE timed f16A launch at this width, nothing else
            if t != 2000 { return Ok(()); }
            let xh: Vec<f32> = (0..t * hidden)
                .map(|i| (((i * 2654435761usize) % 1000) as f32 / 500.0) - 1.0)
                .collect();
            let xf = self.stream.memcpy_stod(&xh)?;
            let mut xh16 = self.stream.alloc_zeros::<u16>(t * hidden)?;
            let ni = (t * hidden) as i32;
            unsafe {
                self.stream.launch_builder(&self.f_tof16)
                    .arg(&xf).arg(&mut xh16).arg(&ni).launch(ecfg(t * hidden))?;
            }
            let mut yw = self.stream.alloc_zeros::<f32>(t * inter)?;
            let gw0 = self.layers[0].gate.clone_ref();
            self.gemm_q4_wg(&gw0, &xh16, &mut yw, t, inter, hidden, true, false, None)?;
            self.sync()?;
            self.gemm_q4_wg(&gw0, &xh16, &mut yw, t, inter, hidden, true, false, None)?;
            self.sync()?;
            return Ok(());
        }
        let n = 2 * inter;
        // deterministic, non-degenerate activations: zeros would hide an indexing error entirely
        let xh: Vec<f32> = (0..t * hidden)
            .map(|i| (((i * 2654435761usize) % 1000) as f32 / 500.0) - 1.0)
            .collect();
        let xf = self.stream.memcpy_stod(&xh)?;
        let mut xb = self.stream.alloc_zeros::<u16>(t * hidden)?;
        let ni = (t * hidden) as i32;
        unsafe {
            self.stream.launch_builder(&self.f_tobf16)
                .arg(&xf).arg(&mut xb).arg(&ni).launch(ecfg(t * hidden))?;
        }
        let mut yq = self.stream.alloc_zeros::<f32>(t * n)?;
        let mut yr = self.stream.alloc_zeros::<f32>(t * n)?;
        // reference: cuBLAS on the bf16 mirror, f32 out
        Self::gemm_pre(&self.stream, &self.blas,
            self.layers[0].f16_gu.as_ref().unwrap(), &xb, &mut yr, t, n, hidden)?;
        // gate|up is gate then up concatenated by row, so q4 must be driven per matrix
        let mut gw0 = self.layers[0].gate.clone_ref();
        self.ensure_permuted(&mut gw0)?;
        self.gemm_q4_v2(&gw0, &xb, &mut yq, t, inter, hidden)?;
        self.sync()?;
        let (a, b) = (self.stream.memcpy_dtov(&yr)?, self.stream.memcpy_dtov(&yq)?);
        let (mut d, mut sc) = (0f32, 0f32);
        for r in 0..t {
            for c in 0..inter {
                let (x, y) = (a[r * n + c], b[r * inter + c]);
                d = d.max((x - y).abs());
                sc = sc.max(x.abs());
            }
        }
        let rel = d / sc.max(1e-6);
        println!("== Q4 GEMM vs cuBLAS/bf16: max|diff| = {d:.4}, max|val| = {sc:.4} (rel {rel:.5}) => {}",
                 if rel < 0.03 { "OK" } else { "WRONG" });
        // timing, same shape both ways
        let mut gw = self.layers[0].gate.clone_ref();
        self.ensure_permuted(&mut gw)?;
        for _ in 0..3 {
            self.gemm_q4_v2(&gw, &xb, &mut yq, t, inter, hidden)?;
        }
        self.sync()?;
        let t0 = std::time::Instant::now();
        for _ in 0..iters {
            self.gemm_q4_v2(&gw, &xb, &mut yq, t, inter, hidden)?;
        }
        self.sync()?;
        let dq = t0.elapsed().as_secs_f64() / iters as f64;
        let fl = 2.0 * t as f64 * inter as f64 * hidden as f64;
        println!("== Q4 GEMM   [{t}x{hidden}]x[{hidden}x{inter}]: {:.3} ms = {:.0} TFLOP/s, weights {:.1} MB",
                 dq * 1e3, fl / dq / 1e12, (inter * hidden) as f64 * 0.5625 / 1e6);
        // wgmma-q4 arm: correctness against the same cuBLAS reference, then timing
        {
            let mut yw = self.stream.alloc_zeros::<f32>(t * inter)?;
            // staging dump: A tile [64][64] then B tile [128][64], linear [mn][k]
            let mut dbg = self.stream.alloc_zeros::<u16>(64 * 64 + 128 * 64)?;
            let gw0 = self.layers[0].gate.clone_ref();
            self.gemm_q4_wg(&gw0, &xb, &mut yw, t, inter, hidden, false, false, Some(&mut dbg))?;
            self.sync()?;
            if std::env::var("Q4WG_DBG").is_ok() {
                let dv = self.stream.memcpy_dtov(&dbg)?;
                let bf = |u: u16| f32::from_bits((u as u32) << 16);
                // A: rows of xh; B: dequant of gate rows (reference = the f16_gu mirror bytes,
                // but recomputing from xh/scales here keeps it self-contained)
                let xbh = self.stream.memcpy_dtov(&xb)?;
                let mut abad = 0usize;
                for mn in 0..t.min(64) {
                    for k in 0..64 {
                        let got = bf(dv[mn * 64 + k]);
                        let want = bf(xbh[mn * hidden + k]);
                        if (got - want).abs() > 1e-3 * want.abs().max(1.0) {
                            if abad < 4 { println!("  A[{mn}][{k}] got {got} want {want}"); }
                            abad += 1;
                        }
                    }
                }
                let mirror = self.stream.memcpy_dtov(self.layers[0].f16_gu.as_ref().unwrap())?;
                let mut bbad = 0usize;
                for rn in 0..128 {
                    for k in 0..64 {
                        let got = bf(dv[64 * 64 + rn * 64 + k]);
                        let want = bf(mirror[rn * hidden + k]);
                        if (got - want).abs() > 1e-3 * want.abs().max(1.0) {
                            if bbad < 4 { println!("  B[{rn}][{k}] got {got} want {want}"); }
                            bbad += 1;
                        }
                    }
                }
                println!("== Q4WG staging: A mismatches {abad}/{} B mismatches {bbad}/{}",
                         t.min(64) * 64, 128 * 64);
            }
            let ww = self.stream.memcpy_dtov(&yw)?;
            let (mut d2, mut s2) = (0f32, 0f32);
            for r in 0..t {
                for c in 0..inter {
                    let (x, y2) = (a[r * n + c], ww[r * inter + c]);
                    d2 = d2.max((x - y2).abs());
                    s2 = s2.max(x.abs());
                }
            }
            let rel2 = d2 / s2.max(1e-6);
            println!("== Q4 WGMMA vs cuBLAS/bf16: max|diff| = {d2:.4} (rel {rel2:.5}) => {}",
                     if rel2 < 0.03 { "OK" } else { "WRONG" });
            for _ in 0..3 { self.gemm_q4_wg(&gw0, &xb, &mut yw, t, inter, hidden, false, false, None)?; }
            self.sync()?;
            let tw = std::time::Instant::now();
            for _ in 0..iters { self.gemm_q4_wg(&gw0, &xb, &mut yw, t, inter, hidden, false, false, None)?; }
            self.sync()?;
            let dw = tw.elapsed().as_secs_f64() / iters as f64;
            println!("== Q4 WGMMA  [{t}x{hidden}]x[{hidden}x{inter}]: {:.3} ms = {:.0} TFLOP/s",
                     dw * 1e3, fl / dw / 1e12);
            // f16-native activations (the production shape once the norm kernels emit f16)
            let mut xh16 = self.stream.alloc_zeros::<u16>(t * hidden)?;
            unsafe {
                self.stream.launch_builder(&self.f_tof16)
                    .arg(&xf).arg(&mut xh16).arg(&ni).launch(ecfg(t * hidden))?;
            }
            for _ in 0..3 { self.gemm_q4_wg(&gw0, &xh16, &mut yw, t, inter, hidden, true, false, None)?; }
            self.sync()?;
            let tw2 = std::time::Instant::now();
            for _ in 0..iters { self.gemm_q4_wg(&gw0, &xh16, &mut yw, t, inter, hidden, true, false, None)?; }
            self.sync()?;
            let dw2 = tw2.elapsed().as_secs_f64() / iters as f64;
            println!("== Q4 WGMMA/f16A                          {:.3} ms = {:.0} TFLOP/s",
                     dw2 * 1e3, fl / dw2 / 1e12);
            // repacked quants: shift-free f16x2 extraction (q4_repack_wg, offline)
            let g = &self.layers[0].gate;
            let nb_total = g.rows * g.nblk;
            let mut rq = self.stream.alloc_zeros::<u32>(nb_total * 4)?;
            let nbl = nb_total as i64;
            unsafe {
                self.stream.launch_builder(&self.f_q4rep_wg)
                    .arg(&g.quants).arg(&mut rq).arg(&nbl).launch(ecfg(nb_total))?;
            }
            let grp = Q4Dev { scales: g.scales.clone(), quants: rq, quants_p: None,
                              rows: g.rows, nblk: g.nblk };
            self.gemm_q4_wg(&grp, &xh16, &mut yw, t, inter, hidden, true, true, None)?;
            self.sync()?;
            let wr = self.stream.memcpy_dtov(&yw)?;
            let mut d3 = 0f32;
            for r in 0..t { for c in 0..inter {
                d3 = d3.max((a[r * n + c] - wr[r * inter + c]).abs());
            } }
            let rel3 = d3 / sc.max(1e-6);
            println!("== Q4 WGMMA/repack vs cuBLAS: max|diff| = {d3:.4} (rel {rel3:.5}) => {}",
                     if rel3 < 0.03 { "OK" } else { "WRONG" });
            for _ in 0..3 { self.gemm_q4_wg(&grp, &xh16, &mut yw, t, inter, hidden, true, true, None)?; }
            self.sync()?;
            let tw3 = std::time::Instant::now();
            for _ in 0..iters { self.gemm_q4_wg(&grp, &xh16, &mut yw, t, inter, hidden, true, true, None)?; }
            self.sync()?;
            let dw3 = tw3.elapsed().as_secs_f64() / iters as f64;
            println!("== Q4 WGMMA/repack+f16A                   {:.3} ms = {:.0} TFLOP/s",
                     dw3 * 1e3, fl / dw3 / 1e12);
            // Marlin-structure kernel
            let gml = self.layers[0].gate.clone_ref();
            self.gemm_q4_ml(&gml, &xh16, &mut yw, t, inter, hidden)?;
            self.sync()?;
            let wm = self.stream.memcpy_dtov(&yw)?;
            let mut d4 = 0f32;
            for r in 0..t { for c in 0..inter {
                d4 = d4.max((a[r * n + c] - wm[r * inter + c]).abs());
            } }
            let rel4 = d4 / sc.max(1e-6);
            println!("== Q4 MARLIN vs cuBLAS: max|diff| = {d4:.4} (rel {rel4:.5}) => {}",
                     if rel4 < 0.03 { "OK" } else { "WRONG" });
            for _ in 0..3 { self.gemm_q4_ml(&gml, &xh16, &mut yw, t, inter, hidden)?; }
            self.sync()?;
            let tm = std::time::Instant::now();
            for _ in 0..iters { self.gemm_q4_ml(&gml, &xh16, &mut yw, t, inter, hidden)?; }
            self.sync()?;
            let dm = tm.elapsed().as_secs_f64() / iters as f64;
            println!("== Q4 MARLIN [{t}x{hidden}]x[{hidden}x{inter}]: {:.3} ms = {:.0} TFLOP/s",
                     dm * 1e3, fl / dm / 1e12);
        }
        let t1 = std::time::Instant::now();
        for _ in 0..iters {
            Self::gemm_pre(&self.stream, &self.blas,
                self.layers[0].f16_gu.as_ref().unwrap(), &xb, &mut yr, t, inter, hidden)?;
        }
        self.sync()?;
        let db = t1.elapsed().as_secs_f64() / iters as f64;
        println!("== bf16 GEMM (same shape):                {:.3} ms = {:.0} TFLOP/s, weights {:.1} MB",
                 db * 1e3, fl / db / 1e12, (inter * hidden) as f64 * 2.0 / 1e6);
        Ok(())
    }

    /// Capture prefill as a CUDA graph.
    ///
    /// Prefill issues ~340 launches (four GEMMs, two norms, rope, V-transpose, attention and silu
    /// per layer, 28 layers). Measured in isolation cuBLAS runs the FFN GEMMs at 641 TFLOP/s, but
    /// in context they achieve 437 — the difference is the gap in front of each launch, and the
    /// GPU drawing 90 W of a 350 W budget at full clock says the same thing. A graph collapses the
    /// whole prefill into one submission.
    ///
    /// The graph is shape-specific: valid only for the token count it was captured at.
    pub fn capture_prefill(&mut self, tokens: &[u32]) -> Result<()> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        // Warm up first: cuBLAS picks algorithms and allocates its workspace on the first call for
        // a shape, and an allocation during capture aborts it.
        self.prefill(tokens)?;
        self.sync()?;
        let t = tokens.len();
        let mut p = self.pf.take().unwrap();
        self.stream
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
        let r = self.prefill_inner(&mut p, tokens, t, None, t);
        let g = self.stream.end_capture(
            CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
        );
        self.pf = Some(p);
        r.context("during prefill graph capture")?;
        let g = g?.context("stream capture produced no graph")?;
        g.upload()?;
        self.pfgraph = Some((t, g));
        self.sync()
    }

    /// Prefill into a batch slot through a captured CUDA graph, capturing on first sight of a
    /// given (length, slot) and replaying thereafter.
    ///
    /// `prefill_inner` already takes the slot as its last argument — `capture_prefill` simply passes
    /// `None` — so this is the same code path, just captured.
    pub fn prefill_slot_graphed(&mut self, tokens: &[u32], slot: usize) -> Result<()> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        let t = tokens.len();
        let key = t;                       // slot-independent: the offset lives in d_kvoff
        self.set_kvoff(Some(slot))?;
        if !self.pfslot_graphs.contains_key(&key) {
            // Warm up outside capture: cuBLAS allocates its workspace on the first call for a shape,
            // and an allocation during capture aborts it.
            self.prefill_slot(tokens, Some(slot))?;
            self.sync()?;
            let mut p = self.pf.take().unwrap();
            self.stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
            let r = self.prefill_inner(&mut p, tokens, t, Some(slot), t);
            let g = self.stream.end_capture(
                CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
            self.pf = Some(p);
            r.context("during slot prefill graph capture")?;
            let g = g?.context("stream capture produced no graph")?;
            g.upload()?;
            self.pfslot_graphs.insert(key, g);
            // The warm-up call already produced this request's KV and sampled its token — but the
            // burst-admission path reads first tokens from d_firsts, and skipping this snapshot
            // made every FIRST request of a (length, slot) pair emit token id 0 ("!") as its
            // opening token. Found by diffing served output against HF greedy.
            let (dt, df) = (&self.d_token, &mut self.d_firsts);
            self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
            self.sync()?;
            return Ok(());
        }
        let mut p = self.pf.take().unwrap();
        let r = self.stream.memcpy_htod(tokens, &mut p.tok).map_err(anyhow::Error::from);
        self.pf = Some(p);
        r?;
        self.pfslot_graphs[&key].launch()?;
        // Stream-ordered: snapshots this prefill's sampled token before the next replay
        // overwrites d_token. No sync — the caller drains a whole admission burst with one.
        let (dt, df) = (&self.d_token, &mut self.d_firsts);
        self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
        Ok(())
    }

    /// First tokens of every prefill replayed since the last sync, one host readback.
    pub fn first_tokens(&mut self) -> Result<Vec<u32>> {
        self.sync()?;
        Ok(self.stream.memcpy_dtov(&self.d_firsts)?)
    }

    /// DEBUG: what token ids the prefill input buffer actually holds (first n).
    pub fn peek_prefill_tokens(&mut self, n: usize) -> Result<Vec<u32>> {
        let p = self.pf.as_ref().context("no prefill state")?;
        let n = n.min(p.tok.len());
        Ok(self.stream.memcpy_dtov(&p.tok.slice(0..n))?)
    }

    /// Whether the sequential decode graph has been captured yet.
    pub fn has_seq_graph(&self) -> bool {
        self.graph.is_some()
    }

    /// Point the sequential decode head at (token, pos) — used when a sequence that has been
    /// stepping on the BATCH path hands back to the sequential graph.
    pub fn seed_seq_head(&mut self, token: u32, pos: usize) -> Result<()> {
        self.stream.memcpy_htod(&[token], &mut self.d_token)?;
        self.stream.memcpy_htod(&[pos as i32], &mut self.d_pos)?;
        Ok(())
    }

    /// Host-side reservation so a burst can claim several slots before any device work runs.
    pub fn reserve_slot(&mut self, slot: usize) {
        if let Some(b) = self.batch.as_mut() {
            if slot < b.mmax { b.occupied[slot] = true; }
        }
    }

    /// Upload a burst's concatenated prompt ids in ONE host->device copy.
    pub fn stage_burst(&mut self, concat: &[u32]) -> Result<()> {
        anyhow::ensure!(concat.len() <= self.d_stage.len(),
            "admission burst of {} ids exceeds the staging buffer", concat.len());
        self.stream.memcpy_htod(concat, &mut self.d_stage.slice_mut(0..concat.len()))?;
        Ok(())
    }

    /// prefill_slot_graphed, but the tokens come from `stage_burst`'s buffer by DEVICE copy — no
    /// synchronizing host transfer between replays, so a whole burst queues back-to-back. The
    /// host slice is still needed for the cold (capture) path.
    pub fn prefill_slot_staged(&mut self, tokens: &[u32], off: usize, slot: usize) -> Result<()> {
        let t = tokens.len();
        if std::env::var("GRAPH_PF").as_deref() == Ok("0") {
            // FULLY eager (first use included): the diagnostic fork must never capture, or
            // sync-instrumented prefill breaks the capture instead of running eagerly.
            self.prefill_slot(tokens, Some(slot))?;
            let (dt, df) = (&self.d_token, &mut self.d_firsts);
            self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
            return Ok(());
        }
        if !self.pfslot_graphs.contains_key(&t) {
            return self.prefill_slot_graphed(tokens, slot);
        }
        self.set_kvoff(Some(slot))?;
        // GRAPH_PF=0: run the warm path EAGERLY. Diagnostic fork for the replay corruption (the
        // captured graph's replay writes f32 garbage over p.tok when many decode steps ran since
        // capture; eager execution of the same work is the discriminator) — and, while that bug
        // lives, the correctness hotfix: eager prefill costs ~2 ms of launch overhead per
        // admission and nothing else.
        if std::env::var("GRAPH_PF").as_deref() == Ok("0") {
            self.prefill_slot(tokens, Some(slot))?;
            let (dt, df) = (&self.d_token, &mut self.d_firsts);
            self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
            return Ok(());
        }
        let mut p = self.pf.take().unwrap();
        let src = self.d_stage.slice(off..off + t);
        let r = self.stream.memcpy_dtod(&src, &mut p.tok.slice_mut(0..t))
            .map_err(anyhow::Error::from);
        self.pf = Some(p);
        r?;
        if std::env::var("EMIT_TRACE").is_ok() {
            self.sync()?;
            let st = self.stream.memcpy_dtov(&self.d_stage.slice(off..off + 4.min(t)))?;
            let pt = self.stream.memcpy_dtov(&self.pf.as_ref().unwrap().tok.slice(0..4.min(t)))?;
            eprintln!("[staged pre-launch] d_stage={st:?} p.tok={pt:?}");
        }
        self.pfslot_graphs[&t].launch()?;
        let (dt, df) = (&self.d_token, &mut self.d_firsts);
        self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
        Ok(())
    }

    /// Stream-ordered occupy: marks the slot active at `pos` without a synchronizing host copy.
    /// The input token is patched separately once the burst's first tokens are known.
    pub fn occupy_slot_async(&mut self, slot: usize, pos: usize) -> Result<()> {
        self.ctx_hint = self.ctx_hint.max(pos);
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(slot < b.mmax, "slot out of range");
        b.occupied[slot] = true;
        let (sl, tk, pv) = (slot as i32, 0u32, pos as i32);
        let zero = 0i32;
        unsafe {
            self.stream.launch_builder(&self.f_slot_ctl)
                .arg(&mut b.tok).arg(&mut b.pos).arg(&mut b.active)
                .arg(&sl).arg(&tk).arg(&pv).arg(&zero)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    /// Stream-ordered token patch for a slot occupied by `occupy_slot_async`.
    pub fn set_slot_token_async(&mut self, slot: usize, token: u32) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        let (sl, pv, one) = (slot as i32, 0i32, 1i32);
        unsafe {
            self.stream.launch_builder(&self.f_slot_ctl)
                .arg(&mut b.tok).arg(&mut b.pos).arg(&mut b.active)
                .arg(&sl).arg(&token).arg(&pv).arg(&one)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    /// Replay the captured prefill for `tokens`, which must be the length it was captured at.
    pub fn run_prefill_graph(&mut self, tokens: &[u32]) -> Result<()> {
        let (n, _) = self.pfgraph.as_ref().context("capture_prefill first")?;
        anyhow::ensure!(*n == tokens.len(),
            "prefill graph was captured for {n} tokens, got {}", tokens.len());
        let mut p = self.pf.take().unwrap();
        let r = self.stream.memcpy_htod(tokens, &mut p.tok).map_err(anyhow::Error::from);
        self.pf = Some(p);
        r?;
        self.pfgraph.as_ref().unwrap().1.launch()?;
        Ok(())
    }

    /// Time the FFN GEMMs back-to-back, with nothing interleaved. This separates two very
    /// different diagnoses of a 445 TFLOP/s in-context rate: cuBLAS being at its ceiling for the
    /// shape, versus launch gaps and dependency stalls between our kernels.
    pub fn bench_gemm(&mut self, t: usize, iters: usize) -> Result<()> {
        anyhow::ensure!(self.layers[0].f16_gu.is_some(), "call init_batch first");
        let (hidden, inter) = (self.cfg.hidden, self.cfg.intermediate);
        let mut xb = self.stream.alloc_zeros::<u16>(t * hidden)?;
        let mut y = self.stream.alloc_zeros::<u16>(t * 2 * inter)?;
        let mut y2 = self.stream.alloc_zeros::<f32>(t * hidden)?;
        let mut g = self.stream.alloc_zeros::<u16>(t * inter)?;
        for _ in 0..3 {
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[0].f16_gu.as_ref().unwrap(), &xb, &mut y, t, 2 * inter, hidden)?;
        }
        self.sync()?;
        let fl = 2.0 * t as f64 * (2 * inter) as f64 * hidden as f64;
        let t0 = std::time::Instant::now();
        for _ in 0..iters {
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[0].f16_gu.as_ref().unwrap(), &xb, &mut y, t, 2 * inter, hidden)?;
        }
        self.sync()?;
        let dt = t0.elapsed().as_secs_f64() / iters as f64;
        println!("== gate|up GEMM  [{t} x {hidden}] x [{hidden} x {}]: {:.3} ms = {:.0} TFLOP/s  (ONE layer, weights stay in L2)",
                 2 * inter, dt * 1e3, fl / dt / 1e12);
        // Same GEMM, but walking every layer's weights: 50 MB per layer against a 50 MB L2, which
        // is what prefill actually does. If the rate collapses, the ceiling is weight streaming,
        // not cuBLAS.
        let nl = self.cfg.layers;
        for li in 0..nl.min(4) {
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[li].f16_gu.as_ref().unwrap(), &xb, &mut y, t, 2 * inter, hidden)?;
        }
        self.sync()?;
        let tc = std::time::Instant::now();
        for i in 0..iters {
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[i % nl].f16_gu.as_ref().unwrap(), &xb, &mut y, t, 2 * inter, hidden)?;
        }
        self.sync()?;
        let dtc = tc.elapsed().as_secs_f64() / iters as f64;
        println!("== gate|up GEMM  cycling all {nl} layers:                {:.3} ms = {:.0} TFLOP/s",
                 dtc * 1e3, fl / dtc / 1e12);
        // FP8 CORRECTNESS on REAL weights. Every FP8 number so far was timed on ZEROED buffers, so
        // no output value had ever been checked — which is exactly why two unrelated scaling fixes
        // failed identically. Same activations into both paths, then diff.
        {
            let wbf = self.layers[0].f16_gu.as_ref().unwrap();
            let nw = (2 * inter * hidden) as i64;
            let nx = (t * hidden) as i64;
            // x = 1.0 everywhere (bf16 0x3F80): the result is a row-sum of real weights
            let xh: Vec<u16> = vec![0x3F80u16; t * hidden];
            let xr = self.stream.memcpy_stod(&xh)?;
            let mut ybf = self.stream.alloc_zeros::<u16>(t * 2 * inter)?;
            Self::gemm_pre_c16(&self.stream, &self.blas, wbf, &xr, &mut ybf, t, 2 * inter, hidden)?;

            let mut wamax = self.stream.alloc_zeros::<f32>(1)?;
            let mut xamax = self.stream.alloc_zeros::<f32>(1)?;
            let (mut ws8, mut xs8) = (self.stream.alloc_zeros::<f32>(1)?, self.stream.alloc_zeros::<f32>(1)?);
            let mut w8 = self.stream.alloc_zeros::<u8>(nw as usize)?;
            let mut x8 = self.stream.alloc_zeros::<u8>(nx as usize)?;
            let mut y8 = self.stream.alloc_zeros::<u16>(t * 2 * inter)?;
            let one = LaunchConfig { grid_dim: (1,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 };
            let tiny = LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: 0 };
            unsafe {
                self.stream.launch_builder(&self.f_absmax_bf).arg(wbf).arg(&mut wamax).arg(&nw)
                    .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
                self.stream.launch_builder(&self.f_mkscale).arg(&wamax).arg(&mut ws8).launch(tiny)?;
                self.stream.launch_builder(&self.f_tofp8).arg(wbf).arg(&mut w8).arg(&wamax).arg(&nw)
                    .launch(ecfg((nw as usize).div_ceil(2)))?;
                self.stream.launch_builder(&self.f_absmax_bf).arg(&xr).arg(&mut xamax).arg(&nx)
                    .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
                self.stream.launch_builder(&self.f_mkscale).arg(&xamax).arg(&mut xs8).launch(tiny)?;
                self.stream.launch_builder(&self.f_tofp8).arg(&xr).arg(&mut x8).arg(&xamax).arg(&nx)
                    .launch(ecfg((nx as usize).div_ceil(2)))?;
            }
            let (xsp, _gx) = xs8.device_ptr(&self.stream);
            let plan = Self::fp8_plan(&self.stream, self.lt, t, 2 * inter, hidden, 1 << 20, &ws8, xsp, false, None)?;
            let mut wsb = self.stream.alloc_zeros::<u8>(1 << 20)?;
            Self::fp8_matmul_planned(&self.stream, self.lt, &plan, &mut wsb, 1 << 20,
                                     &w8, &x8, &mut y8, 1.0)?;
            self.sync()?;
            let a = self.stream.memcpy_dtov(&ybf)?;
            let b2 = self.stream.memcpy_dtov(&y8)?;
            let bf = |u: u16| f32::from_bits((u as u32) << 16);
            let (mut worst, mut mag) = (0f32, 0f32);
            for i in 0..a.len() {
                let (x, y) = (bf(a[i]), bf(b2[i]));
                worst = worst.max((x - y).abs());
                mag = mag.max(x.abs());
            }
            println!("== FP8 vs BF16 on REAL weights: max|diff| = {worst:.4}, max|val| = {mag:.4} (rel {:.5}) => {}",
                     worst / mag.max(1e-9), if worst / mag.max(1e-9) < 0.05 { "OK" } else { "WRONG" });
        }

        // FP8 through our own cuBLASLt wrapper — proves the RUST path, not just the C probe.
        {
            let w8 = self.stream.alloc_zeros::<u8>(2 * inter * hidden)?;
            let x8 = self.stream.alloc_zeros::<u8>(t * hidden)?;
            let mut d8 = self.stream.alloc_zeros::<u16>(t * 2 * inter)?;
            let mut wsb = self.stream.alloc_zeros::<u8>(64 << 20)?;
            let (sxp2, _g2) = self.fp8_sx.device_ptr(&self.stream);
            let bplan = Self::fp8_plan(&self.stream, self.lt, t, 2 * inter, hidden, 64 << 20,
                                       &self.fp8_sx, sxp2, false, None)?;
            let r = (0..3).try_for_each(|_| Self::fp8_matmul_planned(&self.stream, self.lt, &bplan,
                        &mut wsb, 64 << 20, &w8, &x8, &mut d8, 1.0));
            match r {
                Err(e) => println!("== FP8 (Lt, rust): {e}"),
                Ok(()) => {
                    self.sync()?;
                    let tf = std::time::Instant::now();
                    for _ in 0..iters {
                        Self::fp8_matmul_planned(&self.stream, self.lt, &bplan, &mut wsb,
                                                 64 << 20, &w8, &x8, &mut d8, 1.0)?;
                    }
                    self.sync()?;
                    let dtf = tf.elapsed().as_secs_f64() / iters as f64;
                    println!("== FP8  gate|up GEMM (cuBLASLt, rust) [{t} x {hidden}] x [{hidden} x {}]: {:.3} ms = {:.0} TFLOP/s  ({:.2}x bf16)",
                             2 * inter, dtf * 1e3, fl / dtf / 1e12, dt / dtf);
                }
            }
        }

        // FP8 on the SAME shape. H100 FP8 tensor cores are 2x BF16 and vLLM's prefill is BF16, so
        // this is the one route that ends AHEAD rather than at parity. Measure the ceiling before
        // building any FP8 mirrors: if the 2x is not there on these shapes, the plan is dead.
        {
            let w8 = self.stream.alloc_zeros::<u8>(2 * inter * hidden)?;
            let x8 = self.stream.alloc_zeros::<u8>(t * hidden)?;
            let mut y8 = self.stream.alloc_zeros::<u16>(t * 2 * inter)?;
            let (alpha, beta) = (1.0f32, 0.0f32);
            let mut ok = true;
            for warm in 0..(iters + 3) {
                let (pw, _a) = w8.device_ptr(&self.stream);
                let (px, _b) = x8.device_ptr(&self.stream);
                let (py, _c) = y8.device_ptr_mut(&self.stream);
                let r = unsafe {
                    blas::gemm_ex(
                        *self.blas.handle(),
                        blas_sys::cublasOperation_t::CUBLAS_OP_T,
                        blas_sys::cublasOperation_t::CUBLAS_OP_N,
                        (2 * inter) as i32, t as i32, hidden as i32,
                        (&alpha) as *const f32 as *const _,
                        pw as *const std::ffi::c_void,
                        blas_sys::cudaDataType_t::CUDA_R_8F_E4M3, hidden as i32,
                        px as *const std::ffi::c_void,
                        blas_sys::cudaDataType_t::CUDA_R_8F_E4M3, hidden as i32,
                        (&beta) as *const f32 as *const _,
                        py as *mut std::ffi::c_void,
                        blas_sys::cudaDataType_t::CUDA_R_16BF, (2 * inter) as i32,
                        blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                        blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
                    )
                };
                if let Err(e) = r {
                    println!("== FP8 gate|up GEMM: UNSUPPORTED via cublasGemmEx ({e:?}) — needs cublasLt");
                    ok = false;
                    break;
                }
                if warm == 2 { self.sync()?; }
            }
            if ok {
                self.sync()?;
                let tf = std::time::Instant::now();
                for _ in 0..iters {
                    let (pw, _a) = w8.device_ptr(&self.stream);
                    let (px, _b) = x8.device_ptr(&self.stream);
                    let (py, _c) = y8.device_ptr_mut(&self.stream);
                    unsafe {
                        let _ = blas::gemm_ex(
                            *self.blas.handle(),
                            blas_sys::cublasOperation_t::CUBLAS_OP_T,
                            blas_sys::cublasOperation_t::CUBLAS_OP_N,
                            (2 * inter) as i32, t as i32, hidden as i32,
                            (&alpha) as *const f32 as *const _,
                            pw as *const std::ffi::c_void,
                            blas_sys::cudaDataType_t::CUDA_R_8F_E4M3, hidden as i32,
                            px as *const std::ffi::c_void,
                            blas_sys::cudaDataType_t::CUDA_R_8F_E4M3, hidden as i32,
                            (&beta) as *const f32 as *const _,
                            py as *mut std::ffi::c_void,
                            blas_sys::cudaDataType_t::CUDA_R_16BF, (2 * inter) as i32,
                            blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                            blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
                        );
                    }
                }
                self.sync()?;
                let dtf = tf.elapsed().as_secs_f64() / iters as f64;
                println!("== FP8  gate|up GEMM [{t} x {hidden}] x [{hidden} x {}]: {:.3} ms = {:.0} TFLOP/s  ({:.2}x the bf16 rate)",
                         2 * inter, dtf * 1e3, fl / dtf / 1e12, dt / dtf);
            }
        }
        let t1 = std::time::Instant::now();
        for _ in 0..iters {
            Self::gemm_pre(&self.stream, &self.blas,
                self.layers[0].f16_down.as_ref().unwrap(), &g, &mut y2, t, hidden, inter)?;
        }
        self.sync()?;
        let dt2 = t1.elapsed().as_secs_f64() / iters as f64;
        let fl2 = 2.0 * t as f64 * hidden as f64 * inter as f64;
        println!("== down GEMM     [{t} x {inter}] x [{inter} x {hidden}]: {:.3} ms = {:.0} TFLOP/s",
                 dt2 * 1e3, fl2 / dt2 / 1e12);
        // and the two chained, which is what prefill actually does
        let t2 = std::time::Instant::now();
        for _ in 0..iters {
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[0].f16_gu.as_ref().unwrap(), &xb, &mut y, t, 2 * inter, hidden)?;
            Self::gemm_pre(&self.stream, &self.blas,
                self.layers[0].f16_down.as_ref().unwrap(), &g, &mut y2, t, hidden, inter)?;
        }
        self.sync()?;
        let dt3 = t2.elapsed().as_secs_f64() / iters as f64;
        println!("== chained       {:.3} ms = {:.0} TFLOP/s  (sum of the two alone: {:.3} ms)",
                 dt3 * 1e3, (fl + fl2) / dt3 / 1e12, (dt + dt2) * 1e3);
        drop((xb, y, y2, g));
        Ok(())
    }

    /// Batched prefill of `tokens` (dense arm). Projections are one GEMM each on the resident f16
    /// weights; attention is two cuBLAS strided-batched GEMMs per layer so the T^2 work is on
    /// tensor cores. Leaves the KV cache filled and `d_pos` at T, ready for decode.
    pub fn prefill(&mut self, tokens: &[u32]) -> Result<()> {
        self.prefill_slot(tokens, None)
    }

    /// Size the lazily-allocated single-stream KV slabs (bench/verify paths only — the server's
    /// seq path rides batch slot 0 and never touches these).
    fn grow_seq_kv(&mut self) -> Result<()> {
        if self.kc[0].len() > 1 {
            return Ok(());
        }
        let st = self.stream.clone();
        for li in 0..self.cfg.layers {
            let (lhd, lnkv, _, _) = self.cfg.lgeom(li);
            self.kc[li] = st.alloc_zeros::<u16>(self.max_seq * lnkv * lhd)?;
            self.vc[li] = st.alloc_zeros::<u16>(self.max_seq * lnkv * lhd)?;
        }
        Ok(())
    }

    /// (Re)allocate the batched-MoE scratch for up to `rows` rows (batch width or prefill T).
    /// Reallocation orphans EVERY captured graph that touches these buffers — decode widths and
    /// prefill lengths both — so they are all dropped for lazy recapture.
    fn moe_grow(&mut self, rows: usize) -> Result<()> {
        if self.cfg.num_experts == 0 || rows <= self.moe_rows {
            return Ok(());
        }
        let (e_n, k, mi, hidden) = (self.cfg.num_experts, self.cfg.topk, self.cfg.moe_inter, self.cfg.hidden);
        let st = self.stream.clone();
        self.mrlog = st.alloc_zeros::<f32>(rows * e_n)?;
        self.meidx = st.alloc_zeros::<u32>(rows * k)?;
        self.mewt = st.alloc_zeros::<f32>(rows * k)?;
        self.meg = st.alloc_zeros::<f32>(rows * k * 2 * mi)?;
        self.meu = st.alloc_zeros::<f32>(rows * k * mi)?;
        self.medn = st.alloc_zeros::<f32>(rows * k * hidden)?;
        self.mexs = st.alloc_zeros::<f32>(2 * rows * k)?;
        self.mcnt = st.alloc_zeros::<i32>(e_n)?;
        self.moff = st.alloc_zeros::<i32>(e_n)?;
        self.mcur = st.alloc_zeros::<i32>(e_n)?;
        self.mplist = st.alloc_zeros::<u32>(rows * k)?;
        // worst case: every expert contributes a partial tile, plus all pairs tiled
        self.mwork = st.alloc_zeros::<u32>(e_n + rows * k / 16 + 8)?;
        self.mwn = st.alloc_zeros::<i32>(1)?;
        self.moe_rows = rows;
        self.bgraphs.clear();
        self.pfslot_graphs.clear();
        self.pfgraph = None;
        Ok(())
    }

    /// (Re)allocate the prefill scratch for prompts up to `t` rows. Shared by the serial and
    /// wave prefills — growth ORPHANS every captured prefill graph (raw device pointers), so the
    /// graphs are dropped here and recapture lazily at the new geometry.
    fn grow_prefill(&mut self, t: usize) -> Result<()> {
        self.moe_grow(t)?;
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let per = nh + 2 * nkv;
        {
            // Growing the Prefill buffers ORPHANS every captured prefill graph: a CUDA graph
            // holds raw device pointers, so replaying one after this reallocation reads the OLD
            // p.tok (serving the CAPTURED prompt's continuation for a NEW prompt) and, once the
            // freed blocks are reused, writes f32 activations over whatever lives there now.
            // Found end-to-end by diffing served ids against transformers greedy: a warm replay
            // after a longer prompt's capture returned the first prompt's output verbatim.
            // Dropping the graphs here makes them recapture lazily at the new geometry — one
            // ~30 ms recapture per length-class growth, and fixed-length workloads never pay it.
            self.pfslot_graphs.clear();
            self.pfslot_graphs2.clear();
            self.pfgraph = None;
            let st = &self.stream;
            // qkv16 and vt16 are padded to a whole number of flash-attention tiles and zero-filled,
            // so every cp.async is unconditional. Stale rows from a longer previous prompt are
            // harmless: the causal mask forces P = 0 there.
            let tp = t.next_multiple_of(64);
            self.pf = Some(Prefill {
                tmax: t,
                x: st.alloc_zeros::<f32>(t * hidden)?,
                h: st.alloc_zeros::<f32>(t * hidden)?,
                qkv: st.alloc_zeros::<f32>(t * per * hd)?,
                qkv16: st.alloc_zeros::<f16>(tp * per * hd)?,
                vt16: st.alloc_zeros::<f16>(nkv * hd * tp)?,
                score16: st.alloc_zeros::<f16>(nh * t * t)?,
                rowsum: st.alloc_zeros::<f32>(nh * t)?,
                // Doubles as [t, hidden] scratch for the final norm's output (the last-row
                // slice feeds the LM head). nh*hd == hidden on Qwen so the reuse was invisible;
                // Mistral has nh*hd (4096) < hidden (5120) and the slice panicked out of range
                // in cudarc on the FIRST batched prefill this model ever ran.
                attn: st.alloc_zeros::<f32>(t * (nh * hd).max(hidden))?,
                g: st.alloc_zeros::<f32>(t * inter)?,
                u: st.alloc_zeros::<f32>(t * inter)?,
                hb16: st.alloc_zeros::<u16>(t * hidden)?,
                attn16: st.alloc_zeros::<u16>(t * nh * hd)?,
                abuf16: st.alloc_zeros::<u16>(1)?,   // batched decode only
                gu16: st.alloc_zeros::<u16>(t * 2 * inter)?,
                gu32: st.alloc_zeros::<f32>(1)?,   // q4 arm is decode-only; prefill uses mirrors
                g16: st.alloc_zeros::<u16>(t * inter)?,
                logits: st.alloc_zeros::<f32>(c.vocab)?,
                tok: st.alloc_zeros::<u32>(t)?,
            });
        }
        Ok(())
    }

    /// WAVE PREFILL: B same-length prompts as ONE [B*T]-row pass. Every projection/FFN GEMM
    /// amortises across the wave; attention runs per-(segment, head) in attn_wave; KV lands in
    /// each request's slot via the d_kvoffs array; d_firsts[slot] gets each segment's first
    /// sampled token — so the server's existing first_tokens()/occupy flow works unchanged.
    /// Eager (no graph) — ~224 launches ≈ 1.8 ms host, against the ~3.5 ms * B it replaces.
    /// Non-sandwich models take the fused attn_wave path; sandwich (Gemma) routes through
    /// prefill_inner in wave mode — the score path with a segment-masked softmax, which gives
    /// per-layer geometry, dual rope, vderive and the fp8 arms for free.
    /// Caller guarantees: dense arm, tseg <= 64, all prompts exactly tseg long.
    pub fn prefill_wave(&mut self, flat: &[u32], tseg: usize, slots: &[usize]) -> Result<()> {
        let b = slots.len();
        let s_rows = b * tseg;
        anyhow::ensure!(self.cfg.num_experts == 0, "wave: dense arm only");
        anyhow::ensure!(flat.len() == s_rows && b >= 1, "wave shape");
        // tseg > 64 needs the per-segment FA arm (attn_wave's smem holds seg_t^2 scores)
        anyhow::ensure!(tseg <= 64 || (self.fa && self.cfg.head_dim == 128
            && !self.cfg.attn_scale_one) || self.cfg.sandwich, "wave: long segments need FA or score path");
        anyhow::ensure!(self.layers[0].f16_qkv.is_some() || !self.fp8_qkv.is_empty(),
            "init_batch first");
        let bmax = self.batch.as_ref().map(|bt| bt.max_seq).context("init_batch first")?;
        // per-segment KV slab offsets
        let offs: Vec<i32> = slots.iter().map(|&sl| (sl * bmax) as i32).collect();
        self.stream.memcpy_htod(&offs, &mut self.d_kvoffs.slice_mut(0..b))?;
        // grow the prefill scratch to S rows (reuses the same growth rules as prefill_slot)
        if self.pf.as_ref().map(|p| p.tmax) < Some(s_rows) {
            self.grow_prefill(s_rows)?;
        }
        let mut p = self.pf.take().unwrap();
        let sandwich = self.cfg.sandwich;
        let r = self.stream.memcpy_htod(flat, &mut p.tok.slice_mut(0..s_rows))
            .map_err(anyhow::Error::from)
            .and_then(|()| if sandwich {
                self.prefill_wave_score(&mut p, flat, tseg, b, slots)
            } else {
                self.prefill_wave_inner(&mut p, tseg, b, slots)
            });
        self.pf = Some(p);
        r
    }

    /// Sandwich score-path wave: one eager prefill_inner pass over all B*tseg rows (segment
    /// masking inside softmax_causal; KV offsets via d_kvoffs in rope_prefill), then the tail —
    /// gather each segment's LAST x/h rows, run the final sandwich fold on those B rows, one
    /// [B, vocab] LM GEMM, argmax_b, scatter to d_firsts.
    fn prefill_wave_score(&mut self, p: &mut Prefill, flat: &[u32], tseg: usize, b: usize,
                          slots: &[usize]) -> Result<()> {
        let s_rows = b * tseg;
        self.prefill_inner(p, flat, s_rows, None, tseg)?;
        let c = self.cfg.clone();
        let (hidden, eps) = (c.hidden, c.eps);
        let hi = hidden as i32;
        if self.wave_x.len() < b * hidden {
            self.wave_x = self.stream.alloc_zeros::<f32>(64 * hidden)?;
            self.wave_logits = self.stream.alloc_zeros::<f32>(64 * c.vocab)?;
            self.wave_first = self.stream.alloc_zeros::<u32>(64)?;
        }
        if self.wave_h.len() < b * hidden {
            self.wave_h = self.stream.alloc_zeros::<f32>(64 * hidden)?;
            self.wave_b16 = self.stream.alloc_zeros::<u16>(64 * hidden)?;
        }
        for sgi in 0..b {
            let last = (sgi + 1) * tseg - 1;
            let sx = p.x.slice(last * hidden..(last + 1) * hidden);
            let mut dx = self.wave_x.slice_mut(sgi * hidden..(sgi + 1) * hidden);
            self.stream.memcpy_dtod(&sx, &mut dx)?;
            let sh = p.h.slice(last * hidden..(last + 1) * hidden);
            let mut dh = self.wave_h.slice_mut(sgi * hidden..(sgi + 1) * hidden);
            self.stream.memcpy_dtod(&sh, &mut dh)?;
        }
        let lastl = c.layers - 1;
        let pls = self.layers[lastl].ls;
        unsafe {
            self.stream.launch_builder(&self.f_sandwich)
                .arg(&mut self.wave_x).arg(&self.wave_h)
                .arg(self.layers[lastl].post_ffw_norm.as_ref().unwrap())
                .arg(&self.final_norm)
                .arg(&mut self.wave_b16).arg(&hi).arg(&eps).arg(&pls)
                .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
        }
        {
            let (wb, wl) = (&self.wave_b16, &mut self.wave_logits);
            Self::gemm_pre(&self.stream, &self.blas, self.f16_lm.as_ref().unwrap(),
                wb, wl, b, c.vocab, hidden)?;
        }
        let v = c.vocab as i32;
        unsafe {
            self.stream.launch_builder(&self.f_argmax_b)
                .arg(&self.wave_logits).arg(&mut self.wave_first).arg(&v)
                .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
        }
        for (sgi, &sl) in slots.iter().enumerate() {
            let src = self.wave_first.slice(sgi..sgi + 1);
            let mut dst = self.d_firsts.slice_mut(sl..sl + 1);
            self.stream.memcpy_dtod(&src, &mut dst)?;
        }
        Ok(())
    }

    fn prefill_wave_inner(&mut self, p: &mut Prefill, tseg: usize, b: usize,
                          slots: &[usize]) -> Result<()> {
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let (per, eps) = (nh + 2 * nkv, c.eps);
        let s_rows = b * tseg;
        // WAVE_PROF=1: per-stage sync-timing, same shape as PREFILL_PROF (costs syncs; opt-in)
        let prof = std::env::var("WAVE_PROF").is_ok();
        let mut tp = [0f64; 8];
        let mut tk = std::time::Instant::now();
        macro_rules! mark { ($i:expr) => { if prof {
            self.stream.synchronize()?;
            let n = std::time::Instant::now();
            tp[$i] += n.duration_since(tk).as_secs_f64();
            tk = n;
        } } }
        let (si, hi, nhi, nkvi, hdi, tsi) =
            (s_rows as i32, hidden as i32, nh as i32, nkv as i32, hd as i32, tseg as i32);
        let esc = c.embed_scale;
        let rowsg = LaunchConfig { grid_dim: (s_rows as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
        self.fp8_prefill_plans(s_rows)?;
        unsafe {
            self.stream.launch_builder(&self.f_embed_p)
                .arg(&self.embed).arg(&mut p.x).arg(&mut p.h).arg(&p.tok).arg(&hi).arg(&si).arg(&esc)
                .launch(ecfg(s_rows * hidden))?;
        }
        let pos0 = 0i32;
        let awsh = ((2 * tseg * hd + tseg * tseg) * 4) as u32;
        for li in 0..c.layers {
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].attn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            mark!(0);
            if !self.fp8_qkv.is_empty() {
                let nx = (s_rows * hidden) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.hb16).arg(&mut self.fp8_part).arg(&nx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.hb16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                        .launch(ecfg((s_rows * hidden).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_qkv[li], ltws2,
                    1 << 20, &self.fp8_qkv[li].0, &self.fp8_px, &mut p.qkv, 1.0)?;
            } else {
                Self::gemm_pre(&self.stream, &self.blas,
                    self.layers[li].f16_qkv.as_ref().unwrap(), &p.hb16, &mut p.qkv,
                    s_rows, per * hd, hidden)?;
            }
            mark!(1);
            unsafe {
                let bt = self.batch.as_mut().unwrap();
                let bmaxi = bt.max_seq as i32;
                self.stream.launch_builder(if self.fp8kv { &self.f_rope_p_f8 } else { &self.f_rope_p })
                    .arg(&mut p.qkv).arg(&mut p.qkv16).arg(&mut bt.kc[li]).arg(&mut bt.vc[li])
                    .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&si).arg(&self.inv_freq).arg(&self.qkni).arg(&eps)
                    .arg(&pos0).arg(&self.d_kvoffs).arg(&tsi).arg(&bmaxi)
                    .launch(LaunchConfig { grid_dim: ((s_rows * per).div_ceil(4) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                if tseg <= 64 {
                    self.stream.launch_builder(&self.f_attn_wave)
                        .arg(&p.qkv16).arg(&mut p.attn16)
                        .arg(&tsi).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&(1.0f32 / (hd as f32).sqrt()))
                        .launch(LaunchConfig { grid_dim: (b as u32, nh as u32, 1),
                                               block_dim: (128, 1, 1), shared_mem_bytes: awsh })?;
                } else {
                    // LONG-SEGMENT WAVE: the GEMMs stay batched over all B*T rows (that is the
                    // whole point), and attention runs the FLASH kernel once per segment on that
                    // segment's row slice — segments are independent, so this is exactly the
                    // per-request attention with none of the per-request GEMM inefficiency.
                    let (tsegi, tpi) = (tseg as i32, tseg.next_multiple_of(64) as i32);
                    let scale = 1.0f32 / (hd as f32).sqrt();
                    for sg in 0..b {
                        let qoff = sg * tseg * per * hd;
                        let aoff = sg * tseg * nh * hd;
                        let qsl = p.qkv16.slice(qoff..qoff + tseg * per * hd);
                        self.stream.launch_builder(&self.f_vt)
                            .arg(&qsl).arg(&mut p.vt16).arg(&tsegi).arg(&tpi).arg(&nhi).arg(&nkvi).arg(&hdi)
                            .launch(LaunchConfig { grid_dim: (tseg.div_ceil(32) as u32, (hd / 32) as u32,
                                                              nkv as u32),
                                                   block_dim: (32, 8, 1), shared_mem_bytes: 0 })?;
                        let mut asl = p.attn16.slice_mut(aoff..aoff + tseg * nh * hd);
                        self.stream.launch_builder(&self.f_fa)
                            .arg(&qsl).arg(&p.vt16).arg(&mut asl).arg(&tsegi).arg(&tpi)
                            .arg(&nhi).arg(&nkvi).arg(&scale)
                            .launch(LaunchConfig { grid_dim: (tseg.div_ceil(64) as u32, nh as u32, 1),
                                                   block_dim: (128, 1, 1), shared_mem_bytes: FA_SHARED })?;
                    }
                }
            }
            mark!(2);
            if !self.fp8_o.is_empty() {
                let nx = (s_rows * nh * hd) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.attn16).arg(&mut self.fp8_part).arg(&nx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.attn16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                        .launch(ecfg((s_rows * nh * hd).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_o[li], ltws2,
                    1 << 20, &self.fp8_o[li].0, &self.fp8_px, &mut p.h, 1.0)?;
            } else {
                Self::gemm_pre(&self.stream, &self.blas,
                    self.layers[li].f16_o.as_ref().unwrap(), &p.attn16, &mut p.h,
                    s_rows, hidden, nh * hd)?;
            }
            mark!(3);
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].ffn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            if !self.fp8_gu.is_empty() {
                let nx = (s_rows * hidden) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.hb16).arg(&mut self.fp8_part).arg(&nx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.hb16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                        .launch(ecfg((s_rows * hidden).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned(&self.stream, lt2, &self.fp8_pplans[li], ltws2,
                    1 << 20, &self.fp8_gu[li].0, &self.fp8_px, &mut p.gu16, 1.0)?;
            } else {
                Self::gemm_pre_c16(&self.stream, &self.blas,
                    self.layers[li].f16_gu.as_ref().unwrap(), &p.hb16, &mut p.gu16,
                    s_rows, 2 * inter, hidden)?;
            }
            mark!(4);
            let (ii, n8, act) = (inter as i32, (s_rows * inter / 8) as i32, c.gelu as i32);
            unsafe {
                self.stream.launch_builder(&self.f_silu_bf)
                    .arg(&p.gu16).arg(&mut p.g16).arg(&ii).arg(&n8).arg(&act).launch(ecfg(s_rows * inter / 8))?;
            }
            mark!(5);
            if !self.fp8_down.is_empty() {
                let ndx = (s_rows * inter) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.g16).arg(&mut self.fp8_part).arg(&ndx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.g16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&ndx)
                        .launch(ecfg((s_rows * inter).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_down[li], ltws2,
                    1 << 20, &self.fp8_down[li].0, &self.fp8_px, &mut p.h, 1.0)?;
            } else {
                Self::gemm_pre(&self.stream, &self.blas,
                    self.layers[li].f16_down.as_ref().unwrap(), &p.g16, &mut p.h, s_rows, hidden, inter)?;
            }
            mark!(6);
        }
        // tail: final norm over all rows, gather each segment's LAST row, one [B, vocab] LM GEMM
        unsafe {
            self.stream.launch_builder(&self.f_rms_add_b)
                .arg(&mut p.x).arg(&p.h).arg(&self.final_norm)
                .arg(&mut p.attn).arg(&mut self.xs).arg(&hi).arg(&eps).launch(rowsg)?;
        }
        if self.wave_x.len() < b * hidden {
            self.wave_x = self.stream.alloc_zeros::<f32>(64 * hidden)?;
            self.wave_logits = self.stream.alloc_zeros::<f32>(64 * c.vocab)?;
            self.wave_first = self.stream.alloc_zeros::<u32>(64)?;
        }
        // gemm_bf16 stages b*hidden bf16 activations through xf16, which boots at one row's worth
        if self.xf16.len() < b * hidden {
            self.xf16 = self.stream.alloc_zeros::<u16>(64 * hidden)?;
            self.xf16b = self.stream.alloc_zeros::<u16>(64 * hidden)?;
        }
        for sgi in 0..b {
            let last = (sgi + 1) * tseg - 1;
            let src = p.attn.slice(last * hidden..(last + 1) * hidden);
            let mut dst = self.wave_x.slice_mut(sgi * hidden..(sgi + 1) * hidden);
            self.stream.memcpy_dtod(&src, &mut dst)?;
        }
        {
            let (wx, wl) = (&self.wave_x, &mut self.wave_logits);
            Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
                self.f16_lm.as_ref().unwrap(), &mut self.xf16, wx, wl, b, c.vocab, hidden)?;
        }
        let v = c.vocab as i32;
        unsafe {
            self.stream.launch_builder(&self.f_argmax_b)
                .arg(&self.wave_logits).arg(&mut self.wave_first).arg(&v)
                .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
        }
        for (sgi, &sl) in slots.iter().enumerate() {
            let src = self.wave_first.slice(sgi..sgi + 1);
            let mut dst = self.d_firsts.slice_mut(sl..sl + 1);
            self.stream.memcpy_dtod(&src, &mut dst)?;
        }
        mark!(7);
        if prof {
            let names = ["rmsnorm", "qkv-gemm", "rope+attn_wave", "o-gemm", "rms+gu-gemm",
                         "silu", "down-gemm", "tail(norm+lm+argmax)"];
            let tot: f64 = tp.iter().sum();
            println!("== WAVE PROFILE (b={b} tseg={tseg}, sum {:.2} ms)", tot * 1e3);
            for (n, v) in names.iter().zip(tp.iter()) {
                println!("   {n:<20} {:>7.2} ms  {:>5.1}%", v * 1e3, v / tot * 100.0);
            }
        }
        Ok(())
    }

    /// CHUNKED PREFILL: process `tokens` as the chunk starting at absolute position `pos0` of
    /// the sequence in `slot`. The chunk's K/V land in the slab at their absolute positions
    /// (rope_prefill already takes pos0), and attention is `attn_strip` — chunk queries against
    /// slab keys [0, pos0+r]. Only the LAST chunk needs the LM head, so `sample` gates it.
    /// Dense, non-sandwich, hd=128 models: the arm that has an FA-shaped attention to replace.
    pub fn prefill_chunk(&mut self, tokens: &[u32], slot: usize, pos0: usize, sample: bool)
        -> Result<()> {
        let c = self.cfg.clone();
        let t = tokens.len();
        anyhow::ensure!(c.num_experts == 0 && !c.sandwich && c.head_dim == 128,
            "chunked prefill: dense hd=128 arm only");
        anyhow::ensure!(!self.fp8kv, "chunked prefill reads bf16 KV slabs (kv_gather)");
        anyhow::ensure!(self.layers[0].f16_qkv.is_some() || !self.fp8_qkv.is_empty(),
            "init_batch first");
        let bmax = self.batch.as_ref().map(|b| b.max_seq).context("init_batch first")?;
        // The FA arm runs attention over rows [0, pos0+t), so the scratch must cover that —
        // the same size a one-shot prefill of this prompt would need.
        let fa_arm = std::env::var("CHUNKFA").as_deref() != Ok("0") && self.fa;
        let need = if fa_arm { pos0 + t } else { t };
        if self.pf.as_ref().map(|p| p.tmax) < Some(need) {
            self.grow_prefill(need)?;
        }
        self.set_kvoff(Some(slot))?;
        let mut p = self.pf.take().unwrap();
        let r = self.stream.memcpy_htod(tokens, &mut p.tok.slice_mut(0..t))
            .map_err(anyhow::Error::from)
            .and_then(|()| self.prefill_chunk_inner(&mut p, t, slot, pos0, bmax, sample, fa_arm));
        self.pf = Some(p);
        r
    }

    #[allow(clippy::too_many_arguments)]
    fn prefill_chunk_inner(&mut self, p: &mut Prefill, t: usize, slot: usize, pos0: usize,
                           bmax: usize, sample: bool, fa_arm: bool) -> Result<()> {
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let (per, eps) = (nh + 2 * nkv, c.eps);
        let (ti, hi, nhi, nkvi, hdi) = (t as i32, hidden as i32, nh as i32, nkv as i32, hd as i32);
        let (p0i, kvoi) = (pos0 as i32, (slot * bmax) as i32);
        let scale = 1.0f32 / (hd as f32).sqrt();
        let rowsg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
        // qs + acc (f32) + sc + rm/rl/rc, then the staged K and V tiles (bf16, 64 x hd each)
        let ssh = (((32 * hd * 2 + 32 * 64 + 96) * 4) + (2 * 64 * hd * 2)) as u32;
        let esc = c.embed_scale;
        unsafe {
            self.stream.launch_builder(&self.f_embed_p)
                .arg(&self.embed).arg(&mut p.x).arg(&mut p.h).arg(&p.tok).arg(&hi).arg(&ti).arg(&esc)
                .launch(ecfg(t * hidden))?;
        }
        for li in 0..c.layers {
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].attn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (t, self.layers[li].qkv.rows, hidden),
                    self.layers[li].f16_qkv.as_ref().unwrap(), &p.hb16, &mut p.qkv,
                    t, per * hd, hidden)?;
            }
            unsafe {
                let b = self.batch.as_mut().unwrap();
                // FA arm: this chunk's Q/K/V belong at ROWS [pos0, pos0+t) of the buffer, so the
                // kernel gets a view starting there. (The KV-slab write is by absolute position
                // either way.) The strip arm keeps the chunk at rows [0, t).
                let qoff = if fa_arm { pos0 * (nh + 2 * nkv) * hd } else { 0 };
                let mut q16v = p.qkv16.slice_mut(qoff..qoff + t * (nh + 2 * nkv) * hd);
                let bmaxi = bmax as i32;
                self.stream.launch_builder(&self.f_rope_p)
                    .arg(&mut p.qkv).arg(&mut q16v).arg(&mut b.kc[li]).arg(&mut b.vc[li])
                    .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&ti).arg(&self.inv_freq).arg(&self.qkni)
                    .arg(&eps).arg(&p0i).arg(&self.d_kvoff).arg(&ti).arg(&bmaxi)
                    .launch(LaunchConfig { grid_dim: ((t * per).div_ceil(4) as u32, 1, 1),
                                           block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                let bt = self.batch.as_ref().unwrap();
                if fa_arm {
                    // prefix K/V from the slab into the buffer's K/V slots, then ONE FA over
                    // rows [0, pos0+t); only rows [pos0, pos0+t) are consumed downstream.
                    let tall = (pos0 + t) as i32;
                    let tallp = (pos0 + t).next_multiple_of(64) as i32;
                    if pos0 > 0 {
                        self.stream.launch_builder(&self.f_kv_gather)
                            .arg(&bt.kc[li]).arg(&bt.vc[li]).arg(&mut p.qkv16)
                            .arg(&p0i).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&kvoi).arg(&bmaxi)
                            .launch(LaunchConfig { grid_dim: (256, 1, 1), block_dim: (256,1,1),
                                                   shared_mem_bytes: 0 })?;
                    }
                    self.stream.launch_builder(&self.f_vt)
                        .arg(&p.qkv16).arg(&mut p.vt16).arg(&tall).arg(&tallp).arg(&nhi).arg(&nkvi).arg(&hdi)
                        .launch(LaunchConfig { grid_dim: ((pos0 + t).div_ceil(32) as u32,
                                                          (hd / 32) as u32, nkv as u32),
                                               block_dim: (32, 8, 1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_fa)
                        .arg(&p.qkv16).arg(&p.vt16).arg(&mut p.attn16).arg(&tall).arg(&tallp)
                        .arg(&nhi).arg(&nkvi).arg(&scale)
                        .launch(LaunchConfig { grid_dim: ((pos0 + t).div_ceil(64) as u32, nh as u32, 1),
                                               block_dim: (128, 1, 1), shared_mem_bytes: FA_SHARED })?;
                } else {
                    self.stream.launch_builder(&self.f_attn_strip)
                        .arg(&p.qkv16).arg(&bt.kc[li]).arg(&bt.vc[li]).arg(&mut p.attn16)
                        .arg(&ti).arg(&p0i).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&kvoi).arg(&scale).arg(&bmaxi)
                        .launch(LaunchConfig { grid_dim: (t.div_ceil(32) as u32, nh as u32, 1),
                                               block_dim: (256,1,1), shared_mem_bytes: ssh })?;
                }
            }
            if li == 0 && std::env::var("ATTN_DUMP").is_ok() {
                self.stream.synchronize()?;
                let rows: Vec<usize> = [0usize, 1, 5, 31, 32, t - 1].into_iter()
                    .filter(|&r| r < t).collect();
                self.dump_attn("strip", &p.attn16, &rows, nh * hd)?;
            }
            if fa_arm && pos0 > 0 {
                // FA wrote attention for rows [0, pos0+t); the o-GEMM wants only THIS chunk's
                // rows. Stage them at the front of hb16 (free until the FFN norm writes it a
                // few lines below) rather than teaching every GEMM helper about offsets.
                let src = p.attn16.slice(pos0 * nh * hd..(pos0 + t) * nh * hd);
                let mut dst = p.hb16.slice_mut(0..t * nh * hd);
                self.stream.memcpy_dtod(&src, &mut dst)?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                let src: &CudaSlice<u16> = if fa_arm && pos0 > 0 { &p.hb16 } else { &p.attn16 };
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (t, hidden, self.layers[li].o.nblk * 32),
                    self.layers[li].f16_o.as_ref().unwrap(), src, &mut p.h,
                    t, hidden, nh * hd)?;
            }
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].ffn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[li].f16_gu.as_ref().unwrap(), &p.hb16, &mut p.gu16,
                t, 2 * inter, hidden)?;
            let (ii, n8, act) = (inter as i32, (t * inter / 8) as i32, c.gelu as i32);
            unsafe {
                self.stream.launch_builder(&self.f_silu_bf)
                    .arg(&p.gu16).arg(&mut p.g16).arg(&ii).arg(&n8).arg(&act)
                    .launch(ecfg(t * inter / 8))?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans, (t, hidden, inter),
                    self.layers[li].f16_down.as_ref().unwrap(), &p.g16, &mut p.h,
                    t, hidden, inter)?;
            }
        }
        if !sample {
            return Ok(());   // interior chunk: its only product is the KV it wrote
        }
        unsafe {
            self.stream.launch_builder(&self.f_rms_add_b)
                .arg(&mut p.x).arg(&p.h).arg(&self.final_norm)
                .arg(&mut p.attn).arg(&mut self.xs).arg(&hi).arg(&eps).launch(rowsg)?;
        }
        self.stream.memcpy_dtod(&p.attn.slice((t - 1) * hidden..t * hidden), &mut self.h)?;
        Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
            self.f16_lm.as_ref().unwrap(), &mut self.xf16, &self.h, &mut p.logits, 1, c.vocab, hidden)?;
        let v = c.vocab as i32;
        unsafe {
            self.stream.launch_builder(&self.f_argmax)
                .arg(&p.logits).arg(&mut self.d_token).arg(&v)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
        }
        let (dt, df) = (&self.d_token, &mut self.d_firsts);
        self.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
        // The sequential decode head self-feeds from d_pos/d_token. prefill_inner sets d_pos at
        // its tail; omitting it here left the seq path resuming from a STALE position — the
        // first sampled token was right (prefill's own logits) and every token after it was
        // wrong, which reads exactly like a broken attention kernel and is not one.
        let endi = (pos0 + t) as i32;
        unsafe {
            self.stream.launch_builder(&self.f_set_i32)
                .arg(&mut self.d_pos).arg(&endi)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    /// GATED DELTANET, one decode step for `m` slots — the CUDA arm of the hybrid layers.
    ///
    /// `mixed` is the fused qkv projection [m, conv_dim] (the caller's GEMM), `z/b/a` the
    /// three side projections; `state` is [m, nv, dk, dv] and `ring` [m, conv_dim, kn], both
    /// persistent across steps. Writes `core` [m, nv*dv], which the caller feeds to out_proj.
    /// Byte-for-byte the spec in `deltanet.rs`; `dn_gate` checks it against that reference.
    #[allow(clippy::too_many_arguments)]
    pub fn deltanet_step(
        &mut self, state: &mut CudaSlice<f32>, ring: &mut CudaSlice<f32>,
        mixed: &mut CudaSlice<f32>, z: &CudaSlice<f32>, b: &CudaSlice<f32>, a: &CudaSlice<f32>,
        a_log: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, norm_w: &CudaSlice<f32>,
        conv_w: &CudaSlice<f32>, core: &mut CudaSlice<f32>,
        nk: usize, nv: usize, dk: usize, dv: usize, kn: usize, m: usize, eps: f32,
    ) -> Result<()> {
        anyhow::ensure!(dv <= 1024 && dk <= dv.max(1024), "deltanet: dv must fit one block");
        let conv_dim = 2 * nk * dk + nv * dv;
        let (cdi, kni, mi) = (conv_dim as i32, kn as i32, m as i32);
        let (nki, nvi, dki, dvi) = (nk as i32, nv as i32, dk as i32, dv as i32);
        unsafe {
            self.stream.launch_builder(&self.f_dn_conv)
                .arg(&mut *mixed).arg(&mut *ring).arg(conv_w).arg(&cdi).arg(&kni).arg(&mi)
                .launch(ecfg(m * conv_dim))?;
            // shared = S tile + q + k + a 32-slot reduction scratch
            let shb = ((dk * dv + 2 * dk + 32) * 4) as u32;
            self.stream.launch_builder(&self.f_dn_step)
                .arg(&mut *state).arg(&*mixed).arg(z).arg(b).arg(a)
                .arg(a_log).arg(dt_bias).arg(norm_w).arg(&mut *core)
                .arg(&nki).arg(&nvi).arg(&dki).arg(&dvi).arg(&eps)
                .launch(LaunchConfig { grid_dim: (nv as u32, m as u32, 1),
                                       block_dim: (dv as u32, 1, 1), shared_mem_bytes: shb })?;
        }
        Ok(())
    }

    /// DN_GATE=1: run `deltanet_step` against the CPU reference on pseudo-random inputs and
    /// report the max absolute divergence of `core` AND of the evolved state. A linear-attention
    /// layer carries its state forward forever, so a gate that only checks the output would
    /// pass a kernel whose state drifts — this checks both, over several steps.
    pub fn dn_gate(&mut self, nk: usize, nv: usize, dk: usize, dv: usize, kn: usize,
                   steps: usize) -> Result<(f32, f32)> {
        use crate::deltanet::{DeltaNetRef, silu};
        let conv_dim = 2 * nk * dk + nv * dv;
        let mut seed = 0x2545F491_4F6CDD1Du64;
        let mut rnd = move || {
            seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17;
            ((seed >> 40) as f32 / 8388608.0) - 1.0
        };
        let r = DeltaNetRef {
            nk, nv, dk, dv, kernel: kn, eps: 1e-6,
            w_qkv: vec![], w_z: vec![], w_b: vec![], w_a: vec![], w_out: vec![],
            conv_w: (0..conv_dim * kn).map(|_| rnd() * 0.5).collect(),
            a_log: (0..nv).map(|_| rnd()).collect(),
            dt_bias: (0..nv).map(|_| rnd()).collect(),
            norm_w: (0..dv).map(|_| 1.0 + rnd() * 0.2).collect(),
        };
        let mut st = r.fresh_state();
        let st_d = self.stream.clone();
        let mut d_state = st_d.alloc_zeros::<f32>(nv * dk * dv)?;
        let mut d_ring = st_d.alloc_zeros::<f32>(conv_dim * kn)?;
        let d_cw = st_d.memcpy_stod(&r.conv_w)?;
        let d_al = st_d.memcpy_stod(&r.a_log)?;
        let d_dt = st_d.memcpy_stod(&r.dt_bias)?;
        let d_nw = st_d.memcpy_stod(&r.norm_w)?;
        let mut d_core = st_d.alloc_zeros::<f32>(nv * dv)?;
        let (mut worst_core, mut worst_state) = (0f32, 0f32);
        for _ in 0..steps {
            let mixed: Vec<f32> = (0..conv_dim).map(|_| rnd()).collect();
            let z: Vec<f32> = (0..nv * dv).map(|_| rnd()).collect();
            let bb: Vec<f32> = (0..nv).map(|_| rnd()).collect();
            let aa: Vec<f32> = (0..nv).map(|_| rnd()).collect();
            let want = r.core(&mut st, &mixed, &z, &bb, &aa);
            let mut d_mixed = st_d.memcpy_stod(&mixed)?;
            let d_z = st_d.memcpy_stod(&z)?;
            let d_b = st_d.memcpy_stod(&bb)?;
            let d_a = st_d.memcpy_stod(&aa)?;
            self.deltanet_step(&mut d_state, &mut d_ring, &mut d_mixed, &d_z, &d_b, &d_a,
                               &d_al, &d_dt, &d_nw, &d_cw, &mut d_core,
                               nk, nv, dk, dv, kn, 1, 1e-6)?;
            self.sync()?;
            let got: Vec<f32> = self.stream.memcpy_dtov(&d_core)?;
            let gs: Vec<f32> = self.stream.memcpy_dtov(&d_state)?;
            for (g, w) in got.iter().zip(&want) { worst_core = worst_core.max((g - w).abs()); }
            for (g, w) in gs.iter().zip(&st.s) { worst_state = worst_state.max((g - w).abs()); }
        }
        let _ = silu(0.0);
        Ok((worst_core, worst_state))
    }

    /// Whether `step_mixed` can serve this model: the same dense hd=128 bf16-mirror arm as
    /// chunked prefill (mixed steps run EVERY row — decode included — through the bf16 GEMMs;
    /// pure-decode steps keep their own arm and graphs).
    pub fn mix_capable(&self) -> bool {
        self.cfg.num_experts == 0 && !self.cfg.sandwich && self.cfg.head_dim == 128
            && !self.cfg.v_norm && self.layers[0].f16_qkv.is_some() && self.fa
            && !self.fp8kv   // kv_gather + the chunk FA read bf16 slabs
    }

    /// HIGHEST free slot. Mixed prefill parks its cursor here on purpose: `qk_rope_b` writes K/V
    /// for every row below the decode width unconditionally, so a mid-prefill slot INSIDE the
    /// width would get garbage K/V written over the chunk's rows. Above the width it is
    /// untouchable; compaction never moves a sequence upward, so it stays above until it
    /// completes and joins the live set (whereupon compaction pulls it down).
    pub fn free_slot_high(&self) -> Option<usize> {
        self.batch.as_ref()?.occupied.iter().rposition(|o| !o)
    }

    /// QUIESCE a slot reserved for an in-progress prefill: zero its device `active` flag.
    /// move_slot deliberately leaves a vacated slot's active=1 (nothing used to read it), but
    /// the rope kernel's inactive-row guard keys off it, and a stale 1 on a mid-prefill slot
    /// would re-open the garbage-KV-write hole the guard exists to close.
    pub fn quiesce_slot(&mut self, slot: usize) -> Result<()> {
        let b = self.batch.as_mut().context("init_batch first")?;
        anyhow::ensure!(slot < b.mmax, "slot out of range");
        self.stream.memcpy_htod(&[0u32], &mut b.active.slice_mut(slot..slot + 1))?;
        Ok(())
    }

    /// Stage a mixed-prefill prompt's ids on device, keyed by slot. One pageable upload at
    /// admission (where such uploads already live) so each mixed step's chunk read is a device
    /// slice instead of an htod that drains the stream mid-pipeline. Also quiesces the slot.
    pub fn mix_stage(&mut self, slot: usize, ids: &[u32]) -> Result<()> {
        self.quiesce_slot(slot)?;
        let d = self.stream.memcpy_stod(ids)?;
        self.mixtok.insert(slot, d);
        Ok(())
    }

    /// Whether overlapped prefill (a captured one-shot replaying on stream2 while the decode
    /// graph free-runs on the main stream) can serve this model. Beyond mix_capable, the fp8
    /// weight arms are excluded: their prefill and decode share the fp8_x/amax scratch, which
    /// two streams would race on.
    pub fn overlap_capable(&self) -> bool {
        self.mix_capable() && self.fp8_qkv.is_empty()
    }

    /// Run `f` with the engine pointed at the OVERLAP stream: self.stream, the cuBLAS handle's
    /// stream, the stream2 graph map, and the shared tail scratch (xf16/h/xs/Lt workspace) are
    /// all swapped for the duration. Everything the existing prefill path enqueues — captures
    /// included — lands on stream2 with no code threaded through.
    fn on_stream2<R>(&mut self, f: impl FnOnce(&mut Self) -> Result<R>) -> Result<R> {
        let main = self.stream.clone();
        self.stream = self.stream2.clone();
        std::mem::swap(&mut self.pfslot_graphs, &mut self.pfslot_graphs2);
        std::mem::swap(&mut self.xf16, &mut self.xf16b);
        std::mem::swap(&mut self.h, &mut self.hb2);
        std::mem::swap(&mut self.xs, &mut self.xsb);
        std::mem::swap(&mut self.ltws, &mut self.ltws2);
        unsafe {
            let _ = cudarc::cublas::result::set_stream(*self.blas.handle(),
                                                       self.stream.cu_stream() as _);
        }
        let r = f(self);
        unsafe {
            let _ = cudarc::cublas::result::set_stream(*self.blas.handle(), main.cu_stream() as _);
        }
        std::mem::swap(&mut self.ltws, &mut self.ltws2);
        std::mem::swap(&mut self.xs, &mut self.xsb);
        std::mem::swap(&mut self.h, &mut self.hb2);
        std::mem::swap(&mut self.xf16, &mut self.xf16b);
        std::mem::swap(&mut self.pfslot_graphs, &mut self.pfslot_graphs2);
        self.stream = main;
        r
    }

    /// Start a long prefill on the overlap stream and return its completion event. The caller
    /// guarantees the mixed-scheduling slot invariants (reserved above the decode width, device
    /// active quiesced) and that no other Prefill-buffer user runs until the event completes.
    /// Cold lengths warm up eagerly on stream2 (cuBLAS workspace allocation aborts captures),
    /// then capture; warm lengths are one htod + graph replay — the same flow as the staged
    /// admission path, just on the other stream.
    pub fn prefill_overlap_start(&mut self, tokens: &[u32], slot: usize)
        -> Result<cudarc::driver::CudaEvent> {
        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
        anyhow::ensure!(self.overlap_capable(), "overlap prefill: bf16-mirror dense arm only");
        let t = tokens.len();
        // Grow BEFORE entering stream2 work: grow reallocates the Prefill buffers and clears
        // both graph maps, which must never happen mid-swap.
        if self.pf.as_ref().map(|p| p.tmax) < Some(t) {
            self.grow_prefill(t)?;
        }
        self.on_stream2(|eng| {
            eng.set_kvoff(Some(slot))?;
            if !eng.pfslot_graphs.contains_key(&t) {   // swapped: this IS the stream2 map
                eng.prefill_slot(tokens, Some(slot))?;
                eng.sync()?;                           // stream2 under the swap
                let mut p = eng.pf.take().unwrap();
                eng.stream.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
                let r = eng.prefill_inner(&mut p, tokens, t, Some(slot), t);
                let g = eng.stream.end_capture(
                    CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
                eng.pf = Some(p);
                r.context("during overlap prefill graph capture")?;
                let g = g?.context("stream capture produced no graph")?;
                g.upload()?;
                eng.pfslot_graphs.insert(t, g);
                // the eager warm-up above already wrote this request's KV and sampled its token
            } else {
                {
                    let mut p = eng.pf.take().context("no prefill state")?;
                    let r = eng.stream.memcpy_htod(tokens, &mut p.tok.slice_mut(0..t));
                    eng.pf = Some(p);
                    r?;
                }
                eng.pfslot_graphs[&t].launch()?;
            }
            let (dt, df) = (&eng.d_token, &mut eng.d_firsts);
            eng.stream.memcpy_dtod(dt, &mut df.slice_mut(slot..slot + 1))?;
            eng.stream.record_event(None).map_err(anyhow::Error::from)
        })
    }

    /// Order the MAIN stream after an overlap prefill's completion event, so the slot's first
    /// decode step cannot read its KV slab before the flight finished writing it.
    pub fn wait_overlap(&self, ev: &cudarc::driver::CudaEvent) -> Result<()> {
        self.stream.wait(ev)?;
        Ok(())
    }

    pub fn mix_unstage(&mut self, slot: usize) {
        self.mixtok.remove(&slot);
    }

    /// One MIXED step: `m` decode rows (slots [0, m), each at its own position) PLUS one prefill
    /// chunk — rows [pos0, pos0+n) of the prompt staged for `slot` — through ONE layer stack.
    ///
    /// The GEMMs concatenate trivially ([m+n, hidden], decode rows first); everything else
    /// dispatches per row-group: decode rows through qk_rope_b + attn_partial_b/reduce against
    /// their slabs, chunk rows through rope_prefill + the gathered-prefix FA (byte-identical to
    /// prefill_chunk's arm). The decode rows ride the chunk's weight streaming for free — that
    /// is the entire point: a standalone decode step is weight-bound, and the chunk pass is
    /// already reading every weight.
    ///
    /// Runs EAGERLY — never captured. The decode CUDA graph bakes `m`, and a (m, n) pair space
    /// would explode; a mixed step already amortizes a whole chunk, so eager launch cost is
    /// noise. Decode-row sampling lands in b.tok/b.pos/b.out exactly as a graph step would, so
    /// pure-decode graphs resume seamlessly next step. With `sample`, the chunk's last row runs
    /// the LM head into d_firsts[slot] and d_pos/d_token, exactly like prefill_chunk's tail.
    ///
    /// `m == 0` is legal (pure chunk) but the scheduler routes that case to the graphed one-shot
    /// prefill instead, which is strictly faster (captured graph + tuned plans).
    pub fn step_mixed(&mut self, slot: usize, pos0: usize, n: usize, sample: bool, m: usize)
        -> Result<()> {
        anyhow::ensure!(self.mix_capable(), "step_mixed: dense hd=128 bf16-mirror arm only");
        anyhow::ensure!(n > 0, "step_mixed: empty chunk");
        let bmax = self.batch.as_ref().map(|b| b.max_seq).context("init_batch first")?;
        anyhow::ensure!(pos0 + n <= bmax, "step_mixed: chunk beyond the KV slab");
        anyhow::ensure!(self.mixtok.get(&slot).map(|s| s.len()) >= Some(pos0 + n),
            "step_mixed: mix_stage first");
        // Scratch must cover BOTH row spaces: [m+n) combined rows for the GEMM buffers and
        // [0, pos0+n) absolute rows for the FA's qkv16/vt16/attn16.
        let need = (pos0 + n).max(m + n);
        if self.pf.as_ref().map(|p| p.tmax) < Some(need) {
            self.grow_prefill(need)?;
        }
        self.set_kvoff(Some(slot))?;
        let mut p = self.pf.take().unwrap();
        let mut b = self.batch.take().unwrap();
        let r = self.step_mixed_inner(&mut p, &mut b, slot, pos0, n, sample, m);
        self.pf = Some(p);
        self.batch = Some(b);
        r
    }

    #[allow(clippy::too_many_arguments)]
    fn step_mixed_inner(&mut self, p: &mut Prefill, b: &mut Batch, slot: usize, pos0: usize,
                        n: usize, sample: bool, m: usize) -> Result<()> {
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let (per, eps) = (nh + 2 * nkv, c.eps);
        let (bmax, nsplit) = (b.max_seq, b.nsplit);
        let rr = m + n;
        let (ni, hi, nhi, nkvi, hdi, nsi, bmi) =
            (n as i32, hidden as i32, nh as i32, nkv as i32, hd as i32, nsplit as i32, bmax as i32);
        let (mi, p0i, kvoi) = (m as i32, pos0 as i32, (slot * bmax) as i32);
        let roti = c.rot_dim as i32;
        let scale = if c.attn_scale_one { 1.0f32 } else { 1.0 / (hd as f32).sqrt() };
        let win0 = 0i32;
        let esc = c.embed_scale;
        // 256-thread row blocks like the prefill path (chunk rows dominate), and the SAME block
        // size for the decode rows in this step — the norm's reduction order must not depend on
        // which rows share the launch.
        let rowsg = LaunchConfig { grid_dim: (rr as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
        unsafe {
            if m > 0 {
                self.stream.launch_builder(&self.f_embed_b)
                    .arg(&self.embed).arg(&mut p.x).arg(&mut p.h).arg(&b.tok).arg(&hi).arg(&mi).arg(&esc)
                    .launch(ecfg(m * hidden))?;
            }
            let stok = self.mixtok.get(&slot).context("mix_stage first")?;
            let ctok = stok.slice(pos0..pos0 + n);
            let xl = p.x.len();
            let mut cx = p.x.try_slice_mut(m * hidden..rr * hidden).with_context(|| format!("mixed cx m={m} rr={rr} len={xl}"))?;
            let hl = p.h.len();
            let mut ch = p.h.try_slice_mut(m * hidden..rr * hidden).with_context(|| format!("mixed ch m={m} rr={rr} len={hl}"))?;
            self.stream.launch_builder(&self.f_embed_p)
                .arg(&self.embed).arg(&mut cx).arg(&mut ch).arg(&ctok).arg(&hi).arg(&ni).arg(&esc)
                .launch(ecfg(n * hidden))?;
        }
        for li in 0..c.layers {
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].attn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (rr, self.layers[li].qkv.rows, hidden),
                    self.layers[li].f16_qkv.as_ref().unwrap(), &p.hb16, &mut p.qkv,
                    rr, per * hd, hidden)?;
            }
            unsafe {
                // decode rows: rope at each slot's own position, K/V appended to its slab
                if m > 0 {
                    self.stream.launch_builder(&self.f_rope_b)
                        .arg(&mut p.qkv).arg(&mut b.kc[li]).arg(&mut b.vc[li])
                        .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&self.inv_freq)
                        .arg(&self.qkni).arg(&eps).arg(&bmi).arg(&b.active).arg(&roti)
                        .launch(LaunchConfig { grid_dim: ((m * per) as u32, 1, 1),
                                               block_dim: (128, 1, 1), shared_mem_bytes: 0 })?;
                }
                // chunk rows: rope at absolute positions pos0+i, K/V into `slot`'s slab, roped Q
                // into qkv16 at ABSOLUTE rows for the gathered-prefix FA
                let qoff = pos0 * per * hd;
                let q16l = p.qkv16.len();
                let mut q16v = p.qkv16.try_slice_mut(qoff..qoff + n * per * hd).with_context(|| format!("mixed q16v qoff={qoff} n={n} len={q16l}"))?;
                let qkl = p.qkv.len();
                let mut cqkv = p.qkv.try_slice_mut(m * per * hd..rr * per * hd).with_context(|| format!("mixed cqkv m={m} rr={rr} len={qkl}"))?;
                self.stream.launch_builder(&self.f_rope_p)
                    .arg(&mut cqkv).arg(&mut q16v).arg(&mut b.kc[li]).arg(&mut b.vc[li])
                    .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&ni).arg(&self.inv_freq).arg(&self.qkni)
                    .arg(&eps).arg(&p0i).arg(&self.d_kvoff).arg(&ni).arg(&bmi)
                    .launch(LaunchConfig { grid_dim: ((n * per).div_ceil(4) as u32, 1, 1),
                                           block_dim: (128, 1, 1), shared_mem_bytes: 0 })?;
                // decode attention: partial + reduce straight into rows [0, m) of the o-GEMM
                // input (its previous content — the qkv norm — is dead once the GEMM above ran)
                if m > 0 {
                    if self.attnt_now() {
                        let yt = (self.attnb_tgt.div_ceil((m * nkv).max(1)))
                            .clamp(1, nsplit.div_ceil(4).max(1)) as u32;
                        self.stream.launch_builder(&self.f_attn_part_t)
                            .arg(&p.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                            .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                            .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                            .arg(&scale).arg(&win0)
                            .launch(LaunchConfig { grid_dim: ((m * nkv) as u32, yt, 1),
                                                   block_dim: (128, 1, 1),
                                                   shared_mem_bytes: 143616 })?;
                    } else {
                    self.stream.launch_builder(&self.f_attn_part_b)
                        .arg(&p.qkv).arg(&b.kc[li]).arg(&b.vc[li])
                        .arg(&mut b.pm).arg(&mut b.pl).arg(&mut b.pacc)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&b.pos).arg(&nsi).arg(&bmi)
                        .arg(&scale).arg(&win0)
                        .launch(LaunchConfig { grid_dim: ((m * nh) as u32,
                                                          (self.attnb_tgt.div_ceil(m * nh)).clamp(4, nsplit.max(4)) as u32,
                                                          1),
                                               block_dim: (256, 1, 1), shared_mem_bytes: ((hd + 256) * 4) as u32 })?;
                    }
                    let hbl = p.hb16.len();
                    let mut dred = p.hb16.try_slice_mut(0..m * nh * hd).with_context(|| format!("mixed dred m={m} len={hbl}"))?;
                    self.stream.launch_builder(&self.f_attn_red_b16)
                        .arg(&b.pm).arg(&b.pl).arg(&b.pacc).arg(&mut dred).arg(&hdi).arg(&nsi)
                        .arg(&nhi).arg(&b.pos).arg(&win0)
                        .launch(LaunchConfig { grid_dim: ((m * nh) as u32, 1, 1),
                                               block_dim: (128, 1, 1), shared_mem_bytes: 0 })?;
                }
                // chunk attention: prefix K/V gathered from the slab, one FA over [0, pos0+n)
                let tall = (pos0 + n) as i32;
                let tallp = (pos0 + n).next_multiple_of(64) as i32;
                if pos0 > 0 {
                    self.stream.launch_builder(&self.f_kv_gather)
                        .arg(&b.kc[li]).arg(&b.vc[li]).arg(&mut p.qkv16)
                        .arg(&p0i).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&kvoi).arg(&bmi)
                        .launch(LaunchConfig { grid_dim: (256, 1, 1), block_dim: (256, 1, 1),
                                               shared_mem_bytes: 0 })?;
                }
                self.stream.launch_builder(&self.f_vt)
                    .arg(&p.qkv16).arg(&mut p.vt16).arg(&tall).arg(&tallp).arg(&nhi).arg(&nkvi).arg(&hdi)
                    .launch(LaunchConfig { grid_dim: ((pos0 + n).div_ceil(32) as u32,
                                                      (hd / 32) as u32, nkv as u32),
                                           block_dim: (32, 8, 1), shared_mem_bytes: 0 })?;
                // Same wgmma dispatch ladder as the one-shot prefill — a mixed pass must not pay
                // a slower attention kernel than the pass it replaces (ws > hop > mma).
                let hop = self.hop && nh % 2 == 0 && (nh / nkv) % 2 == 0;
                if hop && self.ws {
                    self.stream.memset_zeros(&mut self.fw_ticket)?;
                    self.stream.launch_builder(&self.f_fa_ws)
                        .arg(&p.qkv16).arg(&p.vt16).arg(&mut p.attn16).arg(&tall).arg(&tallp)
                        .arg(&nhi).arg(&nkvi).arg(&scale)
                        .arg(&mut self.fw_ticket)
                        .launch(LaunchConfig { grid_dim: (self.sm_count as u32, 1, 1),
                                               block_dim: (384, 1, 1),
                                               shared_mem_bytes: FW_SHARED })?;
                } else {
                    let (fk, fsh, gy, bx) = if hop { (&self.f_fa_hop, FH_SHARED, nh / 2, 256) }
                                            else { (&self.f_fa, FA_SHARED, nh, 128) };
                    self.stream.launch_builder(fk)
                        .arg(&p.qkv16).arg(&p.vt16).arg(&mut p.attn16).arg(&tall).arg(&tallp)
                        .arg(&nhi).arg(&nkvi).arg(&scale)
                        .launch(LaunchConfig { grid_dim: ((pos0 + n).div_ceil(64) as u32, gy as u32, 1),
                                               block_dim: (bx, 1, 1), shared_mem_bytes: fsh })?;
                }
            }
            // chunk rows of the FA output into rows [m, m+n) of the o-GEMM input
            {
                let src = p.attn16.slice(pos0 * nh * hd..(pos0 + n) * nh * hd);
                let hbl2 = p.hb16.len();
                let mut dst = p.hb16.try_slice_mut(m * nh * hd..rr * nh * hd).with_context(|| format!("mixed adst m={m} rr={rr} len={hbl2}"))?;
                self.stream.memcpy_dtod(&src, &mut dst)?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (rr, hidden, self.layers[li].o.nblk * 32),
                    self.layers[li].f16_o.as_ref().unwrap(), &p.hb16, &mut p.h,
                    rr, hidden, nh * hd)?;
            }
            unsafe {
                self.stream.launch_builder(&self.f_rms_bf)
                    .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].ffn_norm)
                    .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
            }
            Self::gemm_pre_c16(&self.stream, &self.blas,
                self.layers[li].f16_gu.as_ref().unwrap(), &p.hb16, &mut p.gu16,
                rr, 2 * inter, hidden)?;
            let (ii, n8, act) = (inter as i32, (rr * inter / 8) as i32, c.gelu as i32);
            unsafe {
                self.stream.launch_builder(&self.f_silu_bf)
                    .arg(&p.gu16).arg(&mut p.g16).arg(&ii).arg(&n8).arg(&act)
                    .launch(ecfg(rr * inter / 8))?;
            }
            {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans, (rr, hidden, inter),
                    self.layers[li].f16_down.as_ref().unwrap(), &p.g16, &mut p.h,
                    rr, hidden, inter)?;
            }
        }
        unsafe {
            self.stream.launch_builder(&self.f_rms_add_b)
                .arg(&mut p.x).arg(&p.h).arg(&self.final_norm)
                .arg(&mut p.attn).arg(&mut self.xs).arg(&hi).arg(&eps).launch(rowsg)?;
        }
        if m > 0 {
            // decode-row tail exactly as a batch step: LM head over rows [0, m), argmax into
            // b.tok, advance — so the next (pure or mixed) step picks up where a graph step
            // would. The head reads p.attn DIRECTLY (decode rows are its first m rows) — a
            // bounce through b.h looked harmless but b.h/b.abuf swap identities every batch
            // step, and on odd parity b.h was the (possibly smaller) abuf buffer.
            Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
                self.f16_lm.as_ref().unwrap(), &mut self.xf16, &p.attn, &mut b.logits,
                m, c.vocab, hidden)?;
            let v = c.vocab as i32;
            unsafe {
                self.stream.launch_builder(&self.f_argmax_b)
                    .arg(&b.logits).arg(&mut b.tok).arg(&v)
                    .launch(LaunchConfig { grid_dim: (m as u32, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                self.stream.launch_builder(&self.f_advance_b)
                    .arg(&mut b.pos).arg(&b.tok).arg(&mut b.out).arg(&bmi).arg(&mi).arg(&b.active)
                    .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (64, 1, 1), shared_mem_bytes: 0 })?;
            }
        }
        if sample {
            // final chunk: the request's first token, via the same tail as prefill_chunk
            self.stream.memcpy_dtod(&p.attn.slice((rr - 1) * hidden..rr * hidden), &mut self.h)?;
            Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
                self.f16_lm.as_ref().unwrap(), &mut self.xf16, &self.h, &mut p.logits,
                1, c.vocab, hidden)?;
            let v = c.vocab as i32;
            unsafe {
                self.stream.launch_builder(&self.f_argmax)
                    .arg(&p.logits).arg(&mut self.d_token).arg(&v)
                    .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
            }
            let (dt, df) = (&self.d_token, &mut self.d_firsts);
            self.stream.memcpy_dtod(dt, &mut df.try_slice_mut(slot..slot + 1).with_context(|| format!("mixed dfirst slot={slot} len={}", 256))?)?;
            let endi = (pos0 + n) as i32;
            unsafe {
                self.stream.launch_builder(&self.f_set_i32)
                    .arg(&mut self.d_pos).arg(&endi)
                    .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 })?;
            }
        }
        Ok(())
    }

    /// Prefill into a batch slot's KV slab (`None` = the single-stream cache).
    pub fn prefill_slot(&mut self, tokens: &[u32], slot: Option<usize>) -> Result<()> {
        self.set_kvoff(slot)?;
        anyhow::ensure!(self.layers[0].f16_qkv.is_some() || !self.fp8_qkv.is_empty(),
            "call init_batch first (f16 mirrors or fp8 copies)");
        if slot.is_none() && self.batch.is_none() {
            self.grow_seq_kv()?;
        }
        let c = self.cfg.clone();
        let t = tokens.len();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let per = nh + 2 * nkv;
        if self.pf.as_ref().map(|p| p.tmax) < Some(t) {
            self.grow_prefill(t)?;
        }        let mut p = self.pf.take().unwrap();
        let r = self.stream.memcpy_htod(tokens, &mut p.tok)
            .map_err(anyhow::Error::from)
            .and_then(|()| self.prefill_inner(&mut p, tokens, t, slot, t));
        self.pf = Some(p);
        r
    }

    /// Device-side KV offset for the next prefill; a kernel-parameter write so it is safe both
    /// inside and outside graph capture.
    fn set_kvoff(&mut self, slot: Option<usize>) -> Result<()> {
        let off = match slot {
            Some(sl) => (sl * self.batch.as_ref().map(|b| b.max_seq).unwrap_or(0)) as i32,
            None => 0,
        };
        unsafe {
            self.stream.launch_builder(&self.f_set_i32)
                .arg(&mut self.d_kvoff).arg(&off)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    /// `tseg == t`: ordinary single-prompt prefill. `tseg < t`: WAVE mode — the rows are
    /// t/tseg independent segments; rope reads per-segment KV offsets from d_kvoffs, softmax
    /// masks to segments, the FA arm is bypassed (its kernels have no segment mask), and the
    /// caller owns the tail (per-segment last rows -> LM head -> d_firsts).
    fn prefill_inner(&mut self, p: &mut Prefill, tokens: &[u32], t: usize, slot: Option<usize>,
                     tseg: usize) -> Result<()> {
        let c = self.cfg.clone();
        let (hd, nh, nkv, hidden, inter) = (c.head_dim, c.n_heads, c.n_kv, c.hidden, c.intermediate);
        let (per, eps) = (nh + 2 * nkv, c.eps);
        let (ti, hi, nhi, nkvi, hdi) = (t as i32, hidden as i32, nh as i32, nkv as i32, hd as i32);
        // NOTE: p.tok is filled by the caller, NOT here — a graph replay must be able to run a
        // different prompt of the same length, and a host-to-device copy captured into the graph
        // would replay the tokens that were current at capture time.
        // PREFILL_PROF=1 attributes prefill time per stage. Costs one sync per stage per layer
        // (~5 us x 28 x 8), which is noise at 36 ms but tells us exactly what to optimise instead
        // of guessing — the same discipline that localised the f16-scaling bug.
        // NOTE — CUDA events cannot be used to time stages here. `new_event` returns
        // CUDA_ERROR_INVALID_HANDLE because this context runs with `disable_event_tracking()`, which
        // the flash-attention graph capture requires (cudarc's per-slice event tracking otherwise
        // aborts capture with STREAM_CAPTURE_ISOLATION). The two are mutually exclusive.
        //
        // The question events were meant to settle — whether "438 TFLOP/s in context vs 641
        // isolated" is real or an artifact of the syncs below — is answered by arithmetic instead:
        // PROF sums to 23.13 ms against a clean 21.6-22.4 ms, so ~280 syncs cost ~3-5 us each and
        // gate|up's share is ~0.11 ms of its 6.33. Real in-context rate is ~455 TFLOP/s. THE GAP IS
        // REAL, and prefill's remaining 5% is a genuine GEMM-efficiency problem, not mismeasurement.
        let prof = std::env::var("PREFILL_PROF").is_ok();
        let mut tp = [0f64; 10];
        let mut tk = std::time::Instant::now();
        macro_rules! mark { ($i:expr) => { if prof {
            self.stream.synchronize()?;
            let n = std::time::Instant::now();
            tp[$i] += n.duration_since(tk).as_secs_f64();
            tk = n;
        } } }
        let rowsg = LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };

        // FP8 prefill plans are row-count-specific; build all four sets before the layer loop.
        self.fp8_prefill_plans(t)?;
        // bf16 tuned plans for this row count (no-op when already built or under fp8/MoE)
        if self.fp8_gu.is_empty() {
            self.bf16_prefill_plans(t)?;
        }
        let esc = c.embed_scale;
        unsafe {
            self.stream.launch_builder(&self.f_embed_p)
                .arg(&self.embed).arg(&mut p.x).arg(&mut p.h).arg(&p.tok).arg(&hi).arg(&ti).arg(&esc)
                .launch(ecfg(t * hidden))?;
        }
        for li in 0..c.layers {
            // per-layer geometry (Gemma4); uniform models see cfg's values
            let (lhd, lnkv, lwin, lrope2) = c.lgeom(li);
            let (nkvi, hdi) = (lnkv as i32, lhd as i32);
            let per = nh + 2 * lnkv;
            let ascale = if c.attn_scale_one { 1.0f32 } else { 1.0 / (lhd as f32).sqrt() };
            unsafe {
                if c.sandwich {
                    let (postw, pls) = if li == 0 {
                        (&self.layers[0].attn_norm, 1.0f32)
                    } else {
                        (self.layers[li - 1].post_ffw_norm.as_ref().unwrap(), self.layers[li - 1].ls)
                    };
                    self.stream.launch_builder(&self.f_sandwich)
                        .arg(&mut p.x).arg(&p.h).arg(postw).arg(&self.layers[li].attn_norm)
                        .arg(&mut p.hb16).arg(&hi).arg(&eps).arg(&pls).launch(rowsg)?;
                } else {
                    self.stream.launch_builder(&self.f_rms_bf)
                        .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].attn_norm)
                        .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
                }
            }
            mark!(0);
            if !self.fp8_qkv.is_empty() {
                let nx = (t * hidden) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.hb16).arg(&mut self.fp8_part).arg(&nx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.hb16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                        .launch(ecfg((t * hidden).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_qkv[li], ltws2,
                    1 << 20, &self.fp8_qkv[li].0, &self.fp8_px, &mut p.qkv, 1.0)?;
            } else {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                    (t, self.layers[li].qkv.rows, hidden),
                    self.layers[li].f16_qkv.as_ref().unwrap(), &p.hb16, &mut p.qkv,
                    t, per * lhd, hidden)?;
            }
            mark!(1);
            let pos0 = 0i32;
            // KV target: a batch slot's slab, or the single-stream cache. Once a batch exists,
            // "no slot" means SLOT 0 — the sequential step reads/writes batch slot 0 (same
            // addressing), so a None-prefill into self.kc would strand the prompt where the seq
            // path never looks.
            let (kvoff, use_batch) = match slot {
                Some(sl) => ((sl * self.batch.as_ref().map(|b| b.max_seq).unwrap_or(0)) as i32, true),
                None => (0i32, self.batch.is_some()),
            };
            unsafe {
                if c.v_norm {
                    self.stream.launch_builder(&self.f_vderive)
                        .arg(&mut p.qkv).arg(&nhi).arg(&nkvi).arg(&hdi).arg(&eps).arg(&ti)
                        .launch(LaunchConfig { grid_dim: ((t * lnkv) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                }
                let invf = if lrope2 { &self.inv_freq2 } else { &self.inv_freq };
                if use_batch {
                    let b = self.batch.as_mut().unwrap();
                    let tsegi = tseg as i32;
                    let bmaxi = b.max_seq as i32;
                    let kvoffb = if tseg == t { &self.d_kvoff } else { &self.d_kvoffs };
                    self.stream.launch_builder(if self.fp8kv { &self.f_rope_p_f8 } else { &self.f_rope_p })
                        .arg(&mut p.qkv).arg(&mut p.qkv16).arg(&mut b.kc[li]).arg(&mut b.vc[li])
                        .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                        .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&ti).arg(invf).arg(&self.qkni).arg(&eps)
                        .arg(&pos0).arg(kvoffb).arg(&tsegi).arg(&bmaxi)
                        .launch(LaunchConfig { grid_dim: ((t * per).div_ceil(4) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                } else {
                let smaxi = self.max_seq as i32;
                self.stream.launch_builder(&self.f_rope_p)
                    .arg(&mut p.qkv).arg(&mut p.qkv16).arg(&mut self.kc[li]).arg(&mut self.vc[li])
                    .arg(&self.layers[li].q_norm).arg(&self.layers[li].k_norm)
                    .arg(&nhi).arg(&nkvi).arg(&hdi).arg(&ti).arg(invf).arg(&self.qkni).arg(&eps).arg(&pos0).arg(&self.d_kvoff).arg(&ti).arg(&smaxi)
                    .launch(LaunchConfig { grid_dim: ((t * per).div_ceil(4) as u32, 1, 1), block_dim: (128,1,1), shared_mem_bytes: 0 })?;
                }
            }
            mark!(2);
            if self.fa && lhd == 128 && lwin == 0 && !c.attn_scale_one && tseg == t {
                // one fused kernel: no score matrix, no separate softmax, causal blocks skipped
                unsafe {
                    let tpi = t.next_multiple_of(64) as i32;
                    self.stream.launch_builder(&self.f_vt)
                        .arg(&p.qkv16).arg(&mut p.vt16).arg(&ti).arg(&tpi).arg(&nhi).arg(&nkvi).arg(&hdi)
                        .launch(LaunchConfig { grid_dim: (t.div_ceil(32) as u32, (hd / 32) as u32,
                                                          nkv as u32),
                                               block_dim: (32, 8, 1), shared_mem_bytes: 0 })?;
                    // The wgmma kernel pairs two query heads per block to halve K/V traffic, which
                    // is only sound when both land on the same kv head — i.e. an even GQA ratio.
                    let hop = self.hop && nh % 2 == 0 && (nh / nkv) % 2 == 0;
                    if hop && self.ws {
                        // Persistent blocks pulling (query block, head pair) tickets; the counter
                        // is device memory so the zeroing is stream-ordered and graph-capturable.
                        self.stream.memset_zeros(&mut self.fw_ticket)?;
                        self.stream.launch_builder(&self.f_fa_ws)
                            .arg(&p.qkv16).arg(&p.vt16).arg(&mut p.attn16).arg(&ti).arg(&tpi)
                            .arg(&nhi).arg(&nkvi).arg(&(1.0f32 / (hd as f32).sqrt()))
                            .arg(&mut self.fw_ticket)
                            .launch(LaunchConfig { grid_dim: (self.sm_count as u32, 1, 1),
                                                   block_dim: (384, 1, 1),
                                                   shared_mem_bytes: FW_SHARED })?;
                    } else {
                    let (fk, fsh, gy, bx) = if hop {
                                                (&self.f_fa_hop, FH_SHARED, nh / 2, 256) }
                                            else { (&self.f_fa, FA_SHARED, nh, 128) };
                    self.stream.launch_builder(fk)
                        .arg(&p.qkv16).arg(&p.vt16).arg(&mut p.attn16).arg(&ti).arg(&tpi)
                        .arg(&nhi).arg(&nkvi).arg(&(1.0f32 / (hd as f32).sqrt()))
                        .launch(LaunchConfig { grid_dim: (t.div_ceil(64) as u32, gy as u32, 1),
                                               block_dim: (bx, 1, 1),
                                               shared_mem_bytes: fsh })?;
                    }
                }
                mark!(3);
                if li == 0 && std::env::var("ATTN_DUMP").is_ok() {
                    self.stream.synchronize()?;
                    let rows: Vec<usize> = [0usize, 1, 5, 31, 32, t - 1].into_iter()
                        .filter(|&r| r < t).collect();
                    self.dump_attn("fa", &p.attn16, &rows, nh * lhd)?;
                }
                if !self.fp8_o.is_empty() {
                    let nx = (t * nh * lhd) as i64;
                    let g = 64i32;
                    unsafe {
                        self.stream.launch_builder(&self.f_amax_part)
                            .arg(&p.attn16).arg(&mut self.fp8_part).arg(&nx)
                            .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                        self.stream.launch_builder(&self.f_amax_fin)
                            .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                            .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        self.stream.launch_builder(&self.f_tofp8)
                            .arg(&p.attn16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                            .launch(ecfg((t * nh * lhd).div_ceil(2)))?;
                    }
                    let (lt2, ltws2) = (self.lt, &mut self.ltws);
                    Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_o[li], ltws2,
                        1 << 20, &self.fp8_o[li].0, &self.fp8_px, &mut p.h, 1.0)?;
                } else {
                    let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                    Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans,
                        (t, hidden, self.layers[li].o.nblk * 32),
                        self.layers[li].f16_o.as_ref().unwrap(), &p.attn16, &mut p.h,
                        t, hidden, nh * lhd)?;
                }
                mark!(6);
            } else {
                Self::attn_gemms(&self.stream, &self.blas, p, t, nh, lnkv, lhd, per)?;
                mark!(3);
                let wini = lwin as i32;
                unsafe {
                    let tsegi = tseg as i32;
                    self.stream.launch_builder(&self.f_softmax_c)
                        .arg(&mut p.score16).arg(&mut p.rowsum).arg(&ti).arg(&ascale).arg(&wini).arg(&tsegi)
                        .launch(LaunchConfig { grid_dim: (t as u32, nh as u32, 1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                }
                mark!(4);
                Self::attn_pv(&self.stream, &self.blas, p, t, nh, lnkv, lhd, per)?;
                unsafe {
                    self.stream.launch_builder(&self.f_scale_out)
                        .arg(&mut p.attn).arg(&p.rowsum).arg(&ti).arg(&nhi).arg(&hdi)
                        .launch(ecfg(t * nh * lhd))?;
                }
                mark!(5);
                if !self.fp8_o.is_empty() {
                    // score-path o in fp8: p.attn is f32, so the f32 absmax/convert flavours
                    let nx = (t * nh * lhd) as i64;
                    let g = 64i32;
                    unsafe {
                        self.stream.launch_builder(&self.f_amax_f_part)
                            .arg(&p.attn).arg(&mut self.fp8_part).arg(&nx)
                            .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                        self.stream.launch_builder(&self.f_amax_fin)
                            .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                            .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                        self.stream.launch_builder(&self.f_tofp8_f)
                            .arg(&p.attn).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                            .launch(ecfg((t * nh * lhd).div_ceil(2)))?;
                    }
                    let (lt2, ltws2) = (self.lt, &mut self.ltws);
                    Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_o[li], ltws2,
                        1 << 20, &self.fp8_o[li].0, &self.fp8_px, &mut p.h, 1.0)?;
                } else {
                    // the score-matrix path produces f32, so this arm keeps the converting GEMM
                    Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
                        self.layers[li].f16_o.as_ref().unwrap(), &mut self.xf16,
                        &p.attn, &mut p.h, t, hidden, nh * lhd)?;
                }
                // (FA arm's o-GEMM is routed through the tuned plan below)
                mark!(6);
            }
            unsafe {
                if c.sandwich {
                    let lone = 1.0f32;
                    self.stream.launch_builder(&self.f_sandwich)
                        .arg(&mut p.x).arg(&p.h)
                        .arg(self.layers[li].post_attn_norm.as_ref().unwrap())
                        .arg(&self.layers[li].ffn_norm)
                        .arg(&mut p.hb16).arg(&hi).arg(&eps).arg(&lone).launch(rowsg)?;
                } else if c.num_experts > 0 {
                    // MoE: expert GEMVs consume f32 rows + per-row absmax; p.attn is dead here
                    // (the o-GEMM consumed attn16), so it doubles as the normed-row buffer.
                    self.stream.launch_builder(&self.f_rms_add_b)
                        .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].ffn_norm)
                        .arg(&mut p.attn).arg(&mut self.xs).arg(&hi).arg(&eps)
                        .launch(LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                } else {
                    self.stream.launch_builder(&self.f_rms_bf)
                        .arg(&mut p.x).arg(&p.h).arg(&self.layers[li].ffn_norm)
                        .arg(&mut p.hb16).arg(&hi).arg(&eps).launch(rowsg)?;
                }
            }
            if c.num_experts > 0 {
                // BATCHED MoE FFN over t rows — the same per-pair pipeline as the decode step.
                let (e_n, k, mi) = (c.num_experts, c.topk, c.moe_inter);
                let pairs = t * k;
                // Tile height follows the routing load. Pairs land ~evenly across the experts a
                // batch activates, so pairs is a direct proxy for rows-per-expert, and the sort
                // must emit work items of the same height the kernel consumes.
                // PREFILL, but only when the chunk is big enough to fill a tall tile. At T=2000
                // there are ~125 pairs per expert and the 64-row tile's 4x arithmetic intensity
                // lands: +17% on 8x2k, +7% on long-context decode. A short prompt is the decode
                // regime in disguise -- 240 pairs spread over the experts leaves the tall tile
                // mostly empty, and taking it unconditionally put 60ms on short-prompt TTFT and
                // cost 6-13% on every M cell. tsh must match: the sort emits tiles of 16*MT.
                let tsh: i32 = if pairs >= 1024 { 6 } else { 4 };
                let nih = (t * hidden) as i32;
                unsafe {
                    self.stream.launch_builder(&self.f_tobf16)
                        .arg(&p.attn).arg(&mut p.hb16).arg(&nih).launch(ecfg(t * hidden))?;
                }
                Self::gemm_pre(&self.stream, &self.blas, &self.router16[li], &p.hb16,
                    &mut self.mrlog, t, e_n, hidden)?;

                let (ei, ki) = (e_n as i32, k as i32);
                let normi = if c.norm_topk { 1i32 } else { 0i32 };
                let (nb1, mi2, mii, hsi) = ((hidden / 32) as i32, (2 * mi) as i32, mi as i32, hidden as i32);
                let (xs1, xdivk, one_d, sh0) = (hidden as i32, k as i32, 1i32, 0i32);
                unsafe {
                    self.stream.launch_builder(&self.f_topk_b)
                        .arg(&self.mrlog).arg(&mut self.meidx).arg(&mut self.mewt)
                        .arg(&ei).arg(&ki).arg(&normi)
                        .launch(LaunchConfig { grid_dim: (t as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    let grouped = self.moet || std::env::var("MOEG").as_deref() != Ok("0");
                    let pi32 = pairs as i32;
                    if grouped {
                        // ONE launch for the whole routing sort (was five: memset, count,
                        // scan, scatter, worklist). At 48 layers that is 240 launches a step
                        // saved on a path DECODE_PROF showed to be launch-bound.
                        let cap = self.mwork.len() as i32;
                        let ssh = (3 * e_n * 4) as u32;
                        self.stream.launch_builder(&self.f_moe_sort)
                            .arg(&self.meidx).arg(&mut self.mcnt).arg(&mut self.moff)
                            .arg(&mut self.mplist).arg(&mut self.mwork).arg(&mut self.mwn)
                            .arg(&pi32).arg(&ei).arg(&cap).arg(&tsh)
                            .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1024, 1, 1),
                                                   shared_mem_bytes: ssh })?;
                    }
                    // Pick the tile width from the work available. See gemm_q4_moe_t64.
                    let (f_t, bn) = if pairs >= 1024 { (&self.f_moe_th4, 128usize) }
                                    else if pairs <= 128 { (&self.f_moe_t64, 64usize) }
                                    else { (&self.f_moe_t, 128usize) };
                    let w1 = self.layers[li].w1.as_ref().context("MoE w1")?;
                    if self.moet && self.moet_pf {
                        // PREFILL runs the grouped tensor kernel too. It did not, and that was
                        // the whole MoE long-context deficit: at T = 2000 the per-pair GEMV
                        // walks 16 000 (row, expert) pairs, which measured 1 365 ms of TTFT
                        // against vLLM's 62.
                        self.stream.launch_builder(f_t)
                            .arg(&w1.scales).arg(w1.quants_p.as_ref().unwrap())
                            .arg(&p.attn).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&hsi)
                            .arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&xs1).arg(&xdivk).arg(&self.mwork).arg(&self.mwn)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(bn) as u32, 132, 1),
                                                   block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    } else if grouped {
                        self.stream.launch_builder(&self.f_gemv_moe_g)
                            .arg(&w1.scales).arg(&w1.quants)
                            .arg(&p.attn).arg(&self.xs).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&xs1).arg(&sh0).arg(&xdivk)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(32) as u32, e_n as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    } else {
                        self.stream.launch_builder(&self.f_gemv_moe)
                            .arg(&w1.scales).arg(&w1.quants)
                            .arg(&p.attn).arg(&self.xs).arg(&mut self.meg)
                            .arg(&nb1).arg(&mi2).arg(&self.meidx).arg(&xs1).arg(&sh0).arg(&xdivk)
                            .launch(LaunchConfig { grid_dim: ((2 * mi).div_ceil(32) as u32, pairs as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    }
                    self.stream.launch_builder(&self.f_silu_moe)
                        .arg(&self.meg).arg(&mut self.meu).arg(&mut self.mexs).arg(&mii)
                        .launch(LaunchConfig { grid_dim: (pairs as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    let w2 = self.layers[li].w2.as_ref().context("MoE w2")?;
                    let nb2 = (mi / 32) as i32;
                    if self.moet {
                        self.stream.launch_builder(f_t)
                            .arg(&w2.scales).arg(w2.quants_p.as_ref().unwrap())
                            .arg(&self.meu).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&mii)
                            .arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&mii).arg(&one_d).arg(&self.mwork).arg(&self.mwn)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(bn) as u32, 132, 1),
                                                   block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
                    } else if grouped {
                        self.stream.launch_builder(&self.f_gemv_moe_g)
                            .arg(&w2.scales).arg(&w2.quants)
                            .arg(&self.meu).arg(&self.mexs).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&self.mplist).arg(&self.moff).arg(&self.mcnt)
                            .arg(&mii).arg(&sh0).arg(&one_d)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(32) as u32, e_n as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    } else {
                        self.stream.launch_builder(&self.f_gemv_moe)
                            .arg(&w2.scales).arg(&w2.quants)
                            .arg(&self.meu).arg(&self.mexs).arg(&mut self.medn)
                            .arg(&nb2).arg(&hsi).arg(&self.meidx).arg(&mii).arg(&sh0).arg(&one_d)
                            .launch(LaunchConfig { grid_dim: (hidden.div_ceil(32) as u32, pairs as u32, 1), block_dim: (1024, 1, 1), shared_mem_bytes: 0 })?;
                    }
                    let rowsi = t as i32;
                    self.stream.launch_builder(&self.f_combine_b)
                        .arg(&self.medn).arg(&self.mewt).arg(&mut p.h).arg(&hsi).arg(&ki).arg(&rowsi)
                        .launch(ecfg(t * hidden))?;
                }
                if let Ok(rs) = std::env::var("MOE_ROWDUMP") {
                    // per-layer row fingerprints; ONLY under GRAPH_PF=0 (syncs break capture)
                    if let Ok(r) = rs.parse::<usize>() {
                        if r < t {
                            self.stream.synchronize()?;
                            let sum = |v: &[f32]| -> (f32, f32) {
                                let s: f32 = v.iter().sum();
                                let a = v.iter().fold(0f32, |m, x| m.max(x.abs()));
                                (s, a)
                            };
                            let xin: Vec<f32> = self.stream.memcpy_dtov(&p.attn.slice(r * hidden..(r + 1) * hidden))?;
                            let xs2: Vec<f32> = self.stream.memcpy_dtov(&self.xs.slice(2 * r..2 * r + 2))?;
                            let lgr: Vec<f32> = self.stream.memcpy_dtov(&self.mrlog.slice(r * e_n..(r + 1) * e_n))?;
                            let eid: Vec<u32> = self.stream.memcpy_dtov(&self.meidx.slice(r * k..(r + 1) * k))?;
                            let ewt2: Vec<f32> = self.stream.memcpy_dtov(&self.mewt.slice(r * k..(r + 1) * k))?;
                            let out: Vec<f32> = self.stream.memcpy_dtov(&p.h.slice(r * hidden..(r + 1) * hidden))?;
                            let (xi, xa) = sum(&xin);
                            let (lg2, la) = sum(&lgr);
                            let (os, oa) = sum(&out);
                            eprintln!("[rowdump l{li} r{r}] in {xi:.4}/{xa:.4} xs {:.4},{:.6} rlog {lg2:.4}/{la:.4} eid {eid:?} ewt {:?} out {os:.4}/{oa:.4}",
                                xs2[0], xs2[1],
                                ewt2.iter().map(|w| (w * 1e4).round() / 1e4).collect::<Vec<_>>());
                        }
                    }
                }
                mark!(7); mark!(8); mark!(9);
            } else if !self.fp8_gu.is_empty() {
                // FP8 PREFILL. At T=2000 the GEMM is compute-bound, so the 1.67x FP8 math rate
                // applies directly — and the per-GEMM absmax amortises over a huge matmul instead of
                // a tiny one, which is the opposite of the decode case where it did not pay.
                // (plans for all four prefill GEMMs are rebuilt together in fp8_prefill_plans)
                let nx = (t * hidden) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.hb16).arg(&mut self.fp8_part).arg(&nx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.hb16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&nx)
                        .launch(ecfg((t * hidden).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned(&self.stream, lt2, &self.fp8_pplans[li], ltws2, 1 << 20,
                    &self.fp8_gu[li].0, &self.fp8_px, &mut p.gu16, 1.0)?;
            } else {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_c16_t(&self.stream, &self.blas, lt, ltws, plans,
                    (t, 2 * inter, hidden),
                    self.layers[li].f16_gu.as_ref().unwrap(), &p.hb16, &mut p.gu16,
                    t, 2 * inter, hidden)?;
            }
            mark!(7);
            if c.num_experts == 0 {
            let (ii, n8, act) = (inter as i32, (t * inter / 8) as i32, c.gelu as i32);
            unsafe {
                self.stream.launch_builder(&self.f_silu_bf)
                    .arg(&p.gu16).arg(&mut p.g16).arg(&ii).arg(&n8).arg(&act).launch(ecfg(t * inter / 8))?;
            }
            mark!(8);
            if !self.fp8_down.is_empty() {
                let ndx = (t * inter) as i64;
                let g = 64i32;
                unsafe {
                    self.stream.launch_builder(&self.f_amax_part)
                        .arg(&p.g16).arg(&mut self.fp8_part).arg(&ndx)
                        .launch(LaunchConfig { grid_dim: (64,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_amax_fin)
                        .arg(&self.fp8_part).arg(&mut self.fp8_psmax[li]).arg(&mut self.fp8_psx[li]).arg(&g)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (64,1,1), shared_mem_bytes: 0 })?;
                    self.stream.launch_builder(&self.f_tofp8)
                        .arg(&p.g16).arg(&mut self.fp8_px).arg(&self.fp8_psmax[li]).arg(&ndx)
                        .launch(ecfg((t * inter).div_ceil(2)))?;
                }
                let (lt2, ltws2) = (self.lt, &mut self.ltws);
                Self::fp8_matmul_planned_f32(&self.stream, lt2, &self.fp8_pplans_down[li], ltws2,
                    1 << 20, &self.fp8_down[li].0, &self.fp8_px, &mut p.h, 1.0)?;
            } else {
                let (lt, ltws, plans) = (self.lt, &mut self.ltws, &self.bf16_plans);
                Self::gemm_pre_t(&self.stream, &self.blas, lt, ltws, plans, (t, hidden, inter),
                    self.layers[li].f16_down.as_ref().unwrap(), &p.g16, &mut p.h,
                    t, hidden, inter)?;
            }
            mark!(9);
            }
        }
        if prof {
            let names = ["rmsnorm", "qkv-gemm", "rope+cvt", "attention", "softmax",
                         "PV-gemm+scale", "o-gemm", "gate|up-gemm", "silu", "down-gemm"];
            let tot: f64 = tp.iter().sum();
            println!("== PREFILL PROFILE (T={t}, sum {:.2} ms)", tot * 1e3);
            for (n, v) in names.iter().zip(tp.iter()) {
                println!("   {n:<14} {:>7.2} ms  {:>5.1}%", v * 1e3, v / tot * 100.0);
            }
        }
        if tseg != t {
            // WAVE: the caller gathers each segment's last row and runs the [B, vocab] LM head.
            return Ok(());
        }
        if c.sandwich {
            // Only the LAST row feeds the first sampled token, and prefill's x is not read
            // again after this, so the final sandwich (post-FFN norm + residual + layer_scalar
            // + final norm) runs on that single row. Softcap skipped: 30·tanh(z/30) is
            // monotonic, greedy argmax unaffected.
            let last = c.layers - 1;
            let pls = self.layers[last].ls;
            let mut xrow = p.x.slice_mut((t - 1) * hidden..t * hidden);
            let hrow = p.h.slice((t - 1) * hidden..t * hidden);
            let onerow = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                self.stream.launch_builder(&self.f_sandwich)
                    .arg(&mut xrow).arg(&hrow)
                    .arg(self.layers[last].post_ffw_norm.as_ref().unwrap())
                    .arg(&self.final_norm)
                    .arg(&mut self.xf16).arg(&hi).arg(&eps).arg(&pls).launch(onerow)?;
            }
            Self::gemm_pre(&self.stream, &self.blas, self.f16_lm.as_ref().unwrap(),
                &self.xf16, &mut p.logits, 1, c.vocab, hidden)?;
        } else {
        unsafe {
            self.stream.launch_builder(&self.f_rms_add_b)
                .arg(&mut p.x).arg(&p.h).arg(&self.final_norm)
                .arg(&mut p.attn).arg(&mut self.xs).arg(&hi).arg(&eps).launch(rowsg)?;
        }
        // Only the LAST row feeds the first sampled token, so copy it into the single-token buffer
        // rather than running the 151936-wide LM head over all T rows.
        self.stream.memcpy_dtod(&p.attn.slice((t - 1) * hidden..t * hidden), &mut self.h)?;
        Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16,
            self.f16_lm.as_ref().unwrap(), &mut self.xf16, &self.h, &mut p.logits, 1, c.vocab, hidden)?;
        if let Ok(rs) = std::env::var("MOE_ROWDUMP") {
            if let Ok(r) = rs.parse::<usize>() {
                if r < t {
                    self.stream.synchronize()?;
                    let pn: Vec<f32> = self.stream.memcpy_dtov(&p.attn.slice(r * hidden..(r + 1) * hidden))?;
                    let hrow: Vec<f32> = self.stream.memcpy_dtov(&self.h.slice(0..hidden))?;
                    let lgv: Vec<f32> = self.stream.memcpy_dtov(&p.logits.slice(0..c.vocab))?;
                    let fp = |v: &[f32]| (v.iter().sum::<f32>(), v.iter().fold(0f32, |m, x| m.max(x.abs())));
                    let (ps, pa) = fp(&pn);
                    let (hs, ha) = fp(&hrow);
                    let mut idx: Vec<usize> = (0..c.vocab).collect();
                    idx.sort_by(|&a, &b| lgv[b].partial_cmp(&lgv[a]).unwrap());
                    eprintln!("[taildump t{t} r{r}] postnorm {ps:.4}/{pa:.4} hrow {hs:.4}/{ha:.4} top3 {:?}",
                        idx[..3].iter().map(|&i| (i, (lgv[i] * 1e3).round() / 1e3)).collect::<Vec<_>>());
                }
            }
        }
        }
        let v = c.vocab as i32;
        unsafe {
            self.stream.launch_builder(&self.f_argmax)
                .arg(&p.logits).arg(&mut self.d_token).arg(&v)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1024,1,1), shared_mem_bytes: 0 })?;
        }
        // NOT a host memcpy: this line is captured into the per-T prefill graph, and a captured
        // pageable H2D records the (stack) source pointer — replays then read freed memory. The
        // batch path never consumes d_pos so the corruption stayed invisible until the sequential
        // server path read it. A kernel-parameter write replays correctly by construction.
        let ti32 = t as i32;
        unsafe {
            self.stream.launch_builder(&self.f_set_i32)
                .arg(&mut self.d_pos).arg(&ti32)
                .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })?;
        }
        Ok(())
    }

    /// S[h] = Q[h] * K[h]^T for every head, one strided-batched GEMM.
    #[allow(clippy::too_many_arguments)]
    fn attn_gemms(stream: &Arc<CudaStream>, blas: &CudaBlas, p: &mut Prefill,
                  t: usize, nh: usize, nkv: usize, hd: usize, per: usize) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pq, _a) = p.qkv16.device_ptr(stream);
        let (ps, _c) = p.score16.device_ptr_mut(stream);
        let koff = (nh * hd) as u64;
        // GQA: query heads sharing a kv head form contiguous groups, so one strided-batched call
        // per group covers grp heads (a per-head loop cost ~900 launches per prefill).
        let grp = nh / nkv;
        for kvh in 0..nkv {
            unsafe {
                blas::gemm_strided_batched_ex(
                    *blas.handle(),
                    blas_sys::cublasOperation_t::CUBLAS_OP_T,
                    blas_sys::cublasOperation_t::CUBLAS_OP_N,
                    t as i32, t as i32, hd as i32,
                    (&alpha) as *const f32 as *const _,
                    (pq + (koff + (kvh * hd) as u64) * 2) as *const std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_16F, (per * hd) as i32, 0i64,
                    (pq + ((kvh * grp * hd) as u64) * 2) as *const std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_16F, (per * hd) as i32, hd as i64,
                    (&beta) as *const f32 as *const _,
                    (ps + ((kvh * grp * t * t) as u64) * 2) as *mut std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_16F, t as i32, (t * t) as i64,
                    grp as i32,
                    blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                    blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
                ).map_err(|e| anyhow::anyhow!("qk gemm: {e:?}"))?;
            }
        }
        Ok(())
    }

    /// O[h] = P[h] * V[h]; column-major identity O_cm = V_cm * P_cm, so OP_N/OP_N.
    #[allow(clippy::too_many_arguments)]
    fn attn_pv(stream: &Arc<CudaStream>, blas: &CudaBlas, p: &mut Prefill,
               t: usize, nh: usize, nkv: usize, hd: usize, per: usize) -> Result<()> {
        let (alpha, beta) = (1.0f32, 0.0f32);
        let (pq, _a) = p.qkv16.device_ptr(stream);
        let (ps, _b) = p.score16.device_ptr(stream);
        let (po, _c) = p.attn.device_ptr_mut(stream);
        let voff = ((nh + nkv) * hd) as u64;
        let grp = nh / nkv;
        for kvh in 0..nkv {
            unsafe {
                blas::gemm_strided_batched_ex(
                    *blas.handle(),
                    blas_sys::cublasOperation_t::CUBLAS_OP_N,
                    blas_sys::cublasOperation_t::CUBLAS_OP_N,
                    hd as i32, t as i32, t as i32,
                    (&alpha) as *const f32 as *const _,
                    (pq + (voff + (kvh * hd) as u64) * 2) as *const std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_16F, (per * hd) as i32, 0i64,
                    (ps + ((kvh * grp * t * t) as u64) * 2) as *const std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_16F, t as i32, (t * t) as i64,
                    (&beta) as *const f32 as *const _,
                    (po + ((kvh * grp * hd) as u64) * 4) as *mut std::ffi::c_void,
                    blas_sys::cudaDataType_t::CUDA_R_32F, (nh * hd) as i32, hd as i64,
                    grp as i32,
                    blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                    blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
                ).map_err(|e| anyhow::anyhow!("pv gemm: {e:?}"))?;
            }
        }
        Ok(())
    }

    /// Unit test of the tensor-core arm: run layer 0's qkv projection through BOTH the q4_0 GEMV
    /// and the f16-mirror cuBLAS GEMM on the same activation, and report the largest disagreement.
    /// They dequantise the same weights, so they must agree to f16 rounding; a large delta means the
    /// mirror, the GEMM layout, or the per-row scaling is wrong.
    pub fn verify_gemm(&mut self) -> Result<(f32, f32)> {
        anyhow::ensure!(self.layers[0].f16_qkv.is_some(), "call init_batch first");
        let hidden = self.cfg.hidden;
        let rows = self.layers[0].qkv.rows;
        // deterministic activation: normalise the embedding of token 100
        let (tok, hi) = (100i32, hidden as i32);
        let dt = self.stream.memcpy_stod(&[100u32])?;
        let esc = self.cfg.embed_scale;
        unsafe {
            self.stream.launch_builder(&self.f_embed)
                .arg(&self.embed).arg(&mut self.x).arg(&mut self.h).arg(&dt).arg(&hi).arg(&esc)
                .launch(ecfg(hidden))?;
        }
        let _ = tok;
        let eps = self.cfg.eps;
        Self::rms_add(&self.stream, &self.f_rms_add, &mut self.x, &self.h,
                      &self.layers[0].attn_norm, &mut self.abuf, &mut self.xs, hidden, eps)?;
        std::mem::swap(&mut self.h, &mut self.abuf);

        let mut y_gemv = self.stream.alloc_zeros::<f32>(rows)?;
        Self::gemv(&self.stream, &self.f_gemv, &self.layers[0].qkv, &self.h, &self.xs,
                   &mut y_gemv, false)?;
        let mut y_gemm = self.stream.alloc_zeros::<f32>(rows)?;
        Self::gemm_bf16(&self.stream, &self.blas, &self.f_tobf16, self.layers[0].f16_qkv.as_ref().unwrap(),
                       &mut self.xf16, &self.h, &mut y_gemm, 1, rows, hidden)?;
        self.sync()?;
        let a = self.stream.memcpy_dtov(&y_gemv)?;
        let b = self.stream.memcpy_dtov(&y_gemm)?;
        let mut maxabs = 0f32;
        let mut maxdiff = 0f32;
        for (x, y) in a.iter().zip(b.iter()) {
            maxabs = maxabs.max(x.abs());
            maxdiff = maxdiff.max((x - y).abs());
        }
        eprintln!("  gemv[0..8] = {:?}", &a[..8.min(a.len())]);
        eprintln!("  gemm[0..8] = {:?}", &b[..8.min(b.len())]);
        // a constant ratio => scaling bug; shuffled values => layout/permutation bug
        let ratios: Vec<f32> = a.iter().zip(b.iter()).take(6)
            .map(|(x, y)| if y.abs() > 1e-6 { x / y } else { f32::NAN }).collect();
        eprintln!("  gemv/gemm ratios = {ratios:?}");
        Ok((maxdiff, maxabs))
    }

    /// Read back the q4 bytes for one block and dequantise on the HOST, then compare against what
    /// the device mirror holds. Isolates `dequant_q4_f16` from the GEMM entirely.
    pub fn verify_dequant(&mut self) -> Result<()> {
        anyhow::ensure!(self.layers[0].f16_qkv.is_some(), "call init_batch first");
        let w = &self.layers[0].qkv;
        let sc: Vec<f16> = self.stream.memcpy_dtov(&w.scales.slice(0..2))?;
        let qz: Vec<u32> = self.stream.memcpy_dtov(&w.quants.slice(0..8))?;
        let mir: Vec<u16> = self
            .stream
            .memcpy_dtov(&self.layers[0].f16_qkv.as_ref().unwrap().slice(0..64))?;
        // host dequant of block 0, exactly as weights::quantize_q4_0 packs it
        let d = sc[0].to_f32();
        let mut want = [0f32; 32];
        for wi in 0..4usize {
            let word = qz[wi];
            for byte in 0..4usize {
                let bi = wi * 4 + byte;
                let v = (word >> (byte * 8)) & 0xFF;
                want[bi] = d * (((v & 0xF) as i32 - 8) as f32);
                want[bi + 16] = d * ((((v >> 4) & 0xF) as i32 - 8) as f32);
            }
        }
        // bf16 bits -> f32: bf16 is just the top 16 bits of an f32
        let got: Vec<f32> = mir[..32]
            .iter()
            .map(|&b| f32::from_bits((b as u32) << 16))
            .collect();
        let mut worst = 0f32;
        for i in 0..32 {
            worst = worst.max((want[i] - got[i]).abs());
        }
        eprintln!("  scale d = {d:e}");
        eprintln!("  want[0..8] = {:?}", &want[..8]);
        eprintln!("  got [0..8] = {:?}", &got[..8]);
        eprintln!("  block-0 worst |diff| = {worst:e}  => {}",
                  if worst < d.abs() * 0.1 { "dequant OK" } else { "DEQUANT WRONG" });
        Ok(())
    }

    pub fn sync(&self) -> Result<()> {
        self.stream.synchronize()?;
        Ok(())
    }

    /// (min, max, mean, nan_count) of a device f32 buffer — the correctness probe.
    fn stats(&self, b: &CudaSlice<f32>, n: usize) -> Result<(f32, f32, f32, usize)> {
        let v = self.stream.memcpy_dtov(&b.slice(0..n))?;
        let nan = v.iter().filter(|x| !x.is_finite()).count();
        let fin: Vec<f32> = v.iter().copied().filter(|x| x.is_finite()).collect();
        if fin.is_empty() {
            return Ok((f32::NAN, f32::NAN, f32::NAN, nan));
        }
        let mn = fin.iter().copied().fold(f32::INFINITY, f32::min);
        let mx = fin.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        Ok((mn, mx, fin.iter().sum::<f32>() / fin.len() as f32, nan))
    }

    pub fn cfg(&self) -> &Cfg {
        &self.cfg
    }
}

const KERNELS: &str = r#"
// ---- half helpers via inline PTX (NVRTC ships no cuda_fp16.h; requiring CUDA headers on a
// serving host would defeat a dlopen-only fast path) ----
__device__ __forceinline__ float h2f(unsigned short h) {
    float f; asm("{ .reg .f16 t; mov.b16 t, %1; cvt.f32.f16 %0, t; }" : "=f"(f) : "h"(h)); return f;
}

__device__ __forceinline__ float bf2f(unsigned short h) {
    return __uint_as_float(((unsigned int)h) << 16);
}

// act = 0: silu (Qwen/Llama/Mistral). act = 1: gelu_pytorch_tanh (Gemma4 GeGLU).
__device__ __forceinline__ float ffn_act(const float v, const int act) {
    if (act) {
        const float c = 0.7978845608028654f * (v + 0.044715f * v * v * v);
        return 0.5f * v * (1.f + tanhf(c));
    }
    return v / (1.f + __expf(-v));
}

__device__ __forceinline__ unsigned short f2h(float v) {
    unsigned short r;
    asm("{ .reg .f16 t; cvt.rn.f16.f32 t, %1; mov.b16 %0, t; }" : "=h"(r) : "f"(v));
    return r;
}
__device__ __forceinline__ unsigned int pack_h2(float lo, float hi) {
    unsigned int d;
    asm("{ .reg .f16 l, h; cvt.rn.f16.f32 l, %1; cvt.rn.f16.f32 h, %2; mov.b32 %0, {l,h}; }"
        : "=r"(d) : "f"(lo), "f"(hi));
    return d;
}
__device__ __forceinline__ unsigned int hadd2(unsigned int a, unsigned int b) {
    unsigned int d; asm("add.f16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); return d;
}
__device__ __forceinline__ unsigned int hfma2(unsigned int a, unsigned int b, unsigned int c) {
    unsigned int d; asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(c)); return d;
}
// f32 -> bf16 (round-to-nearest-even) by bit manipulation. bf16 keeps f32's exponent range, which
// is why the GEMM arm needs NO activation scaling: silu reaching ~1.5e4 is unremarkable in bf16,
// whereas f16 required per-row absmax scaling plus a rescale (3 extra kernels per GEMM, measured to
// cost ~15% of concurrency throughput). Done with integer ops so this still compiles at compute_70
// instead of needing sm_80's cvt.rn.bf16.f32.
__device__ __forceinline__ unsigned short f2bf(float v) {
    unsigned int u = __float_as_uint(v);
    return (unsigned short)((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16);
}

__device__ __forceinline__ void unpack_h2(unsigned int a, float* lo, float* hi) {
    asm("{ .reg .f16 l, h; mov.b32 {l,h}, %2; cvt.f32.f16 %0, l; cvt.f32.f16 %1, h; }"
        : "=f"(*lo), "=f"(*hi) : "r"(a));
}

// Activation scaling. f16 products overflow on real residual streams: measured at layer 2 of
// Qwen3-1.7B, silu(gate)*up reaches -1.5e4, so ONE product against a q4 weight (|q-8| <= 8) is
// ~1.2e5 — past f16's 65504 ceiling before any accumulation, which is what produced the NaNs.
// So normalise: stage x * (1/max|x|) into shared (values in [-1,1], products <= 8), accumulate,
// and multiply max|x| back at the end. Exactly equivalent, and the f16 path stays safe.
// MEASURED: 13.07 us x ~28 per M=1 step. It runs on ONE block, so the whole GPU is one SM's
// worth of threads; the only lever is threads-per-block and barrier count. Warp-shuffle max
// (one barrier instead of log2(blockDim)) and red[] sized by warp count, not by blockDim — so
// this is correct at 256 or 1024 threads and cannot repeat the red[256] out-of-bounds trap.
extern "C" __global__ void absmax(const float* __restrict__ x, float* __restrict__ out, const int n) {
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float m = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) m = fmaxf(m, fabsf(x[i]));
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_down_sync(0xffffffffu, m, o));
    if (lane == 0) red[wid] = m;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t = fmaxf(t, __shfl_down_sync(0xffffffffu, t, o));
        if (lane == 0) { out[0] = t; out[1] = (t > 0.f) ? (1.f / t) : 0.f; }
    }
}

#define TILEB 64
// q4_0 GEMV, the 1829 GB/s kernel (91% of H100 peak): activations staged in shared and PERMUTED
// to match the packing, weights dequantised with bit ops instead of int->float cvt.
//
// Packing (weights::quantize_q4_0): word wi, byte b holds q[wi*4+b] in the low nibble and
// q[wi*4+b+16] in the high nibble. Masking a word with 0x000F000F therefore yields elements
// (4wi + r/2 + 16*(r&1)) and that +2, for r = 0..3 — which is exactly the pairing staged below.
extern "C" __global__ void gemv_q4(
    const unsigned short* __restrict__ scales,
    const uint4*  __restrict__ quants,
    const float*  __restrict__ x,
    const float*  __restrict__ xsc,   /* [max|x|, 1/max|x|] from absmax */
    float*        __restrict__ y,
    const int nblk,
    const int rows,
    const int accumulate)   /* 1 => y[row] += result (residual add, one launch fewer) */
{
    __shared__ unsigned int xh[TILEB * 17];
    const float inv_s = xsc[1];
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int row  = blockIdx.x * 32 + warp;
    const long base = (long)row * (long)nblk;
    float acc = 0.f;
    const unsigned int bias = 0xE408E408u;   /* half2(-1032) = -(1024 + 8) */

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        for (int i = threadIdx.x; i < TILEB * 16; i += blockDim.x) {
            const int b = i >> 4, j = i & 15;
            const int wi = j >> 2, r = j & 3;
            const int a = 4 * wi + (r >> 1) + 16 * (r & 1);   /* element A of the pair */
            const int gi = ((t0 + b) << 5);
            xh[b * 17 + j] = ((t0 + b) < nblk)
                ? pack_h2(x[gi + a] * inv_s, x[gi + a + 2] * inv_s)
                : 0u;
        }
        __syncthreads();
        if (row < rows) {
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                const uint4 q = quants[base + t0 + b];
                const float s = h2f(scales[base + t0 + b]);
                const unsigned int* xp = xh + b * 17;
                unsigned int w[4] = { q.x, q.y, q.z, q.w };
                // Accumulate per WORD (8 products) in f16x2, then fold into f32. Summing all 32
                // products of a q-block in f16 OVERFLOWS: the residual stream reaches O(100s) by
                // layer 2, products hit ~4e3, and 32 of them exceed f16's 65504 -> inf -> NaN
                // (observed: clean through L1, 3 NaNs at L2, everything NaN at L3). Per-word keeps
                // the peak ~8x under the ceiling. Memory traffic is IDENTICAL — this is ALU only.
                float blk = 0.f;
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const unsigned int v = w[i];
                    unsigned int sw = 0u;
                    #pragma unroll
                    for (int r = 0; r < 4; ++r) {
                        const unsigned int m = ((v >> (4 * r)) & 0x000F000Fu) | 0x64006400u;
                        sw = hfma2(hadd2(m, bias), xp[i * 4 + r], sw);
                    }
                    float lo, hi; unpack_h2(sw, &lo, &hi);
                    blk += lo + hi;
                }
                acc = fmaf(s, blk, acc);
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0 && row < rows) y[row] = acc * xsc[0];
}

// Router GEMV kept in f32: it is tiny ([E, hidden]) and routing accuracy decides which experts
// run, so quantising it would trade real quality for no bandwidth.
extern "C" __global__ void gemv_f32(const float* __restrict__ w, const float* __restrict__ x,
                                    float* __restrict__ y, const int rows, const int cols) {
    const int row = blockIdx.x;
    if (row >= rows) return;
    __shared__ float red[256];
    float a = 0.f;
    for (int i = threadIdx.x; i < cols; i += blockDim.x) a = fmaf(w[(long)row * cols + i], x[i], a);
    red[threadIdx.x] = a;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o];
        __syncthreads();
    }
    if (threadIdx.x == 0) y[row] = red[0];
}

// softmax over experts -> top-k selection -> (optional) renormalisation of the kept weights.
extern "C" __global__ void moe_topk(const float* __restrict__ logits, unsigned int* __restrict__ idx,
                                    float* __restrict__ wts, const int n_exp, const int k,
                                    const int norm) {
    __shared__ float p[512];
    __shared__ float red[256];
    __shared__ int ired[256];
    __shared__ float wsum_s;
    const int t = threadIdx.x;
    float mx = -1e30f;
    for (int i = t; i < n_exp; i += blockDim.x) mx = fmaxf(mx, logits[i]);
    red[t] = mx; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if (t < o) red[t] = fmaxf(red[t], red[t + o]); __syncthreads(); }
    const float m = red[0];
    __syncthreads();
    float sum = 0.f;
    for (int i = t; i < n_exp; i += blockDim.x) { const float e = __expf(logits[i] - m); p[i] = e; sum += e; }
    red[t] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if (t < o) red[t] += red[t + o]; __syncthreads(); }
    const float inv = 1.f / red[0];
    __syncthreads();
    for (int i = t; i < n_exp; i += blockDim.x) p[i] *= inv;
    if (t == 0) wsum_s = 0.f;
    __syncthreads();
    // k selection rounds, each an argmax reduction across the block
    for (int j = 0; j < k; ++j) {
        float bv = -1.f; int bi = 0;
        for (int i = t; i < n_exp; i += blockDim.x) if (p[i] > bv) { bv = p[i]; bi = i; }
        red[t] = bv; ired[t] = bi; __syncthreads();
        for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
            if (t < o && red[t + o] > red[t]) { red[t] = red[t + o]; ired[t] = ired[t + o]; }
            __syncthreads();
        }
        if (t == 0) {
            idx[j] = (unsigned int)ired[0];
            wts[j] = red[0];
            wsum_s += red[0];
            p[ired[0]] = -1.f;   // consume
        }
        __syncthreads();
    }
    if (norm && t == 0 && wsum_s > 0.f) for (int j = 0; j < k; ++j) wts[j] /= wsum_s;
}

// absmax per SLOT: each selected expert's activation has its own magnitude, so each needs its own
// f16 scaling factor. blockIdx.x = slot; writes (max, 1/max) pairs.
extern "C" __global__ void absmax_slots(const float* __restrict__ x, float* __restrict__ out,
                                        const int n, const int stride) {
    __shared__ float red[256];
    const int slot = blockIdx.x;
    const float* xs = x + (long)slot * stride;
    float m = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) m = fmaxf(m, fabsf(xs[i]));
    red[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) { const float mx = red[0]; out[2*slot] = mx; out[2*slot+1] = (mx > 0.f) ? (1.f/mx) : 0.f; }
}

// The MoE GEMV: one launch covers ALL k selected experts (blockIdx.y = slot), instead of k
// separate launches — at 48 layers x 3 matrices that is the difference between 1152 and 144
// launches per token, and this engine is launch-sensitive.
extern "C" __global__ void gemv_q4_moe(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, const float* __restrict__ xsc, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const unsigned int* __restrict__ eids,
    const int x_stride,
    const int shared_scale,  /* 1 => every slot uses xsc[0..2] (rmsnorm already computed it) */
    const int x_div)         /* BATCH: slot p = row*k + j shares row p/x_div's input + scale */
{
    __shared__ unsigned int xh[TILEB * 17];
    const int slot = blockIdx.y;
    const int e = (int)eids[slot];
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int row = blockIdx.x * 32 + warp;
    const long base = ((long)e * rows_per_expert + row) * (long)nblk;
    const int xrow = slot / x_div;
    const float* xin = x + (long)xrow * x_stride;
    const float inv_s = shared_scale ? xsc[1] : xsc[2*xrow + 1];
    const float sc_mul = shared_scale ? xsc[0] : xsc[2*xrow];
    float acc = 0.f;
    const unsigned int bias = 0xE408E408u;

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        for (int i = threadIdx.x; i < TILEB * 16; i += blockDim.x) {
            const int b = i >> 4, j = i & 15;
            const int wi = j >> 2, r = j & 3;
            const int a = 4 * wi + (r >> 1) + 16 * (r & 1);
            const int gi = ((t0 + b) << 5);
            xh[b * 17 + j] = ((t0 + b) < nblk) ? pack_h2(xin[gi + a] * inv_s, xin[gi + a + 2] * inv_s) : 0u;
        }
        __syncthreads();
        if (row < rows_per_expert) {
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                const uint4 q = quants[base + t0 + b];
                const float sc = h2f(scales[base + t0 + b]);
                const unsigned int* xp = xh + b * 17;
                unsigned int w[4] = { q.x, q.y, q.z, q.w };
                float blk = 0.f;
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const unsigned int v = w[i];
                    unsigned int sw = 0u;
                    #pragma unroll
                    for (int r = 0; r < 4; ++r) {
                        const unsigned int m = ((v >> (4 * r)) & 0x000F000Fu) | 0x64006400u;
                        sw = hfma2(hadd2(m, bias), xp[i * 4 + r], sw);
                    }
                    float lo, hi; unpack_h2(sw, &lo, &hi);
                    blk += lo + hi;
                }
                acc = fmaf(sc, blk, acc);
            }
        }
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
    if (lane == 0 && row < rows_per_expert)
        y[(long)slot * rows_per_expert + row] = acc * sc_mul;
}

// EXPERT-MAJOR GROUPED MoE GEMV. The per-pair kernel reads each pair's expert weights
// independently: at M=64/topk8 that is 512 x expert-bytes per layer against a ~125-distinct-
// expert floor (4x), and it is why the first 30B ladder lost every M>=8 cell. Here pairs are
// counting-sorted by expert on device (moe_count/scan/scatter) and each block owns one
// (expert, 32-row weight tile): every weight read serves up to PBG pairs of that expert.
// Per-pair accumulators are independent of their batch-mates, so results are DETERMINISTIC
// regardless of the atomic scatter order.
#define PBG 8
extern "C" __global__ void moe_count(const unsigned int* __restrict__ eids, int* __restrict__ cnt,
                                     const int pairs) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < pairs) atomicAdd(&cnt[eids[i]], 1);
}
// exclusive prefix sum over <=512 counts, one block; also zero-inits the scatter cursors
extern "C" __global__ void moe_scan(const int* __restrict__ cnt, int* __restrict__ off,
                                    int* __restrict__ cur, const int e_n) {
    __shared__ int sh[512];
    const int t = threadIdx.x;
    sh[t] = (t < e_n) ? cnt[t] : 0;
    __syncthreads();
    for (int d = 1; d < 512; d <<= 1) {
        int v = (t >= d) ? sh[t - d] : 0;
        __syncthreads();
        sh[t] += v;
        __syncthreads();
    }
    if (t < e_n) {
        off[t] = sh[t] - cnt[t];   // inclusive -> exclusive
        cur[t] = 0;
    }
}
// Compact (expert, pair-tile) WORK LIST. Only ~8 of 128 experts are active on a shared-prefix
// batch, so a grid of (columns x experts x pair-tiles) dispatches ~49k blocks to do ~384 blocks
// of work — millions of empty dispatches per step. One thread per expert appends its tiles here
// and the GEMM grid becomes (columns x fixed) blocks that pull from this list.
// FUSED ROUTING SORT: memset + count + scan + scatter + worklist in ONE launch.
//
// The MoE step is launch-bound, not math-bound — DECODE_PROF puts 59% of it in the FFN against
// a ~0.7 ms weight roofline, and the routing chain alone was five launches per layer (240 per
// step at 48 layers) for work that touches a few thousand ints. One block does all of it:
// counts into shared, scans in shared, scatters with a shared cursor, then emits the compact
// (expert, pair-tile) work list. Same outputs, same order guarantees, one launch.
extern "C" __global__ void moe_sort_fused(const unsigned int* __restrict__ eids,
                                          int* __restrict__ cnt, int* __restrict__ off,
                                          unsigned int* __restrict__ plist,
                                          unsigned int* __restrict__ work, int* __restrict__ nwork,
                                          const int pairs, const int e_n, const int cap,
                                          const int tsh) {
    extern __shared__ int ssh2[];
    int* scnt = ssh2;              // [e_n]
    int* soff = scnt + e_n;        // [e_n]
    int* scur = soff + e_n;        // [e_n]
    for (int e = threadIdx.x; e < e_n; e += blockDim.x) scnt[e] = 0;
    __syncthreads();
    for (int i = threadIdx.x; i < pairs; i += blockDim.x) atomicAdd(&scnt[(int)eids[i]], 1);
    __syncthreads();
    // exclusive scan, one thread — e_n is 128-256 and this is cheaper than a parallel scan's
    // barriers at that size
    if (threadIdx.x == 0) {
        int acc = 0;
        for (int e = 0; e < e_n; ++e) { soff[e] = acc; scur[e] = acc; acc += scnt[e]; }
        *nwork = 0;
    }
    __syncthreads();
    for (int i = threadIdx.x; i < pairs; i += blockDim.x) {
        const int e = (int)eids[i];
        plist[atomicAdd(&scur[e], 1)] = (unsigned int)i;
    }
    __syncthreads();
    for (int e = threadIdx.x; e < e_n; e += blockDim.x) {
        cnt[e] = scnt[e];
        off[e] = soff[e];
        const int tiles = (scnt[e] + (1 << tsh) - 1) >> tsh;
        for (int t = 0; t < tiles; ++t) {
            const int slot = atomicAdd(nwork, 1);
            if (slot < cap) work[slot] = ((unsigned int)e << 16) | (unsigned int)t;
        }
    }
}

extern "C" __global__ void moe_worklist(const int* __restrict__ cnt, unsigned int* __restrict__ work,
                                        int* __restrict__ nwork, const int e_n, const int cap) {
    const int e = blockIdx.x * blockDim.x + threadIdx.x;
    if (e == 0) *nwork = 0;
    __threadfence();
    __syncthreads();
    if (e >= e_n) return;
    const int c = cnt[e];
    const int tiles = (c + 15) >> 4;
    for (int t = 0; t < tiles; ++t) {
        const int slot = atomicAdd(nwork, 1);
        if (slot < cap) work[slot] = ((unsigned int)e << 16) | (unsigned int)t;
    }
}

extern "C" __global__ void moe_scatter(const unsigned int* __restrict__ eids,
                                       const int* __restrict__ off, int* __restrict__ cur,
                                       unsigned int* __restrict__ plist, const int pairs) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < pairs) {
        const int e = (int)eids[i];
        plist[off[e] + atomicAdd(&cur[e], 1)] = (unsigned int)i;
    }
}

extern "C" __global__ void gemv_q4_moe_g(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, const float* __restrict__ xsc, float* __restrict__ y,
    const int nblk, const int rows_per_expert,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt,
    const int x_stride, const int shared_scale, const int x_div)
{
    const int e = blockIdx.y;
    const int cnt = ecnt[e];
    if (cnt == 0) return;
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int row = blockIdx.x * 32 + warp;
    const long base = ((long)e * rows_per_expert + row) * (long)nblk;
    __shared__ unsigned int xh[PBG][TILEB * 17];
    const unsigned int bias = 0xE408E408u;
    for (int p0 = 0; p0 < cnt; p0 += PBG) {
        const int pb = min(PBG, cnt - p0);
        float acc[PBG];
        #pragma unroll
        for (int i = 0; i < PBG; ++i) acc[i] = 0.f;
        for (int t0 = 0; t0 < nblk; t0 += TILEB) {
            __syncthreads();
            for (int i = threadIdx.x; i < pb * TILEB * 16; i += blockDim.x) {
                const int pi = i / (TILEB * 16), j2 = i - pi * (TILEB * 16);
                const int b = j2 >> 4, j = j2 & 15;
                const int wi = j >> 2, r = j & 3;
                const int a = 4 * wi + (r >> 1) + 16 * (r & 1);
                const int gi = ((t0 + b) << 5);
                const unsigned int pr = plist[eoff[e] + p0 + pi];
                const int xrow = (int)pr / x_div;
                const float inv_s = shared_scale ? xsc[1] : xsc[2 * xrow + 1];
                const float* xin = x + (long)xrow * x_stride;
                xh[pi][b * 17 + j] = ((t0 + b) < nblk)
                    ? pack_h2(xin[gi + a] * inv_s, xin[gi + a + 2] * inv_s) : 0u;
            }
            __syncthreads();
            if (row < rows_per_expert) {
                for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                    const uint4 q = quants[base + t0 + b];
                    const float sc = h2f(scales[base + t0 + b]);
                    unsigned int w[4] = { q.x, q.y, q.z, q.w };
                    #pragma unroll
                    for (int pi = 0; pi < PBG; ++pi) {
                        if (pi >= pb) break;
                        const unsigned int* xp = xh[pi] + b * 17;
                        unsigned int sw = 0u;
                        float blk = 0.f;
                        #pragma unroll
                        for (int i = 0; i < 4; ++i) {
                            const unsigned int v = w[i];
                            unsigned int s2 = 0u;
                            #pragma unroll
                            for (int r = 0; r < 4; ++r) {
                                const unsigned int mm = ((v >> (4 * r)) & 0x000F000Fu) | 0x64006400u;
                                s2 = hfma2(hadd2(mm, bias), xp[i * 4 + r], s2);
                            }
                            float lo, hi; unpack_h2(s2, &lo, &hi);
                            blk += lo + hi;
                        }
                        (void)sw;
                        acc[pi] = fmaf(sc, blk, acc[pi]);
                    }
                }
            }
        }
        #pragma unroll
        for (int pi = 0; pi < PBG; ++pi) {
            if (pi >= pb) break;
            float a = acc[pi];
            #pragma unroll
            for (int off2 = 16; off2 > 0; off2 >>= 1) a += __shfl_down_sync(0xffffffffu, a, off2);
            if (lane == 0 && row < rows_per_expert) {
                const unsigned int pr = plist[eoff[e] + p0 + pi];
                const int xrow = (int)pr / x_div;
                const float sc_mul = shared_scale ? xsc[0] : xsc[2 * xrow];
                y[(long)pr * rows_per_expert + row] = a * sc_mul;
            }
        }
    }
}

// out = sum_slot wts[slot] * dn[slot]
// silu(gate)*up for every slot AND that slot's (max,1/max) in one pass — the down GEMV needs the
// scale anyway and this kernel already walks the data, so it replaces silu_mul + absmax_slots.
extern "C" __global__ void silu_mul_moe(const float* __restrict__ gu, float* __restrict__ out,
                                        float* __restrict__ exs, const int mi) {
    const int s = blockIdx.x;
    __shared__ float red[256];
    const float* g = gu + (long)s * 2 * mi;
    const float* u = g + mi;
    float* o = out + (long)s * mi;
    float m = 0.f;
    for (int i = threadIdx.x; i < mi; i += blockDim.x) {
        const float v = g[i];
        const float r = (v / (1.f + __expf(-v))) * u[i];
        o[i] = r;
        m = fmaxf(m, fabsf(r));
    }
    red[threadIdx.x] = m;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if ((int)threadIdx.x < off) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + off]);
        __syncthreads();
    }
    if (threadIdx.x == 0) { const float mx = red[0]; exs[2*s] = mx; exs[2*s+1] = (mx > 0.f) ? (1.f/mx) : 0.f; }
}

extern "C" __global__ void moe_combine(const float* __restrict__ dn, const float* __restrict__ wts,
                                       float* __restrict__ out, const int hidden, const int k) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= hidden) return;
    float s = 0.f;
    for (int j = 0; j < k; ++j) s = fmaf(wts[j], dn[(long)j * hidden + i], s);
    out[i] = s;
}

// Round f32 values to bf16 precision in place. HF's MoE gate is a bf16 Linear: logits are
// f32-accumulated then STORED bf16, and softmax/top-k read the rounded values. Top-k membership
// at near-ties depends on that rounding, so our routers reproduce it exactly.
extern "C" __global__ void bf16_round_f32(float* __restrict__ x, const int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) x[i] = bf2f(f2bf(x[i]));
}

// BATCHED top-k: blockIdx.x = row. Same softmax + k argmax rounds as moe_topk, over that
// row's logits. p[] holds the whole prob row: n_exp <= 512 (Qwen3.6-35B has 256).
extern "C" __global__ void moe_topk_b(const float* __restrict__ logits, unsigned int* __restrict__ idx,
                                      float* __restrict__ wts, const int n_exp, const int k,
                                      const int norm) {
    __shared__ float p[512];
    __shared__ float red[256];
    __shared__ int ired[256];
    __shared__ float wsum_s;
    const int t = threadIdx.x;
    const float* lg = logits + (long)blockIdx.x * n_exp;
    unsigned int* oid = idx + (long)blockIdx.x * k;
    float* owt = wts + (long)blockIdx.x * k;
    float mx = -1e30f;
    for (int i = t; i < n_exp; i += blockDim.x) mx = fmaxf(mx, lg[i]);
    red[t] = mx; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if (t < o) red[t] = fmaxf(red[t], red[t + o]); __syncthreads(); }
    const float m = red[0];
    __syncthreads();
    float sum = 0.f;
    for (int i = t; i < n_exp; i += blockDim.x) { const float e = __expf(lg[i] - m); p[i] = e; sum += e; }
    red[t] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if (t < o) red[t] += red[t + o]; __syncthreads(); }
    const float inv = 1.f / red[0];
    __syncthreads();
    for (int i = t; i < n_exp; i += blockDim.x) p[i] *= inv;
    if (t == 0) wsum_s = 0.f;
    __syncthreads();
    for (int j = 0; j < k; ++j) {
        float bv = -1.f; int bi = 0;
        for (int i = t; i < n_exp; i += blockDim.x) if (p[i] > bv) { bv = p[i]; bi = i; }
        red[t] = bv; ired[t] = bi; __syncthreads();
        for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
            if (t < o && red[t + o] > red[t]) { red[t] = red[t + o]; ired[t] = ired[t + o]; }
            __syncthreads();
        }
        if (t == 0) {
            oid[j] = (unsigned int)ired[0];
            owt[j] = red[0];
            wsum_s += red[0];
            p[ired[0]] = -1.f;
        }
        __syncthreads();
    }
    if (norm && t == 0 && wsum_s > 0.f) for (int j = 0; j < k; ++j) owt[j] /= wsum_s;
}

// BATCHED combine: out[row, :] = sum_j wts[row*k+j] * dn[(row*k+j), :] — one thread per
// (row, col) element over rows*hidden.
extern "C" __global__ void moe_combine_b(const float* __restrict__ dn, const float* __restrict__ wts,
                                         float* __restrict__ out, const int hidden, const int k,
                                         const int rows) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= rows * hidden) return;
    const int row = i / hidden, col = i - row * hidden;
    float s = 0.f;
    const long pb = (long)row * k;
    for (int j = 0; j < k; ++j) s = fmaf(wts[pb + j], dn[(pb + j) * hidden + col], s);
    out[i] = s;
}

// Also zeroes h: layer 0's norm folds in a residual add, and at that point h would otherwise
// still hold the previous token's MLP output.
// Batched q4_0 GEMV: ONE weight read serves M sequences. Decode is memory-bound, so this is where
// concurrency pays — single-stream reads the full weight set per token, M-way batching divides that
// by M while multiplying only the (much smaller) activation traffic. Measured single-stream traffic
// was 14.2 MB weights + 3.1 MB activations per MoE gate|up GEMV; at M=8 with 32 rows/block it is
// 14.2 + 24.6 MB for EIGHT tokens = ~3.6x less per token.
#define MAXM 8
// Batched q4_0 GEMV — one weight read serves M sequences, which is the whole point of concurrency
// on a memory-bound decode. Two things had to be right:
//
//  (a) COALESCED activation staging. The obvious version computes each half2's permuted element
//      index and reads those two floats from global, so consecutive threads read scattered
//      addresses — and unlike the single-stream case that cost is multiplied by M. Measured: the
//      batched GEMV achieved 157 GB/s at M=8 (8.6% of peak) against a 956 MB/step weight floor
//      that should allow ~15k tok/s aggregate. Now threads walk ELEMENTS linearly (coalesced) and
//      the permutation happens on the shared-memory WRITE, which is free.
//  (b) 64 rows per block (32 warps x 2 rows). Every block re-stages the activation tile, so
//      staging traffic is (rows / rows_per_block) * M * cols * 4 — doubling rows halves it.
extern "C" __global__ void gemv_q4_batch(
    const unsigned short* __restrict__ scales,
    const uint4* __restrict__ quants,
    const float* __restrict__ X,     /* [M, cols] */
    const float* __restrict__ xsc,   /* [2M] */
    float* __restrict__ Y,           /* [M, rows] */
    const int nblk, const int rows, const int M, const int cols)
{
    extern __shared__ unsigned int xh[];            /* [M][TILEB*17] half2 slots */
    unsigned short* xh16 = (unsigned short*)xh;     /* half-granular view for permuted writes */
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int nwarp = blockDim.x >> 5;
    const int row0 = blockIdx.x * (nwarp * 2);
    const unsigned int bias = 0xE408E408u;
    float acc[2][MAXM];
    #pragma unroll
    for (int rr = 0; rr < 2; ++rr)
        #pragma unroll
        for (int m = 0; m < MAXM; ++m) acc[rr][m] = 0.f;

    for (int t0 = 0; t0 < nblk; t0 += TILEB) {
        __syncthreads();
        // coalesced: consecutive threads read consecutive activation elements
        for (int idx = threadIdx.x; idx < M * TILEB * 32; idx += blockDim.x) {
            const int m = idx / (TILEB * 32);
            const int e = idx - m * (TILEB * 32);
            const int b = e >> 5, el = e & 31;
            if ((t0 + b) >= nblk) { xh16[(m * (TILEB * 17) + b * 17) * 2] = 0; continue; }
            const float v = X[(long)m * cols + ((long)(t0 + b) << 5) + el] * xsc[2 * m + 1];
            // where does element `el` live? pairs are (a, a+2) with
            // a = 4*wi + (r>>1) + 16*(r&1); see quantize_q4_0's packing.
            const int hi16 = (el >= 16) ? 1 : 0;
            const int e2 = el - 16 * hi16;         // 0..15
            const int wi = e2 >> 2, off = e2 & 3;
            const int r = ((off & 1) ? 2 : 0) + hi16;
            const int slot = m * (TILEB * 17) + b * 17 + (wi * 4 + r);
            unsigned short h;
            asm("{ .reg .f16 t; cvt.rn.f16.f32 t, %1; mov.b16 %0, t; }" : "=h"(h) : "f"(v));
            xh16[slot * 2 + ((off < 2) ? 0 : 1)] = h;   // low half or high half
        }
        __syncthreads();
        #pragma unroll
        for (int rr = 0; rr < 2; ++rr) {
            const int row = row0 + warp * 2 + rr;
            if (row >= rows) continue;
            const long base = (long)row * (long)nblk;
            for (int b = lane; b < TILEB && (t0 + b) < nblk; b += 32) {
                const uint4 q = quants[base + t0 + b];
                const float sc = h2f(scales[base + t0 + b]);
                unsigned int w[4] = { q.x, q.y, q.z, q.w };
                for (int m = 0; m < M; ++m) {
                    const unsigned int* xp = xh + m * (TILEB * 17) + b * 17;
                    float blk = 0.f;
                    #pragma unroll
                    for (int i = 0; i < 4; ++i) {
                        const unsigned int v = w[i];
                        unsigned int sw = 0u;
                        #pragma unroll
                        for (int r2 = 0; r2 < 4; ++r2) {
                            const unsigned int mm = ((v >> (4 * r2)) & 0x000F000Fu) | 0x64006400u;
                            sw = hfma2(hadd2(mm, bias), xp[i * 4 + r2], sw);
                        }
                        float lo, hi2; unpack_h2(sw, &lo, &hi2);
                        blk += lo + hi2;
                    }
                    acc[rr][m] = fmaf(sc, blk, acc[rr][m]);
                }
            }
        }
    }
    #pragma unroll
    for (int rr = 0; rr < 2; ++rr) {
        const int row = row0 + warp * 2 + rr;
        for (int m = 0; m < M; ++m) {
            float a = acc[rr][m];
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) a += __shfl_down_sync(0xffffffffu, a, o);
            if (lane == 0 && row < rows) Y[(long)m * rows + row] = a * xsc[2 * m];
        }
    }
}

extern "C" __global__ void embed_lookup(const float* __restrict__ e, float* __restrict__ x,
                                        float* __restrict__ h,
                                        const unsigned int* __restrict__ token, const int n,
                                        const float esc) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) { x[i] = e[(long)(*token) * n + i] * esc; h[i] = 0.f; }
}

// Append k/v at the DEVICE-side position, replacing a host-offset memcpy so the KV write can live
// inside the captured graph.
extern "C" __global__ void kv_write(const float* __restrict__ qkv, const float* __restrict__ _unused,
                                    float* __restrict__ kc, float* __restrict__ vc,
                                    const int* __restrict__ pos, const int kvdim,
                                    const int k_off, const int v_off) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= kvdim) return;
    const long off = (long)(*pos) * (long)kvdim + i;
    kc[off] = qkv[k_off + i];
    vc[off] = qkv[v_off + i];
}

// Record the sampled token and advance the position — the only per-step state mutation, on device.
extern "C" __global__ void advance(int* __restrict__ pos, const unsigned int* __restrict__ tok,
                                   unsigned int* __restrict__ out, const int max_out) {
    if (threadIdx.x == 0) {
        const int p = *pos;
        if (p < max_out) out[p] = *tok;
        *pos = p + 1;
    }
}

// Emits the normalised vector AND its (max|.|, 1/max|.|) — the GEMV needs that scale anyway, and
// this kernel already owns a block reduction, so folding it in removes one launch per norm. At 48
// layers x 2 norms that is ~96 launches/token, and this engine is launch-bound.
// x += h, then RMSNorm(x)*w into out, plus (max,1/max) — folds the residual add into the norm
// that always follows it. Saves 2 launches per layer.
extern "C" __global__ void rmsnorm_add(float* __restrict__ x, const float* __restrict__ h,
                                       const float* __restrict__ w, float* __restrict__ out,
                                       float* __restrict__ xs, const int n, const float eps) {
    // MEASURED: 14.68 us x ~56 per M=1 step. One block for the whole row, so the two smem
    // reduction trees (log2(blockDim) barriers each) dominate a kernel that moves ~24 KB.
    // Warp-shuffle both: two barriers total, and red[] sized by warp count so the kernel is
    // correct at any block size.
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = x[i] + h[i]; x[i] = v; s = fmaf(v, v, s); }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float inv = rsqrtf(red[0] / (float)n + eps);
    __syncthreads();
    float m = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        const float v = x[i] * inv * w[i];
        out[i] = v;
        m = fmaxf(m, fabsf(v));
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_down_sync(0xffffffffu, m, o));
    if (lane == 0) red[wid] = m;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t = fmaxf(t, __shfl_down_sync(0xffffffffu, t, o));
        if (lane == 0) { xs[0] = t; xs[1] = (t > 0.f) ? (1.f / t) : 0.f; }
    }
}

extern "C" __global__ void rmsnorm(const float* __restrict__ x, const float* __restrict__ w,
                                   float* __restrict__ out, float* __restrict__ xs,
                                   const int n, const float eps) {
    __shared__ float red[256];
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) s = fmaf(x[i], x[i], s);
    red[threadIdx.x] = s;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o];
        __syncthreads();
    }
    const float inv = rsqrtf(red[0] / (float)n + eps);
    __syncthreads();
    float m = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        const float v = x[i] * inv * w[i];
        out[i] = v;
        m = fmaxf(m, fabsf(v));
    }
    red[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) { const float mx = red[0]; xs[0] = mx; xs[1] = (mx > 0.f) ? (1.f / mx) : 0.f; }
}

// One block per head: per-head RMSNorm over head_dim (Qwen3 q_norm/k_norm), then RoPE at `pos`.
// Blocks [0, n_heads) are query heads; [n_heads, n_heads+n_kv) are key heads.
extern "C" __global__ void qk_norm_rope(float* __restrict__ qkv, float* __restrict__ _unused,
                                        const float* __restrict__ qw, const float* __restrict__ kw,
                                        unsigned short* __restrict__ kc,
                                        unsigned short* __restrict__ vc,
                                        const int n_heads, const int n_kv, const int hd,
                                        const int* __restrict__ pos_p, const float* __restrict__ invf,
                                        const int qkn, const float eps, const int max_seq) {
    const int pos = *pos_p;
    const int hidx = blockIdx.x;
    // v heads: pure copy into the cache, no norm/rope
    if (hidx >= n_heads + n_kv) {
        const int vh = hidx - n_heads - n_kv;
        const float* src = qkv + (long)(n_heads + n_kv + vh) * hd;
        unsigned short* dst = vc + ((long)vh * max_seq + pos) * hd;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) dst[i] = f2bf(src[i]);
        return;
    }
    const bool is_q = hidx < n_heads;
    // q occupies rows [0, n_heads), k the next n_kv — both inside the fused qkv buffer
    float* base = qkv + (long)hidx * hd;
    const float* nw = is_q ? qw : kw;

    __shared__ float red[128];
    if (qkn) {
        float s = 0.f;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) s = fmaf(base[i], base[i], s);
        red[threadIdx.x] = s;
        __syncthreads();
        for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
            if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o];
            __syncthreads();
        }
        const float inv = rsqrtf(red[0] / (float)hd + eps);
        for (int i = threadIdx.x; i < hd; i += blockDim.x) base[i] = base[i] * inv * nw[i];
        __syncthreads();
    }
    // RoPE, NeoX/HF halves layout: pairs (i, i + hd/2)
    const int half = hd >> 1;
    for (int i = threadIdx.x; i < half; i += blockDim.x) {
        const float ang = (float)pos * invf[i];
        const float c = __cosf(ang), sn = __sinf(ang);
        const float a = base[i], b = base[i + half];
        base[i]        = a * c - b * sn;
        base[i + half] = a * sn + b * c;
    }
    if (!is_q) {
        __syncthreads();
        const int kh = hidx - n_heads;
        unsigned short* dst = kc + ((long)kh * max_seq + pos) * hd;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) dst[i] = f2bf(base[i]);
    }
}

// FlashDecoding-style split-K attention.
//
// The previous kernel used ONE block per head (32 blocks on a 132-SM GPU) and, worse, had each
// thread walk a whole 128-dim key: across the warp that is a 2 KB stride, so every score load was
// fully uncoalesced. At 2k context it measured 21.4 tok/s against vLLM's 155.8.
//
// Now each block owns one (head, position-chunk): grid (n_heads, nsplit) gives hundreds of blocks,
// a warp cooperates on a single position with float4 lanes (512 B contiguous per warp), and each
// block emits a PARTIAL softmax (running max, sum, weighted-V) that the reduce kernel merges with
// the standard log-sum-exp rescale.
// Positions per attention chunk.
//
// 256 leaves too little parallelism at long context: at T=2000 that is 8 chunks x 16 heads = 128
// blocks for 132 SMs, and each of a block's 8 warps then walks 32 positions SERIALLY with a
// dependent global load at the head of each. That is why single-stream decode degrades 22% at 2k
// (454 -> 354 tok/s) when the added KV traffic only accounts for a fifth of it. 64 gives 4x the
// blocks and an eighth of the serial walk.
#define ACHUNK 64
extern "C" __global__ void attn_partial(const float* __restrict__ q,
                                        const unsigned short* __restrict__ kc,
                                        const unsigned short* __restrict__ vc,
                                        float* __restrict__ pm, float* __restrict__ pl,
                                        float* __restrict__ pacc,
                                        const int n_heads, const int n_kv, const int hd,
                                        const int* __restrict__ pos_p, const int nsplit,
                                        const int max_seq) {
    const int h = blockIdx.x, sp = blockIdx.y;
    const int T = *pos_p + 1;
    const int kvh = h / (n_heads / n_kv);
    const long kvbase = (long)kvh * max_seq * hd;   // head-major plane in the (slot-0) slab
    const int start = sp * ACHUNK;
    const int end = min(start + ACHUNK, T);

    extern __shared__ float sh[];
    float* qs = sh;              // hd
    float* sc = sh + hd;         // ACHUNK
    for (int i = threadIdx.x; i < hd; i += blockDim.x) qs[i] = q[(long)h * hd + i];
    __syncthreads();
    if (start >= end) {
        if (threadIdx.x == 0) { pm[h * nsplit + sp] = -1e30f; pl[h * nsplit + sp] = 0.f; }
        for (int d = threadIdx.x; d < hd; d += blockDim.x) pacc[((long)h * nsplit + sp) * hd + d] = 0.f;
        return;
    }

    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int nwarp = blockDim.x >> 5;
    const float scale = rsqrtf((float)hd);
    // one position per warp; lanes take 4 contiguous dims each => coalesced 512 B per warp
    for (int t = start + warp; t < end; t += nwarp) {
        const unsigned short* kt = kc + kvbase + (long)t * hd;
        float d = 0.f;
        for (int i = lane * 4; i < hd; i += 128) {
            // four bf16 = 8 bytes; this loop is the long-context cost, so halving it is the point
            const uint2 kk = *reinterpret_cast<const uint2*>(kt + i);
            d = fmaf(qs[i + 0], bf2f((unsigned short)(kk.x & 0xFFFFu)), d);
            d = fmaf(qs[i + 1], bf2f((unsigned short)(kk.x >> 16)), d);
            d = fmaf(qs[i + 2], bf2f((unsigned short)(kk.y & 0xFFFFu)), d);
            d = fmaf(qs[i + 3], bf2f((unsigned short)(kk.y >> 16)), d);
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
        if (lane == 0) sc[t - start] = d * scale;
    }
    __syncthreads();

    __shared__ float red[256];
    const int n = end - start;
    float m = -1e30f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) m = fmaxf(m, sc[i]);
    red[threadIdx.x] = m; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + o]); __syncthreads(); }
    const float mx = red[0];
    __syncthreads();
    float sum = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float e = __expf(sc[i] - mx); sc[i] = e; sum += e; }
    red[threadIdx.x] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o]; __syncthreads(); }
    const float ls = red[0];
    __syncthreads();

    // Weighted V. The obvious loop (thread per dim, scalar, sequential over positions) leaves
    // half the block idle at hd=128 and reads 4 B per thread per step; at 2k context that phase was
    // the long-context bottleneck (measured 267 GB/s on the KV read = 15% of peak).
    // Instead: 8 warp-groups each take every 8th position, a lane owns 4 dims via float4, and the
    // groups are summed at the end. All 256 threads busy, 16 B per lane per step.
    if (hd == 128) {
        __shared__ float vred[8 * 128];
        const int g = threadIdx.x >> 5;
        float4 a4 = make_float4(0.f, 0.f, 0.f, 0.f);
        for (int t = start + g; t < end; t += 8) {
            const float pw = sc[t - start];
            const uint2 vv = *reinterpret_cast<const uint2*>(
                vc + kvbase + (long)t * hd + lane * 4);
            a4.x = fmaf(pw, bf2f((unsigned short)(vv.x & 0xFFFFu)), a4.x);
            a4.y = fmaf(pw, bf2f((unsigned short)(vv.x >> 16)), a4.y);
            a4.z = fmaf(pw, bf2f((unsigned short)(vv.y & 0xFFFFu)), a4.z);
            a4.w = fmaf(pw, bf2f((unsigned short)(vv.y >> 16)), a4.w);
        }
        *reinterpret_cast<float4*>(vred + g * 128 + lane * 4) = a4;
        __syncthreads();
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float acc2 = 0.f;
            #pragma unroll
            for (int gg = 0; gg < 8; ++gg) acc2 += vred[gg * 128 + d];
            pacc[((long)h * nsplit + sp) * hd + d] = acc2;
        }
    } else {
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float a = 0.f;
            for (int t = start; t < end; ++t)
                a = fmaf(sc[t - start], vc[kvbase + (long)t * hd + d], a);
            pacc[((long)h * nsplit + sp) * hd + d] = a;
        }
    }
    if (threadIdx.x == 0) { pm[h * nsplit + sp] = mx; pl[h * nsplit + sp] = ls; }
}

// Merge the per-chunk partials: rescale each by exp(m_i - m_global) and normalise.
extern "C" __global__ void attn_reduce(const float* __restrict__ pm, const float* __restrict__ pl,
                                       const float* __restrict__ pacc, float* __restrict__ out,
                                       const int hd, const int nsplit) {
    const int h = blockIdx.x;
    __shared__ float mg, lg;
    if (threadIdx.x == 0) {
        float m = -1e30f;
        for (int s = 0; s < nsplit; ++s) m = fmaxf(m, pm[h * nsplit + s]);
        float l = 0.f;
        for (int s = 0; s < nsplit; ++s) l += pl[h * nsplit + s] * __expf(pm[h * nsplit + s] - m);
        mg = m; lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    for (int d = threadIdx.x; d < hd; d += blockDim.x) {
        float a = 0.f;
        for (int s = 0; s < nsplit; ++s)
            a = fmaf(__expf(pm[h * nsplit + s] - mg), pacc[((long)h * nsplit + s) * hd + d], a);
        out[(long)h * hd + d] = a / lg;
    }
}

// ---------------- batched (concurrent-sequence) variants ----------------
// M sequences advance in lockstep, each with its own KV slab (sequence m at m*max_seq*kvdim),
// its own token and its own activation scale. Grids carry the batch dimension.
// ---------------- the compute-bound arm: dequant q4_0 -> f16, then cuBLAS on tensor cores ----
// Single-stream decode is memory-bound and the q4_0 GEMV wins there. But PREFILL (M = prompt
// length) and BATCHED decode (M = concurrent sequences) are compute-bound, and the f16x2 FMA path
// tops out near 1.3 TFLOP/s against cuBLAS's measured 558. Measured consequence: we beat vLLM at
// M=1 (324 vs 317) and lose at M=8 (1324 vs 2352). So for M > 1 the weights get dequantised into an
// f16 scratch once per matrix and the GEMM runs on tensor cores; the dequant is amortised over all
// M columns, which is why it only pays off for M > 1.
__device__ __forceinline__ unsigned int hmul2(unsigned int a, unsigned int b) {
    unsigned int d;
    asm("mul.rn.f16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b));
    return d;
}



// silu(gate)*up with bf16 in AND out.
//
// The f32 version cost ~400 MB of activation traffic per layer (two f32 [T,inter] GEMM outputs,
// read back, written, then converted again for the down GEMM). Keeping the FFN interior in bf16 —
// the precision vLLM and SGLang use for activations anyway — cuts that to ~160 MB and lets the
// down GEMM consume this output with no conversion pass at all.
// Layout: `gu` is [T, 2*mi] row-major (one merged gate|up GEMM), gate at +j, up at +mi+j.
// silu(gate)*up straight from a MERGED f32 [T, 2*mi] result into bf16 [T, mi].
//
// The q4 arm ran gate and up as two separate GEMMs and then a separate f32->bf16 conversion before
// the down projection: 3 kernels where the bf16 path uses 2. gate and up share their input, so one
// GEMM over the concatenated weight produces both, and this kernel folds the activation and the
// conversion together — 5 GEMMs + 2 conversions per layer becomes 4 + 1.
extern "C" __global__ void silu_gu_bf(const float* __restrict__ gu,
                                      unsigned short* __restrict__ out,
                                      const int mi, const int n, const int act) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    const int t = i / mi, j = i - t * mi;
    const long base = (long)t * 2 * mi;
    const float g = gu[base + j], u = gu[base + mi + j];
    out[i] = f2bf(ffn_act(g, act) * u);
}

extern "C" __global__ void silu_mul_bf(const unsigned short* __restrict__ gu,
                                       unsigned short* __restrict__ out,
                                       const int mi, const int n8, const int act) {
    // EIGHT halves per thread. The scalar version moved 74 MB per layer at ~1.07 TB/s (half of
    // HBM peak) because every access was 2 bytes; 16-byte access doubles it. n8 = (T*mi)/8.
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n8) return;
    const int per = mi >> 3;
    const int t = i / per, j = i - t * per;
    const long base = (long)t * 2 * mi + (long)j * 8;
    const uint4 gv = *(const uint4*)(gu + base);
    const uint4 uv = *(const uint4*)(gu + base + mi);
    const unsigned int* gw = (const unsigned int*)&gv;
    const unsigned int* uw = (const unsigned int*)&uv;
    unsigned int ow[4];
    #pragma unroll
    for (int k = 0; k < 4; ++k) {
        const float g0 = bf2f((unsigned short)(gw[k] & 0xFFFFu)), u0 = bf2f((unsigned short)(uw[k] & 0xFFFFu));
        const float g1 = bf2f((unsigned short)(gw[k] >> 16)),     u1 = bf2f((unsigned short)(uw[k] >> 16));
        const unsigned int lo = f2bf(ffn_act(g0, act) * u0);
        const unsigned int hi = f2bf(ffn_act(g1, act) * u1);
        ow[k] = lo | (hi << 16);
    }
    *(uint4*)(out + (long)i * 8) = *(const uint4*)ow;
}

// ---------------------------------------------------------------------------------------------
// Flash-attention prefill on tensor cores.
//
// The GEMM-based path materialises S = QK^T as [nh, T, T] f16 (128 MB per layer at T=2000), the
// softmax reads and rewrites it, then PV reads it again — ~13.2 ms of the 31.6 ms prefill, of which
// roughly two thirds is score traffic and half the FLOPs are the masked upper triangle we compute
// and throw away. This kernel never writes S: it tiles (query block x key block), keeps the running
// max/sum and the whole O accumulator in registers, and skips key blocks the causal mask kills.
//
// grid (ceil(T/BQ), n_heads), block 128 threads = 4 warps; warp w owns query rows [16w, 16w+16).
// mma.sync.m16n8k16 exists only in .row.col, so the B operand is indexed [n][k]: for O = P.V that
// means V must sit in shared memory TRANSPOSED, hence sV[dim][key].
//
// The reason no data movement is needed between the two mma chains: the f32 C-fragment layout of
// S (lane l holds rows l>>2 and (l>>2)+8, columns (l&3)*2 + {0,1} of each 8-wide n-tile) is
// exactly the f16 A-fragment layout that P needs. Pairs of S n-tiles pack straight into A.
#define FA_BQ 64
#define FA_BK 32
#define FA_HD 128           // Qwen3/Llama/Mistral all use 128; other head dims take the GEMM path
#define FA_MASKV (-1e30f)
// PADDED row strides. Unpadded, a fragment load addresses row*(128*2) bytes, which is 0 mod the
// 128-byte bank cycle, so all eight `gid` rows of a warp hit the SAME four banks: an 8-way conflict
// on every A/B load (16-way on the transposed V, stride 64 B). +8 halves shifts each row by 16 B =
// 4 banks, spreading the eight rows across all 32. Measured 33.6 ms -> see commit message.
#define FA_QS (FA_HD + 8)
#define FA_KS (FA_HD + 8)
#define FA_VS (FA_BK + 8)

// cp.async: global -> shared without a register staging round-trip, and asynchronous, so the tile
// for key block kb+1 streams in while the mma pipe chews on kb. ncu put the tensor pipe at 13.2%
// active with 15% warp occupancy — the kernel was waiting on loads, not computing.
__device__ __forceinline__ void fa_cp16(void* smem, const void* gmem) {
    unsigned int s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(s), "l"(gmem));
}
// The src-size form: copies `sz` bytes (0 or 16 here) and ZERO-FILLS the rest of the 16-byte
// cp-size. PTX guarantees src is not read when src-size is 0, which is what makes it the
// out-of-bounds guard for tiles that extend past the padded buffer: the wgmma tile width (BK=128)
// exceeds the 64-row padding granularity of qkv16/vt16, so a short prompt's last tile MUST NOT
// dereference the missing rows — but the tile still has to be zeroed so the masked softmax path
// stays branch-free.
__device__ __forceinline__ void fa_cp16z(void* smem, const void* gmem, unsigned sz) {
    unsigned int s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" :: "r"(s), "l"(gmem), "r"(sz));
}
#define FA_CP_COMMIT() asm volatile("cp.async.commit_group;\n" ::)
#define FA_CP_WAIT1()  asm volatile("cp.async.wait_group 1;\n" ::)

__device__ __forceinline__ unsigned int fa_pack2(float a, float b) {
    unsigned int r;
    asm("{ .reg .f16 lo, hi; cvt.rn.f16.f32 lo, %1; cvt.rn.f16.f32 hi, %2; mov.b32 %0, {lo,hi}; }"
        : "=r"(r) : "f"(a), "f"(b));
    return r;
}

#define FA_MMA(d0,d1,d2,d3, a0,a1,a2,a3, b0,b1) \
    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " \
                 "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" \
                 : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) \
                 : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1))

// ldmatrix — ONE instruction for a whole 16x16 fragment, replacing four 32-bit ld.shared plus their
// address arithmetic. This is what the kernel was actually bound on: ncu put ALU (integer/address
// math) at 33.3% as the top pipeline, ahead of the tensor pipe, with L1/TEX at 56% and IPC 2.2/4 —
// 160 scalar shared loads per warp per key block against only 64 mma. ldmatrix takes it to 40.
//
// Each lane supplies the address of ONE row of an 8x8 f16 core matrix and the hardware gathers into
// exactly the register layout mma.m16n8k16 wants — lanes 0-7 address matrix 0, 8-15 matrix 1, and
// so on, with the returned d0..d3 being those four matrices. The matrices are ordered so that a
// 16x16 A operand is (rows 0-7 | rows 8-15 | rows 0-7,cols 8-15 | rows 8-15,cols 8-15), which is
// a0,a1,a2,a3 unchanged; for a B operand the same call covers TWO adjacent n-tiles at once
// (d0,d2 = tile n, d1,d3 = tile n+1).
//
// The existing +8 row padding is what makes this conflict-free: at a 272-byte row stride each of
// the 8 rows a matrix gathers shifts by 16 B = 4 banks, so they tile all 32 banks exactly once
// (the transposed V's 80-byte stride does the same by a different permutation).
__device__ __forceinline__ void fa_ldm4(unsigned int& d0, unsigned int& d1, unsigned int& d2,
                                        unsigned int& d3, const unsigned short* smem) {
    unsigned int s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) : "r"(s));
}

// V transposed once per layer, so flash attention never transposes it per key block.
//
// mma .row.col wants the B operand indexed [n][k]; for O = P.V that is V^T. Building it inside the
// attention kernel meant eight scalar shared stores per 16-byte read, 16-way bank conflicted,
// repeated for every (query block, key block) pair. Doing it once costs one 8 MB pass per layer.
#define VT_TILE 32
extern "C" __global__ void vt_pack(const unsigned short* __restrict__ qkv16,
                                   unsigned short* __restrict__ vt,
                                   const int T, const int Tp, const int n_heads, const int n_kv,
                                   const int hd) {
    // Reads the f16 copy, not the f32 original: rope_prefill has already written the V heads into
    // qkv16 by the time this runs, and the f32 read was half this kernel's traffic.
    __shared__ unsigned short sh[VT_TILE][VT_TILE + 1];
    const int per = n_heads + 2 * n_kv;
    const int kvh = blockIdx.z;
    const int t0 = blockIdx.x * VT_TILE, d0 = blockIdx.y * VT_TILE;
    const int lx = threadIdx.x;
    for (int r = threadIdx.y; r < VT_TILE; r += blockDim.y) {
        const int t = t0 + r, d = d0 + lx;
        unsigned short hh = 0;
        if (t < T && d < hd)
            hh = qkv16[(long)t * per * hd + (long)(n_heads + n_kv + kvh) * hd + d];
        sh[r][lx] = hh;
    }
    __syncthreads();
    for (int r = threadIdx.y; r < VT_TILE; r += blockDim.y) {
        const int d = d0 + r, t = t0 + lx;
        if (d < hd && t < T) vt[((long)kvh * hd + d) * Tp + t] = sh[lx][r];
    }
}

#if __CUDA_ARCH__ >= 800
// 4 warps of 16 query rows each, and the KEY RANGE IS SPLIT across blockIdx.z.
//
// Causal work per query block grows linearly with qb, so with one block per (query block, head)
// the last block does 32x the work of the first. All 512 blocks fit on the GPU at once, which
// means the kernel's makespan was simply the largest block while the SMs holding small blocks sat
// idle — ncu measured 11.3% warps active against 3 resident blocks. Splitting the keys halves the
// largest block and doubles the block count, so the scheduler can pack instead of straggle. Each
// block emits an unnormalised partial (running max, sum, O) that fa_merge combines by log-sum-exp,
// the same structure attn_partial/attn_reduce already use for decode.
//
// MEASURED NEGATIVE, kept as a note: splitting the key range across blockIdx.z (partial max/sum/O
// merged by log-sum-exp, as decode does) made attention SLOWER, 6.46 -> 7.50 ms. The imbalance is
// real, but writing and re-reading 33 MB of f32 partials per layer costs more than the packing
// buys. BQ=128 also measured worse (a spill plus one fewer resident block), and BK=64 needs 89 KB
// of shared, capping the SM at 2 resident blocks.
extern "C" __global__ __launch_bounds__(128, 4) void fa_prefill(
                                      const unsigned short* __restrict__ qkv16,
                                      const unsigned short* __restrict__ vt16,
                                      unsigned short* __restrict__ out,
                                      const int T, const int Tp, const int n_heads, const int n_kv,
                                      const float scale) {
    const int qb = blockIdx.x, h = blockIdx.y;
    const int per = n_heads + 2 * n_kv;
    const int kvh = h / (n_heads / n_kv);
    const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
    const int gid = lane >> 2, t4 = lane & 3;
    // ldmatrix addressing: lane l supplies the address of row (l&7) of core matrix (l>>3), where
    // matrices 0/1 are the two 8-row halves of the 16-row tile and 2/3 are those halves 8 columns
    // over. Loop-invariant, so the whole address collapses to one add inside the k loop.
    const int lrow = 8 * ((lane >> 3) & 1) + (lane & 7);
    const int lcol = 8 * (lane >> 4);
    const int q0 = qb * FA_BQ;

    extern __shared__ unsigned short fash[];
    unsigned short* sQ = fash;                          // [FA_BQ][FA_QS]
    unsigned short* sKb = sQ + FA_BQ * FA_QS;           // [2][FA_BK][FA_KS]
    unsigned short* sVb = sKb + 2 * FA_BK * FA_KS;      // [2][FA_HD][FA_VS], transposed

    // 16-byte (8-half) staging. Scalar 2-byte loads meant 64 dependent global loads per thread per
    // key block; at 128 registers only 4 blocks fit per SM, so that loop was pure latency —
    // ~1 ms per layer, which is exactly what the profile showed.
    // Q is pre-multiplied by 1/sqrt(head_dim) HERE, once, rather than scaling every score in every
    // key block: 32 multiplies per key block per warp removed. Post-rmsnorm q is O(1) and the scale
    // is 0.088, so the f16 product is nowhere near the format's limits.
    for (int i = threadIdx.x; i < FA_BQ * (FA_HD / 8); i += blockDim.x) {
        const int r = i / (FA_HD / 8), c = i - r * (FA_HD / 8);
        const uint4 g = *(const uint4*)(qkv16 + (long)(q0 + r) * per * FA_HD + (long)h * FA_HD + c * 8);
        const unsigned int* gw = (const unsigned int*)&g;
        unsigned int sw[4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const float lo = h2f((unsigned short)(gw[j] & 0xFFFFu)) * scale;
            const float hi = h2f((unsigned short)(gw[j] >> 16)) * scale;
            sw[j] = fa_pack2(lo, hi);
        }
        *(uint4*)(sQ + r * FA_QS + c * 8) = *(const uint4*)sw;
    }

    // O accumulator: 16 n-tiles of 8 dims, 4 f32 each (rows gid, gid+8)
    float o[FA_HD / 8][4];
    #pragma unroll
    for (int nt = 0; nt < FA_HD / 8; ++nt) { o[nt][0] = 0.f; o[nt][1] = 0.f; o[nt][2] = 0.f; o[nt][3] = 0.f; }
    float m0 = FA_MASKV, m1 = FA_MASKV, l0 = 0.f, l1 = 0.f;

    const int row0 = q0 + 16 * w + gid;          // global query row of C-frag regs 0,1
    const int row1 = row0 + 8;                   //                         regs 2,3
    const int kb_max = (q0 + FA_BQ - 1) / FA_BK; // causal: no key block beyond this one contributes
    const int kb_lo = 0, kb_hi = kb_max + 1;

    // Both source buffers are allocated with T rounded up to FA_BK and zero-filled, so a tile read
    // never needs a tail branch — which is what lets every load be an unconditional cp.async.
    #define FA_STAGE(KB, BUF) do { \
        unsigned short* dK = sKb + (BUF) * (FA_BK * FA_KS); \
        unsigned short* dV = sVb + (BUF) * (FA_HD * FA_VS); \
        const int _k0 = (KB) * FA_BK; \
        for (int i = threadIdx.x; i < FA_BK * (FA_HD / 8); i += blockDim.x) { \
            const int r = i / (FA_HD / 8), c = i - r * (FA_HD / 8); \
            fa_cp16(dK + r * FA_KS + c * 8, \
                    qkv16 + (long)(_k0 + r) * per * FA_HD + (long)(n_heads + kvh) * FA_HD + c * 8); \
        } \
        for (int i = threadIdx.x; i < FA_HD * (FA_BK / 8); i += blockDim.x) { \
            const int ch = i & (FA_BK / 8 - 1), d = i / (FA_BK / 8); \
            fa_cp16(dV + d * FA_VS + ch * 8, \
                    vt16 + ((long)kvh * FA_HD + d) * Tp + _k0 + ch * 8); \
        } \
    } while (0)

    FA_STAGE(kb_lo, 0);
    FA_CP_COMMIT();

    for (int kb = kb_lo; kb < kb_hi; ++kb) {
        const int k0 = kb * FA_BK;
        // Parity RELATIVE TO kb_lo: the prologue stages into buffer 0, so a slice that starts on
        // an odd tile would otherwise read the buffer it never filled.
        const int buf = (kb - kb_lo) & 1;
        const unsigned short* sK = sKb + buf * (FA_BK * FA_KS);
        const unsigned short* sV = sVb + buf * (FA_HD * FA_VS);
        // Issue kb+1 BEFORE waiting on kb. An empty commit on the last iteration keeps the
        // wait_group 1 accounting uniform, so the final tile is still guaranteed complete.
        if (kb + 1 < kb_hi) FA_STAGE(kb + 1, buf ^ 1);
        FA_CP_COMMIT();
        FA_CP_WAIT1();
        __syncthreads();

        // ---- S = Q . K^T for this key block: 4 n-tiles of 8 keys, hd/16 k-steps ----
        float sc[FA_BK / 8][4];
        #pragma unroll
        for (int nt = 0; nt < FA_BK / 8; ++nt) { sc[nt][0] = 0.f; sc[nt][1] = 0.f; sc[nt][2] = 0.f; sc[nt][3] = 0.f; }
        #pragma unroll
        for (int ks = 0; ks < FA_HD / 16; ++ks) {
            unsigned int a0, a1, a2, a3;
            fa_ldm4(a0, a1, a2, a3, sQ + (long)(16 * w + lrow) * FA_QS + ks * 16 + lcol);
            // one ldmatrix per PAIR of n-tiles: d0,d2 are tile 2np's two B regs, d1,d3 are 2np+1's
            #pragma unroll
            for (int np = 0; np < FA_BK / 16; ++np) {
                unsigned int b0, b1, b2, b3;
                fa_ldm4(b0, b1, b2, b3, sK + (long)(16 * np + lrow) * FA_KS + ks * 16 + lcol);
                FA_MMA(sc[2 * np][0], sc[2 * np][1], sc[2 * np][2], sc[2 * np][3],
                       a0, a1, a2, a3, b0, b2);
                FA_MMA(sc[2 * np + 1][0], sc[2 * np + 1][1], sc[2 * np + 1][2], sc[2 * np + 1][3],
                       a0, a1, a2, a3, b1, b3);
            }
        }

        // ---- causal mask (only where the block actually straddles) + online softmax ----
        // Warp-uniform: for all but the diagonal blocks every column is <= every row of this warp,
        // so the 32 per-element compares are pure waste. At T=2000 that is ~62 of 64 blocks.
        const bool nomask = (k0 + FA_BK - 1 <= q0 + 16 * w) && (k0 + FA_BK <= T);
        float t0 = FA_MASKV, t1 = FA_MASKV;
        if (nomask) {
            #pragma unroll
            for (int nt = 0; nt < FA_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    t0 = fmaxf(t0, sc[nt][j]); t1 = fmaxf(t1, sc[nt][2 + j]);
                }
            }
        } else {
            #pragma unroll
            for (int nt = 0; nt < FA_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    const int col = k0 + 8 * nt + t4 * 2 + j;
                    float v0 = sc[nt][j], v1 = sc[nt][2 + j];
                    if (col > row0 || col >= T) v0 = FA_MASKV;
                    if (col > row1 || col >= T) v1 = FA_MASKV;
                    sc[nt][j] = v0; sc[nt][2 + j] = v1;
                    t0 = fmaxf(t0, v0); t1 = fmaxf(t1, v1);
                }
            }
        }
        // a row's 32 values live in the four lanes sharing gid, i.e. differing in the low 2 bits
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            t0 = fmaxf(t0, __shfl_xor_sync(0xffffffffu, t0, x));
            t1 = fmaxf(t1, __shfl_xor_sync(0xffffffffu, t1, x));
        }
        const float nm0 = fmaxf(m0, t0), nm1 = fmaxf(m1, t1);
        // both still the sentinel => this key block is entirely masked for that row: contribute 0
        const bool live0 = nm0 > -1e29f, live1 = nm1 > -1e29f;
        const float r0 = __expf(m0 - nm0), r1 = __expf(m1 - nm1);
        float ps0 = 0.f, ps1 = 0.f;
        #pragma unroll
        for (int nt = 0; nt < FA_BK / 8; ++nt) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const float e0 = live0 ? __expf(sc[nt][j] - nm0) : 0.f;
                const float e1 = live1 ? __expf(sc[nt][2 + j] - nm1) : 0.f;
                sc[nt][j] = e0; sc[nt][2 + j] = e1;
                ps0 += e0; ps1 += e1;
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            ps0 += __shfl_xor_sync(0xffffffffu, ps0, x);
            ps1 += __shfl_xor_sync(0xffffffffu, ps1, x);
        }
        l0 = l0 * r0 + ps0; l1 = l1 * r1 + ps1;
        m0 = nm0; m1 = nm1;
        // Measured: gating this on a warp vote (skip when the running max did not move) was SLOWER
        // — 64 FFMAs are cheaper than one warp-wide __any_sync. Kept unconditional deliberately.
        #pragma unroll
        for (int nt = 0; nt < FA_HD / 8; ++nt) {
            o[nt][0] *= r0; o[nt][1] *= r0; o[nt][2] *= r1; o[nt][3] *= r1;
        }

        // ---- O += P . V. The S accumulators ARE the A fragments; only the f32->f16 pack. ----
        #pragma unroll
        for (int ks = 0; ks < FA_BK / 16; ++ks) {
            const unsigned int a0 = fa_pack2(sc[2 * ks][0],     sc[2 * ks][1]);
            const unsigned int a1 = fa_pack2(sc[2 * ks][2],     sc[2 * ks][3]);
            const unsigned int a2 = fa_pack2(sc[2 * ks + 1][0], sc[2 * ks + 1][1]);
            const unsigned int a3 = fa_pack2(sc[2 * ks + 1][2], sc[2 * ks + 1][3]);
            #pragma unroll
            for (int np = 0; np < FA_HD / 16; ++np) {
                unsigned int b0, b1, b2, b3;
                fa_ldm4(b0, b1, b2, b3, sV + (long)(16 * np + lrow) * FA_VS + ks * 16 + lcol);
                FA_MMA(o[2 * np][0], o[2 * np][1], o[2 * np][2], o[2 * np][3],
                       a0, a1, a2, a3, b0, b2);
                FA_MMA(o[2 * np + 1][0], o[2 * np + 1][1], o[2 * np + 1][2], o[2 * np + 1][3],
                       a0, a1, a2, a3, b1, b3);
            }
        }
        __syncthreads();     // buffer (kb&1) is refilled two iterations from now
    }
    #undef FA_STAGE

    // bf16 straight out: the o-projection is a bf16 GEMM, so an f32 result would only be re-read
    // and converted.
    const float i0 = (l0 > 0.f) ? 1.f / l0 : 0.f, i1 = (l1 > 0.f) ? 1.f / l1 : 0.f;
    #pragma unroll
    for (int nt = 0; nt < FA_HD / 8; ++nt) {
        const int d = 8 * nt + t4 * 2;
        if (row0 < T) {
            unsigned short* dst = out + (long)row0 * n_heads * FA_HD + (long)h * FA_HD + d;
            dst[0] = f2bf(o[nt][0] * i0); dst[1] = f2bf(o[nt][1] * i0);
        }
        if (row1 < T) {
            unsigned short* dst = out + (long)row1 * n_heads * FA_HD + (long)h * FA_HD + d;
            dst[0] = f2bf(o[nt][2] * i1); dst[1] = f2bf(o[nt][3] * i1);
        }
    }
}
#endif

#if __CUDA_ARCH__ < 800
// Pre-Ampere: the symbol must exist so module loading succeeds, but `fa` is false there.
extern "C" __global__ void fa_prefill(const unsigned short* __restrict__ qkv16,
                                      const unsigned short* __restrict__ vt16,
                                      unsigned short* __restrict__ out, const int T, const int Tp,
                                      const int n_heads, const int n_kv, const float scale) {}
#endif

// ---------------------------------------------------------------------------------------------
// HOPPER PREFILL ATTENTION — the same algorithm on wgmma instead of mma.sync.
//
// WHY, measured rather than assumed. The mma.sync kernel above sits at ~72 TFLOP/s at T=2000 and
// plateaus at ~93 by T=8000 (where four waves of blocks absorb the causal load imbalance), against
// a ~474 TFLOP/s mma.sync ceiling on this part. ncu says why: ALU — integer and address math — is
// the top pipeline at 33-41%, mem pipes are only 14% busy, IPC is 2.4 of 4, and 39.7% of cycles
// have NO eligible warp. It is issue- and latency-bound, not tensor-bound.
//
// MEASURED NEGATIVE, kept so nobody repeats it: replacing every fragment load with `ldmatrix`
// (160 scalar shared loads per warp per key block down to 40) changed nothing — 6.72 ms against
// 6.39. Fewer instructions do not help when the warps are stalled rather than issuing. That is the
// evidence that sent this kernel to wgmma: it is ASYNCHRONOUS, so a single warpgroup keeps the
// tensor pipe fed without needing warps to hide latency, it reads both operands straight from
// shared memory through a descriptor (no ldmatrix, no address math, no A/B registers), and it runs
// at twice the mma.sync rate. Per key block this is 10 wgmma where there were 256 mma.sync.
//
// The layouts below are not guessed. `bench/probes/wgmma_layout.cu` derives every one of them
// against a CPU reference, and the outcome is the reason this port is small: wgmma's accumulator
// fragment and its A-operand-in-registers fragment are IDENTICAL to the mma.m16n8k16 ones this
// file already uses — warp w owns rows 16w..16w+15, register group nt holds columns
// 8nt + (lane%4)*2 + {0,1} for rows lane/4 and lane/4+8. So the causal mask, the online softmax,
// the P-pack and the output write are all reused verbatim; only the multiply changes.
#if __CUDA_ARCH__ == 900
#define FH_BQ 64            // = the m64 of wgmma.m64nNk16; one warpgroup owns one query block
// 64, not 32. The ablation (ABL=1/2/3 in the history of this file) partitioned the kernel:
// softmax 29%, BOTH wgmma groups only 20%, cp.async staging 22%. The multiply is a fifth of the
// cost — which is why ldmatrix, the wgmma port and intra-warpgroup overlap were each worth ~0.
// The softmax's two 4-lane shuffle reduction chains, the 64-FFMA O rescale, the two barriers and
// the descriptor setup are all paid ONCE PER KEY BLOCK regardless of how wide it is, so doubling
// BK halves how often they are paid. Arithmetic intensity is unchanged (it depends only on BQ),
// so this is orthogonal to the GQA pairing above and composes with it.
#define FH_BK 64
#define FH_HD 128
// TWO query heads per block, one warpgroup each, sharing ONE staged K/V tile.
//
// This is the change that actually mattered, and the profile is what found it: ncu puts L2 at
// 59.8% (the top metric) against 31.9% compute, 7.7% mem-pipes and 57% of cycles with no eligible
// warp, DRAM at 2.6% behind a 96.4% L2 hit rate. The kernel is bound by K/V traffic out of L2, so
// neither ldmatrix nor wgmma could move it — both make the MULTIPLY cheaper.
//
// Arithmetic intensity here is exactly BQ: the work is 4*BQ*BK*HD FLOP against 4*BK*HD bytes of
// K/V, so bytes-per-FLOP depends on the query rows per staged tile and on NOTHING else. GQA hands
// that over for free — the grid ran 16 query heads over 8 KV heads and so read every K/V byte
// twice. Pairing the two heads that share a KV head doubles the intensity with no extra registers
// (each warpgroup keeps its own 64-row accumulator) and, unlike pairing adjacent query BLOCKS,
// both warpgroups have the SAME causal key range, so neither idles.
#define FH_QPB 2
// Depth of the cp.async pipeline, in key-block tiles.
//
// MEASURED NEGATIVE: 4 stages was neutral (6.13 ms against 6.07 for 2), even though the profile
// showed 57% of cycles with no eligible warp. Deeper prefetch does not help because the stall is
// not a shortage of outstanding LOADS — it is that the warpgroup waits on its own wgmma the
// instruction after committing it (see the overlap note at the end of this kernel). Kept at 2, so
// shared memory stays at 64 KB and leaves headroom for a larger tile later.
#define FH_NS 2

// GMMA matrix descriptor, canonical Major-K INTERLEAVE (no swizzle). From CuTe's make_gmma_desc,
// in units of uint128_t the operand must be ((8,n),2):((1,SBO),LBO) — i.e. 8 consecutive MN rows
// are ONE contiguous 128-byte core matrix, the next 8-row group is SBO away, and the next 16 bytes
// of K are LBO away. Packing core matrix (i,j) at (i*nj + j)*128 bytes therefore gives LBO = 128
// and SBO = nj*128, where nj is the operand's full K extent in 8-element units.
//
// No swizzle is needed BECAUSE of that shape: a core matrix is 128 contiguous bytes, so the eight
// rows the hardware gathers span all 32 banks exactly once. The +8 padding the mma.sync kernel
// needs for the same reason would actually break the descriptor, which addresses in 16-byte units.
__device__ __forceinline__ unsigned long long fh_desc(const void* smem, unsigned lbo, unsigned sbo) {
    unsigned s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    return ((unsigned long long)((s >> 4) & 0x3FFFu))
         | ((unsigned long long)((lbo >> 4) & 0x3FFFu) << 16)
         | ((unsigned long long)((sbo >> 4) & 0x3FFFu) << 32);   // layout 0, base_offset 0
}
// Advancing an operand is just arithmetic on the descriptor's low bits: the start address lives in
// bits [0,14) as bytes>>4, so +16 is +256 bytes = one k-step of 16 halves.
#define FH_ADV(d, bytes) ((d) + (unsigned long long)((bytes) >> 4))

#define FH_A16 "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15}"
#define FH_O16(c) "+f"((c)[0]),"+f"((c)[1]),"+f"((c)[2]),"+f"((c)[3]),"+f"((c)[4]),"+f"((c)[5]), \
                  "+f"((c)[6]),"+f"((c)[7]),"+f"((c)[8]),"+f"((c)[9]),"+f"((c)[10]),"+f"((c)[11]), \
                  "+f"((c)[12]),"+f"((c)[13]),"+f"((c)[14]),"+f"((c)[15])
// n=32 everywhere so one accumulator width covers both GEMMs: S is exactly BK=32 keys wide, and
// O's 128 dims are four of these. Both operands K-major, hence the trailing 0s.
#define FH_MMA_SS(c, da, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " FH_A16 \
                 ", %16, %17, 1, 1, 1, 0, 0;\n" : FH_O16(c) : "l"(da), "l"(db))
#define FH_MMA_RS(c, a0, a1, a2, a3, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " FH_A16 \
                 ", {%16,%17,%18,%19}, %20, 1, 1, 1, 0;\n" \
                 : FH_O16(c) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "l"(db))
#define FH_FENCE()  asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory")
#define FH_COMMIT() asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory")
// N = how many committed groups may still be in flight afterwards. Retirement is in commit
// order, so wait_group 1 with [PV, S] outstanding retires PV and leaves S flying.
#define FH_WAIT(N)  asm volatile("wgmma.wait_group.sync.aligned " #N ";\n" ::: "memory")


// One m64n128k16 per k-step instead of four m64n32k16: the accumulator layout of nN wgmma is the
// n32 layout repeated across nt groups, so &c[0][0]..&c[15][3] bind unchanged. Fewer, fatter async
// ops: 8 issues per S tile instead of 32, and the same for PV.
#define FH_A64 "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,%60,%61,%62,%63}"
#define FH_O64(c) \
    "+f"((c)[0]),"+f"((c)[1]),"+f"((c)[2]),"+f"((c)[3]),"+f"((c)[4]),"+f"((c)[5]),"+f"((c)[6]),"+f"((c)[7]), \
    "+f"((c)[8]),"+f"((c)[9]),"+f"((c)[10]),"+f"((c)[11]),"+f"((c)[12]),"+f"((c)[13]),"+f"((c)[14]),"+f"((c)[15]), \
    "+f"((c)[16]),"+f"((c)[17]),"+f"((c)[18]),"+f"((c)[19]),"+f"((c)[20]),"+f"((c)[21]),"+f"((c)[22]),"+f"((c)[23]), \
    "+f"((c)[24]),"+f"((c)[25]),"+f"((c)[26]),"+f"((c)[27]),"+f"((c)[28]),"+f"((c)[29]),"+f"((c)[30]),"+f"((c)[31]), \
    "+f"((c)[32]),"+f"((c)[33]),"+f"((c)[34]),"+f"((c)[35]),"+f"((c)[36]),"+f"((c)[37]),"+f"((c)[38]),"+f"((c)[39]), \
    "+f"((c)[40]),"+f"((c)[41]),"+f"((c)[42]),"+f"((c)[43]),"+f"((c)[44]),"+f"((c)[45]),"+f"((c)[46]),"+f"((c)[47]), \
    "+f"((c)[48]),"+f"((c)[49]),"+f"((c)[50]),"+f"((c)[51]),"+f"((c)[52]),"+f"((c)[53]),"+f"((c)[54]),"+f"((c)[55]), \
    "+f"((c)[56]),"+f"((c)[57]),"+f"((c)[58]),"+f"((c)[59]),"+f"((c)[60]),"+f"((c)[61]),"+f"((c)[62]),"+f"((c)[63])
#define FH_MMA128_SS(c, da, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " FH_A64 \
                 ", %64, %65, 1, 1, 1, 0, 0;\n" : FH_O64(c) : "l"(da), "l"(db))
#define FH_MMA128_RS(c, a0, a1, a2, a3, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " FH_A64 \
                 ", {%64,%65,%66,%67}, %68, 1, 1, 1, 0;\n" \
                 : FH_O64(c) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "l"(db))

// scatter one 16-byte chunk into the core-matrix-blocked layout: element (mn, k) of an [MN][KX]
// tile lands at brick (mn/8, k/8), row mn%8. KX8 is KX/8.
#define FH_OFF(mn, k8, KX8) (((long)((mn) >> 3) * (KX8) + (k8)) * 64 + ((mn) & 7) * 8)

// ---- WIDE-M Q4 GEMM ON WGMMA ----
// C[M, R] = X[M, K] * W^T with W in q4_0. The v2 kernel's per-block dequant ALU stops
// amortising past its register tile (plateau 56-104 TF/s vs cuBLAS 469 at M=2000); here the
// producer traffic is the RAW q4 bytes (4.5 bits/weight) and the dequant happens once per
// stage into a bf16 smem tile the tensor cores then read — dense math on quantised memory.
// bf16 operands because the activations arrive as bf16 (hb16); accumulator layout is the
// n32 wgmma layout repeated 16x, identical binding to FH_O64.
#define FH_MMA128BF_SS(c, da, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " FH_A64 \
                 ", %64, %65, 1, 1, 1, 0, 0;\n" : FH_O64(c) : "l"(da), "l"(db))
#define GW_BM 128
#define GW_BN 128
#define GW_BK 64
// dynamic smem: double-buffered A + B f16 tiles, raw-q4 B staging, scales
#define GW_SMEM ((2 * (GW_BM * GW_BK + GW_BN * GW_BK)) * 2 + 2 * (GW_BN * 2 * 16) + 2 * (GW_BN * 2 * 2) + 256)

// 4-byte cp.async with src-size zero-fill (16-byte helpers are above; scales need 4B)
__device__ __forceinline__ void gw_cp4z(void* smem, const void* gmem, unsigned sz) {
    unsigned int s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n" :: "r"(s), "l"(gmem), "r"(sz));
}

// ---- MARLIN-STRUCTURE Q4 GEMM (gemm_q4_ml) ----
// The wgmma kernel above tops out because wgmma REQUIRES both operands in shared memory: the
// weights must be dequantised INTO smem and read back, and the resulting smem footprint caps the
// pipeline at 2 stages. ncu on it: 18% tensor-active, 24% occupancy, long_scoreboard 46k
// inst/warp — global-load latency with nothing to hide it behind.
//
// This kernel takes the published Marlin shape instead: mma.sync (register operands) so the
// weights are dequantised straight into REGISTERS and never round-trip through smem, which frees
// enough smem for a FOUR-stage cp.async pipeline over the raw q4 bytes. Only A is staged in smem
// (via ldmatrix), and only the raw 4.5-bit weights cross the memory system.
//
// Tiling: BM=64 (all of a decode batch), BN=128 output columns, BK=64. 4 warps, each owning 32 of
// the 128 columns: per warp 4 m-tiles x 4 n-tiles x 4 f32 = 64 accumulators.
#define ML_BM 64
#define ML_BN 128
#define ML_BK 64
#define ML_ST 4                                  // pipeline depth (3 measured a wash)
// smem: A[ML_ST][BM][BK] f16 (4*64*64*2 = 32 KB) + raw q4 B[ML_ST][BN][BK/2 nibbles] (4*128*32 =
// 16 KB) + scales (4*128*2*2) = ~50 KB — three blocks per SM at 228 KB.
#define ML_SMEM (ML_ST * ML_BM * ML_BK * 2 + ML_ST * ML_BN * (ML_BK / 2) + ML_ST * ML_BN * 2 * 2 + 256)

template<int BN, int NT>
__device__ __forceinline__ void ml_body(
        const unsigned short* __restrict__ scales,
        const uint4* __restrict__ quants,
        const unsigned short* __restrict__ xb,     // [M, K] f16 row-major
        float* __restrict__ y,                     // [M, R] f32 row-major
        const int m, const int rows, const int cols, const int nblk,
        const int ksplit, float* __restrict__ part) {
    const int r0 = blockIdx.x * BN;
    const int m0 = blockIdx.y * ML_BM;
    const int tid = threadIdx.x;                   // 128 threads = 4 warps
    const int warp = tid >> 5, lane = tid & 31;
    const int nw0 = warp * (BN / 4);                     // this warp's 32 output columns
    // ldmatrix lane mapping, identical to the attention kernel's proven one
    const int lrow = 8 * ((lane >> 3) & 1) + (lane & 7);
    const int lcol = 8 * (lane >> 4);

    extern __shared__ unsigned short mlsh[];
    unsigned short* base = (unsigned short*)(((unsigned long long)mlsh + 127) & ~127ULL);
    // A tiles: [stage][64 rows][64 k] with +8 row padding for conflict-free ldmatrix
    const int ARS = ML_BK + 8;
    unsigned short* sA = base;                                        // ML_ST*64*(64+8)
    // QRS = 12, not 8: an 8-word row stride put the 8 distinct ncol of a lane-quad on banks
    // 0,8,16,24,0,8,16,24 — a 4-way conflict (ncu counted 311k). 12 words rotates them to
    // 0,12,24,4,16,28,8,20, all distinct, AND stays 16-byte aligned so the cp.async 16 B
    // stores remain legal (a 9-word stride is conflict-free too but misaligns them).
    const int QRS = 12;
    unsigned int* sQ = (unsigned int*)(base + ML_ST * ML_BM * ARS);   // raw q4: [stage][BN][9 u32]
    unsigned short* sS = (unsigned short*)(sQ + ML_ST * BN * QRS);    // scales: [stage][BN][2]

    float acc[4 * NT * 4];
    #pragma unroll
    for (int i = 0; i < 4 * NT * 4; ++i) acc[i] = 0.f;

    // SPLIT-K: BN=128 keeps the 4:16 ldmatrix:mma amortisation (BN=32 threw it away for grid
    // fill and lost 128/256), so instead we slice K across gridDim.z to fill the machine —
    // 48 n-blocks x 4 slices = 192 blocks at decode width, same wide tile.
    const int nk_all = cols / ML_BK;
    const int zs = (int)blockIdx.z;
    const int per_z = (nk_all + ksplit - 1) / ksplit;
    const int k_lo = zs * per_z;
    const int k_hi = min(nk_all, k_lo + per_z);
    const int nk = k_hi - k_lo;
    if (nk <= 0) return;

    #define ML_STAGE(st, ks) do {                                                          \
        const int _k0 = (ks) * ML_BK;                                                      \
        /* A: 64 rows x 64 halves = 8 uint4 per row; 128 threads -> 4 rows each */         \
        for (int c = tid; c < ML_BM * 8; c += 128) {                                       \
            const int mn = c >> 3, k8 = c & 7;                                             \
            const int gr = m0 + mn;                                                        \
            fa_cp16z(sA + (long)(st) * ML_BM * ARS + mn * ARS + k8 * 8,                    \
                     xb + (long)gr * cols + _k0 + k8 * 8, gr < m ? 16u : 0u);              \
        }                                                                                  \
        /* B raw: 128 rows x 2 q4 blocks = 2 uint4 per row */                              \
        for (int c = tid; c < BN * 2; c += 128) {                                       \
            const int rn = c >> 1, kb = c & 1;                                             \
            const int gr = r0 + rn;                                                        \
            fa_cp16z(sQ + (long)(st) * BN * QRS + rn * QRS + kb * 4,                    \
                     (const void*)(quants + (long)gr * nblk + (_k0 >> 5) + kb),            \
                     gr < rows ? 16u : 0u);                                                \
        }                                                                                  \
        for (int c = tid; c < BN; c += 128) {                                           \
            const int gr = r0 + c;                                                         \
            gw_cp4z(sS + (long)(st) * BN * 2 + c * 2,                                   \
                    scales + (long)gr * nblk + (_k0 >> 5), gr < rows ? 4u : 0u);           \
        }                                                                                  \
    } while (0)

    // prologue: fill ML_ST-1 stages
    #pragma unroll
    for (int st = 0; st < ML_ST - 1; ++st) {
        if (st < nk) ML_STAGE(st, k_lo + st);
        FA_CP_COMMIT();
    }

    for (int ks = 0; ks < nk; ++ks) {
        const int st = ks % ML_ST;
        // keep ML_ST-1 groups in flight
        asm volatile("cp.async.wait_group %0;\n" :: "n"(ML_ST - 2));
        __syncthreads();

        const unsigned short* As = sA + (long)st * ML_BM * ARS;
        const unsigned int* Qs = sQ + (long)st * BN * QRS;
        const unsigned short* Ss = sS + (long)st * BN * 2;

        // B fragments IN REGISTERS: each lane owns rows (nw0 + (lane>>2)*? ) per the mma .col
        // B layout — for m16n8k16 .row.col, B is [n][k]: lane l supplies n = l>>2, k pairs from
        // (l&3)*2. Two k-groups (k0..15) per mma; we run 4 k-steps of 16 across BK=64.
        #pragma unroll
        for (int kt = 0; kt < ML_BK / 16; ++kt) {
            // A fragments for the 4 m-tiles at this k-step
            unsigned int a[4][4];
            #pragma unroll
            for (int mt = 0; mt < 4; ++mt)
                fa_ldm4(a[mt][0], a[mt][1], a[mt][2], a[mt][3],
                        As + (long)(mt * 16 + lrow) * ARS + kt * 16 + lcol);
            #pragma unroll
            for (int nt = 0; nt < NT; ++nt) {
                // mma.m16n8k16 B fragment (PTX): lane l supplies column n = l>>2 and rows
                // k = (l&3)*2 + {0,1} in reg0 and the SAME +8 in reg1. Both halves of a pair sit
                // 8 bits apart in one q4 word, so the magic-bias extraction is one mask+or.
                const int ncol = nw0 + nt * 8 + (lane >> 2);
                const int e_lo = kt * 16 + (lane & 3) * 2;
                const int kb   = e_lo >> 5;                       // which 32-wide q4 block
                const float sc = h2f(Ss[ncol * 2 + kb]);
                const unsigned short sh = f2h(sc), bh = f2h(-1032.0f * sc);
                const unsigned int sc2 = (unsigned)sh | ((unsigned)sh << 16);
                const unsigned int bs2 = (unsigned)bh | ((unsigned)bh << 16);
                unsigned int pk[2];
                #pragma unroll
                for (int h2 = 0; h2 < 2; ++h2) {
                    const int ee = (e_lo + 8 * h2) & 31;          // element within the block
                    const unsigned int qw = Qs[ncol * QRS + kb * 4 + ((ee & 15) >> 2)];
                    const int sft = (ee & 3) * 8 + ((ee >= 16) ? 4 : 0);
                    const unsigned int v = qw >> sft;
                    pk[h2] = ((v & 0x0000000Fu) | ((v & 0x00000F00u) << 8)) | 0x64006400u;
                }
                unsigned int b0, b1;
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(b0) : "r"(pk[0]), "r"(sc2), "r"(bs2));
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(b1) : "r"(pk[1]), "r"(sc2), "r"(bs2));
                #pragma unroll
                for (int mt = 0; mt < 4; ++mt) {
                    float* d = &acc[(mt * NT + nt) * 4];
                    FA_MMA(d[0], d[1], d[2], d[3], a[mt][0], a[mt][1], a[mt][2], a[mt][3], b0, b1);
                }
            }
        }
        __syncthreads();
        // RING SLOT = (k index) % depth, NOT the slot just consumed: with depth 4 and a
        // 3-stage prologue (k=0,1,2 in slots 0,1,2), prefetching k=ks+3 into slot `st` left
        // slot 3 never written — every 4th k-stage multiplied garbage.
        if (ks + ML_ST - 1 < nk) ML_STAGE((ks + ML_ST - 1) % ML_ST, k_lo + ks + ML_ST - 1);
        FA_CP_COMMIT();
    }
    #undef ML_STAGE
    asm volatile("cp.async.wait_group 0;\n" ::);

    // epilogue: m16n8k16 C layout — lane l holds rows (l>>2) and (l>>2)+8, cols (l&3)*2 (+1)
    #pragma unroll
    for (int mt = 0; mt < 4; ++mt) {
        const int ra = m0 + mt * 16 + (lane >> 2), rb = ra + 8;
        #pragma unroll
        for (int nt = 0; nt < NT; ++nt) {
            const int col = r0 + nw0 + nt * 8 + (lane & 3) * 2;
            const float* d = &acc[(mt * NT + nt) * 4];
            float* out = (ksplit > 1) ? (part + (long)zs * m * rows) : y;
            if (col < rows) {
                if (ra < m) {
                    out[(long)ra * rows + col] = d[0];
                    if (col + 1 < rows) out[(long)ra * rows + col + 1] = d[1];
                }
                if (rb < m) {
                    out[(long)rb * rows + col] = d[2];
                    if (col + 1 < rows) out[(long)rb * rows + col + 1] = d[3];
                }
            }
        }
    }
}

#define ML_WRAP(NAME, BNV, NTV) \
extern "C" __global__ __launch_bounds__(128, 3) void NAME( \
        const unsigned short* __restrict__ scales, const uint4* __restrict__ quants, \
        const unsigned short* __restrict__ xb, float* __restrict__ y, \
        const int m, const int rows, const int cols, const int nblk, \
        const int ksplit, float* __restrict__ part) { \
    ml_body<BNV, NTV>(scales, quants, xb, y, m, rows, cols, nblk, ksplit, part); \
}
// Three tile widths: decode grids need many narrow blocks to fill 114 SMs (BN=128 gives only
// rows/128 = 48 blocks at Qwen's gate|up), prefill wants the wide tile for A-reuse.
ML_WRAP(gemm_q4_ml, 128, 4)
ML_WRAP(gemm_q4_ml64, 64, 2)
ML_WRAP(gemm_q4_ml32, 32, 1)

// deterministic split-K reduction: fixed slice order, one pass over [ksplit][M][rows]
extern "C" __global__ void gw_ksum(const float* __restrict__ part, float* __restrict__ y,
                                   const int n, const int ksplit) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float a = 0.f;
    for (int z = 0; z < ksplit; ++z) a += part[(long)z * n + i];
    y[i] = a;
}

// n64 wgmma for the GEMM's BN=64 tile: 32 f32 accumulators per thread (the n32 layout twice)
#define GW_A32 "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}"
#define GW_O32(c) \
    "+f"((c)[0]),"+f"((c)[1]),"+f"((c)[2]),"+f"((c)[3]),"+f"((c)[4]),"+f"((c)[5]),"+f"((c)[6]),"+f"((c)[7]), \
    "+f"((c)[8]),"+f"((c)[9]),"+f"((c)[10]),"+f"((c)[11]),"+f"((c)[12]),"+f"((c)[13]),"+f"((c)[14]),"+f"((c)[15]), \
    "+f"((c)[16]),"+f"((c)[17]),"+f"((c)[18]),"+f"((c)[19]),"+f"((c)[20]),"+f"((c)[21]),"+f"((c)[22]),"+f"((c)[23]), \
    "+f"((c)[24]),"+f"((c)[25]),"+f"((c)[26]),"+f"((c)[27]),"+f"((c)[28]),"+f"((c)[29]),"+f"((c)[30]),"+f"((c)[31])
#define GW_MMA64_SS(c, da, db) \
    asm volatile("wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " GW_A32 \
                 ", %32, %33, 1, 1, 1, 0, 0;\n" : GW_O32(c) : "l"(da), "l"(db))


// Repack q4_0 quants for the wgmma GEMM: arrange nibbles so each u32 yields four f16x2 pairs
// with ONE shift+and+or each (p01 = (v & 0x000F000F)|magic, p23 = ((v>>4)&...)|magic, ...).
// Repacked word j of a block holds elements 8j..8j+7 at nibble positions {0,4,1,5,2,6,3,7}.
extern "C" __global__ void q4_repack_wg(const uint4* __restrict__ src, uint4* __restrict__ dst,
                                        const long nblocks) {
    const long b = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= nblocks) return;
    const uint4 q = src[b];
    const unsigned int w[4] = { q.x, q.y, q.z, q.w };
    unsigned int o[4];
    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        const int pos[8] = { 0, 4, 1, 5, 2, 6, 3, 7 };
        unsigned int v = 0u;
        #pragma unroll
        for (int t = 0; t < 8; ++t) {
            const int E = 8 * j + t;
            const int sw = (E < 16) ? (E >> 2) : ((E - 16) >> 2);
            const int sn = (E < 16) ? (E & 3) : (4 + ((E - 16) & 3));
            v |= ((w[sw] >> (4 * sn)) & 0xFu) << (4 * pos[t]);
        }
        o[j] = v;
    }
    dst[b] = make_uint4(o[0], o[1], o[2], o[3]);
}

// Two warpgroups (256 threads): each owns 64 M-rows; BOTH consume the SAME B tile, so the
// dequant cost per FLOP is 4x lower than the 64-row single-warpgroup version. Double-buffered:
// stage k+1's generic stores run while stage k's wgmmas are in flight.
template<int BN, int NACC>
__device__ __forceinline__ void gw_body(
        const unsigned short* __restrict__ scales,
        const uint4* __restrict__ quants,
        const unsigned short* __restrict__ xb,     // [M, K] bf16 row-major
        float* __restrict__ y,                     // [M, R] f32 row-major
        const int m, const int rows, const int cols, const int nblk,
        const int a_is_f16,                        // activations already f16: A becomes a pure copy
        const int repacked,                        // quants pre-permuted by q4_repack_wg
        const int ksplit,                          // K-slices across gridDim.z (1 = off)
        float* __restrict__ part,                  // [ksplit][M][rows] when ksplit > 1
        unsigned short* __restrict__ dbg) {        // non-null: block(0,0) dumps its k0=0 tiles
    const int r0 = blockIdx.x * BN;
    const int m0 = blockIdx.y * GW_BM;
    const int tid = threadIdx.x;
    const int wg = tid >> 7;                       // warpgroup 0/1 -> m-rows [wg*64, wg*64+64)
    const int wtid = tid & 127;
    const int lane = tid & 31, w = wtid >> 5;
    const int gid = lane >> 2, t4 = lane & 3;

    extern __shared__ unsigned short gwsh[];
    unsigned short* base = (unsigned short*)(((unsigned long long)gwsh + 127) & ~127ULL);
    const int ASZ = GW_BM * GW_BK, BSZ = BN * GW_BK;
    unsigned short* sA[2] = { base, base + ASZ + BSZ };
    unsigned short* sB[2] = { base + ASZ, base + ASZ + BSZ + ASZ };
    // raw-q4 staging (cp.async targets) after the f16 tiles: [BN rows][2 blocks] uint4 + scales
    unsigned short* rawbase = base + 2 * (ASZ + BSZ);
    unsigned short* rawB[2] = { rawbase, rawbase + BN * 2 * 8 };
    unsigned short* rawS[2] = { rawbase + 2 * BN * 2 * 8, rawbase + 2 * BN * 2 * 8 + BN * 2 };

    float acc[NACC];
    #pragma unroll
    for (int i = 0; i < NACC; ++i) acc[i] = 0.f;

    const int KX8 = GW_BK / 8;
    // SPLIT-K: at decode widths the (rows/BN) grid is well under one wave and each block walks
    // all K stages serially — ncu long_scoreboard-dominated at 12% occupancy. Slicing K across
    // gridDim.z multiplies resident blocks by ksplit; partials are summed by gw_ksum in a fixed
    // order, so the result stays deterministic.
    const int nk_all = cols / GW_BK;
    const int zslice = (int)blockIdx.z;
    const int nk_per = (nk_all + ksplit - 1) / ksplit;
    const int ks_lo = zslice * nk_per;
    const int ks_hi = min(nk_all, ks_lo + nk_per);
    const int nk = ks_hi - ks_lo;
    if (nk <= 0) return;

    // stage `ks` into buffer `buf` — all 256 threads cooperate
    #define GW_STAGE(buf, ks) do {                                                        \
        const int _k0 = (ks) * GW_BK;                                                     \
        for (int c = tid; c < GW_BM * KX8; c += 256) {                                    \
            const int mn = c / KX8, k8 = c - mn * KX8;                                    \
            const int gr = m0 + mn;                                                       \
            uint4 v = make_uint4(0u, 0u, 0u, 0u);                                         \
            if (gr < m) v = *(const uint4*)(xb + (long)gr * cols + _k0 + k8 * 8);         \
            if (!a_is_f16) {                                                              \
                /* bf16 activations: convert to f16 for the f16 wgmma. Production feeds  */ \
                /* f16 directly (norm kernels emit it) and skips this entirely.          */ \
                unsigned int* vw = (unsigned int*)&v;                                     \
                _Pragma("unroll")                                                         \
                for (int i2 = 0; i2 < 4; ++i2) {                                          \
                    const unsigned int b2 = vw[i2];                                       \
                    const unsigned short a0 = f2h(bf2f((unsigned short)(b2 & 0xFFFFu))); \
                    const unsigned short a1 = f2h(bf2f((unsigned short)(b2 >> 16)));     \
                    vw[i2] = (unsigned)a0 | ((unsigned)a1 << 16);                         \
                }                                                                         \
            }                                                                             \
            *(uint4*)(sA[buf] + FH_OFF(mn, k8, KX8)) = v;                                 \
        }                                                                                 \
        for (int c = tid; c < BN * (GW_BK / 32); c += 256) {                           \
            const int rn = c / (GW_BK / 32), kb = c % (GW_BK / 32);                       \
            const int gr = r0 + rn;                                                       \
            uint4 q = make_uint4(0u, 0u, 0u, 0u);                                         \
            unsigned int sc2 = 0u, bs2 = 0u;                                              \
            if (gr < rows) {                                                              \
                const long blk = (long)gr * nblk + (_k0 >> 5) + kb;                       \
                q = quants[blk];                                                          \
                const float s = h2f(scales[blk]);                                         \
                /* fma2 constants: out = magic*s + (-1032*s), magic = f16(1024+n) */      \
                const unsigned short sh = f2h(s), bh = f2h(-1032.0f * s);                 \
                sc2 = (unsigned)sh | ((unsigned)sh << 16);                                \
                bs2 = (unsigned)bh | ((unsigned)bh << 16);                                \
            }                                                                             \
            unsigned int wq[4] = { q.x, q.y, q.z, q.w };                                  \
            unsigned short tmp[32];                                                       \
            _Pragma("unroll")                                                             \
            for (int i = 0; i < 4; ++i) {                                                 \
                const unsigned int v = wq[i];                                             \
                /* nibble n -> f16(1024+n) by OR into 0x6400's mantissa; one fma2 then */ \
                /* applies scale and the -(1024+8)*s bias: two elements per op.       */  \
                unsigned int p01, p23, h01, h23;                                          \
                if (repacked) {                                                           \
                    /* word i holds elems 8i..8i+7 at nibbles {0,4,1,5,2,6,3,7}: one      */ \
                    /* mask+or per pair. tmp targets: elems 8i.. -> in-block K order.     */ \
                    p01 = (v & 0x000F000Fu) | 0x64006400u;                                \
                    p23 = ((v >> 4)  & 0x000F000Fu) | 0x64006400u;                        \
                    h01 = ((v >> 8)  & 0x000F000Fu) | 0x64006400u;                        \
                    h23 = ((v >> 12) & 0x000F000Fu) | 0x64006400u;                        \
                } else {                                                                  \
                    p01 = ((v & 0x0000000Fu) | ((v & 0x00000F00u) << 8)) | 0x64006400u;   \
                    p23 = (((v >> 16) & 0x0000000Fu) | (((v >> 16) & 0x00000F00u) << 8)) | 0x64006400u; \
                    h01 = (((v >> 4) & 0x0000000Fu) | (((v >> 4) & 0x00000F00u) << 8)) | 0x64006400u; \
                    h23 = (((v >> 20) & 0x0000000Fu) | (((v >> 20) & 0x00000F00u) << 8)) | 0x64006400u; \
                }                                                                         \
                unsigned int r01, r23, r45, r67;                                          \
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r01) : "r"(p01), "r"(sc2), "r"(bs2)); \
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r23) : "r"(p23), "r"(sc2), "r"(bs2)); \
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r45) : "r"(h01), "r"(sc2), "r"(bs2)); \
                asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r67) : "r"(h23), "r"(sc2), "r"(bs2)); \
                if (repacked) {                                                           \
                    /* word i = elems 8i..8i+7 in K order */                              \
                    *(unsigned int*)(tmp + i * 8 + 0) = r01;                              \
                    *(unsigned int*)(tmp + i * 8 + 2) = r23;                              \
                    *(unsigned int*)(tmp + i * 8 + 4) = r45;                              \
                    *(unsigned int*)(tmp + i * 8 + 6) = r67;                              \
                } else {                                                                  \
                    *(unsigned int*)(tmp + i * 4 + 0)  = r01;                             \
                    *(unsigned int*)(tmp + i * 4 + 2)  = r23;                             \
                    *(unsigned int*)(tmp + i * 4 + 16) = r45;                             \
                    *(unsigned int*)(tmp + i * 4 + 18) = r67;                             \
                }                                                                         \
            }                                                                             \
            _Pragma("unroll")                                                             \
            for (int h = 0; h < 4; ++h)                                                   \
                *(uint4*)(sB[buf] + FH_OFF(rn, kb * 4 + h, KX8)) = *(const uint4*)(tmp + h * 8); \
        }                                                                                 \
    } while (0)

    if (a_is_f16) {
        // v7: cp.async pipeline. A (already f16) streams STRAIGHT to its core-matrix smem home
        // with zero register involvement; B's raw q4 + scales cp.async to a small staging area
        // and dequantise from SMEM — the v3 profile was global-load LATENCY with too few warps
        // to hide it, and cp.async removes the register dependency chains entirely.
        #define GW_CP_STAGE(buf, ks) do {                                                 \
            const int _k0 = (ks) * GW_BK;                                                 \
            for (int c = tid; c < GW_BM * KX8; c += 256) {                                \
                const int mn = c / KX8, k8 = c - mn * KX8;                                \
                const int gr = m0 + mn;                                                   \
                fa_cp16z(sA[buf] + FH_OFF(mn, k8, KX8),                                   \
                         xb + (long)gr * cols + _k0 + k8 * 8, gr < m ? 16u : 0u);         \
            }                                                                             \
            for (int c = tid; c < BN * 2; c += 256) {                                  \
                const int rn = c >> 1, kb = c & 1;                                        \
                const int gr = r0 + rn;                                                   \
                const long blk = (long)gr * nblk + (_k0 >> 5) + kb;                       \
                fa_cp16z(rawB[buf] + c * 8, (const void*)(quants + blk),                  \
                         gr < rows ? 16u : 0u);                                           \
            }                                                                             \
            for (int c = tid; c < BN; c += 256) {                                      \
                const int gr = r0 + c;                                                    \
                gw_cp4z(rawS[buf] + c * 2, scales + (long)gr * nblk + (_k0 >> 5),         \
                        gr < rows ? 4u : 0u);                                             \
            }                                                                             \
        } while (0)
        GW_CP_STAGE(0, ks_lo);
        FA_CP_COMMIT();
        for (int ks = 0; ks < nk; ++ks) {
            const int buf = ks & 1;
            if (ks + 1 < nk) {
                // the cp target buf^1 was read by wgmma group ks-1 — retire it (depth 1: the
                // group just committed for ks-1... at loop top the outstanding wgmma groups are
                // {ks-1}; ncu showed WAIT(0)-per-stage exposing the whole wgmma latency at 12%
                // occupancy, so the wait moves HERE, guarding only actual buffer reuse.
                if (ks >= 1) { FH_WAIT(0); }
                GW_CP_STAGE(buf ^ 1, ks_lo + ks + 1);
                FA_CP_COMMIT();
                FA_CP_WAIT1();                 // stage ks (older cp group) has landed
            } else {
                if (ks >= 1) { FH_WAIT(0); }
                asm volatile("cp.async.wait_group 0;\n" ::);
            }
            __syncthreads();
            // dequant stage ks from SMEM raw -> f16 tile (magic-bias, one fma2 per pair)
            for (int c = tid; c < BN * 2; c += 256) {
                const int rn = c >> 1, kb = c & 1;
                const uint4 q = *(const uint4*)(rawB[buf] + c * 8);
                const float s = h2f(rawS[buf][rn * 2 + kb]);
                const unsigned short sh = f2h(s), bh = f2h(-1032.0f * s);
                const unsigned int sc2 = (unsigned)sh | ((unsigned)sh << 16);
                const unsigned int bs2 = (unsigned)bh | ((unsigned)bh << 16);
                const unsigned int wq[4] = { q.x, q.y, q.z, q.w };
                unsigned short tmp[32];
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const unsigned int v = wq[i];
                    const unsigned int p01 = ((v & 0x0000000Fu) | ((v & 0x00000F00u) << 8)) | 0x64006400u;
                    const unsigned int p23 = (((v >> 16) & 0x0000000Fu) | (((v >> 16) & 0x00000F00u) << 8)) | 0x64006400u;
                    const unsigned int h01 = (((v >> 4) & 0x0000000Fu) | (((v >> 4) & 0x00000F00u) << 8)) | 0x64006400u;
                    const unsigned int h23 = (((v >> 20) & 0x0000000Fu) | (((v >> 20) & 0x00000F00u) << 8)) | 0x64006400u;
                    unsigned int r01, r23, r45, r67;
                    asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r01) : "r"(p01), "r"(sc2), "r"(bs2));
                    asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r23) : "r"(p23), "r"(sc2), "r"(bs2));
                    asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r45) : "r"(h01), "r"(sc2), "r"(bs2));
                    asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(r67) : "r"(h23), "r"(sc2), "r"(bs2));
                    *(unsigned int*)(tmp + i * 4 + 0)  = r01;
                    *(unsigned int*)(tmp + i * 4 + 2)  = r23;
                    *(unsigned int*)(tmp + i * 4 + 16) = r45;
                    *(unsigned int*)(tmp + i * 4 + 18) = r67;
                }
                #pragma unroll
                for (int h = 0; h < 4; ++h)
                    *(uint4*)(sB[buf] + FH_OFF(rn, kb * 4 + h, KX8)) = *(const uint4*)(tmp + h * 8);
            }
            __syncthreads();
            asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
            const unsigned long long dA = fh_desc(sA[buf] + (long)wg * 64 * GW_BK, 128,
                                                  (unsigned)(KX8 * 128));
            const unsigned long long dB = fh_desc(sB[buf], 128, (unsigned)(KX8 * 128));
            FH_FENCE();
            #pragma unroll
            for (int s = 0; s < GW_BK / 16; ++s)
                { if (BN == 128) FH_MMA128_SS(acc, FH_ADV(dA, s * 256), FH_ADV(dB, s * 256));
                    else            GW_MMA64_SS(acc, FH_ADV(dA, s * 256), FH_ADV(dB, s * 256)); }
            FH_COMMIT();
            // no wait here: the group retires at the NEXT iteration's loop-top guard, so the
            // whole cp+dequant of stage ks+1 overlaps this stage's wgmmas
        }
        FH_WAIT(0);                            // accumulators are async-written registers
        #undef GW_CP_STAGE
        goto epilogue;
    }

    GW_STAGE(0, ks_lo);
    __syncthreads();
    asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");

    for (int ks = 0; ks < nk; ++ks) {
        const int buf = ks & 1;
        // this warpgroup's A sub-tile: rows wg*64..wg*64+63 = brick rows wg*8.. -> +wg*8*KX8*128 B
        const unsigned long long dA = fh_desc(sA[buf] + (long)wg * 64 * GW_BK, 128,
                                              (unsigned)(KX8 * 128));
        const unsigned long long dB = fh_desc(sB[buf], 128, (unsigned)(KX8 * 128));
        FH_FENCE();
        #pragma unroll
        for (int s = 0; s < GW_BK / 16; ++s)
            { if (BN == 128) FH_MMA128_SS(acc, FH_ADV(dA, s * 256), FH_ADV(dB, s * 256));
                    else            GW_MMA64_SS(acc, FH_ADV(dA, s * 256), FH_ADV(dB, s * 256)); }
        FH_COMMIT();
        if (dbg && ks == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
            for (int c = tid; c < GW_BM * GW_BK; c += 256) {
                const int mn = c / GW_BK, k = c - mn * GW_BK;
                dbg[c] = sA[0][FH_OFF(mn, k >> 3, KX8) + (k & 7)];
            }
            for (int c = tid; c < BN * GW_BK; c += 256) {
                const int mn = c / GW_BK, k = c - mn * GW_BK;
                dbg[GW_BM * GW_BK + c] = sB[0][FH_OFF(mn, k >> 3, KX8) + (k & 7)];
            }
        }
        if (ks + 1 < nk) {
            // stage the NEXT tile into the other buffer while this one's wgmmas fly
            GW_STAGE(buf ^ 1, ks_lo + ks + 1);
        }
        FH_WAIT(0);
        __syncthreads();
        asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
    }
    #undef GW_STAGE

    epilogue:
    // epilogue: rows m0 + wg*64 + w*16 + gid (+8); cols r0 + nt*8 + t4*2 (+1)
    const int ra = m0 + wg * 64 + w * 16 + gid, rb = ra + 8;
    #pragma unroll
    for (int nt = 0; nt < NACC / 4; ++nt) {
        const int col = r0 + nt * 8 + t4 * 2;
        if (col >= rows) continue;
        float* out = (ksplit > 1) ? (part + (long)zslice * m * rows) : y;
        if (ra < m) {
            out[(long)ra * rows + col] = acc[nt * 4 + 0];
            if (col + 1 < rows) out[(long)ra * rows + col + 1] = acc[nt * 4 + 1];
        }
        if (rb < m) {
            out[(long)rb * rows + col] = acc[nt * 4 + 2];
            if (col + 1 < rows) out[(long)rb * rows + col + 1] = acc[nt * 4 + 3];
        }
    }
}

extern "C" __global__ __launch_bounds__(256, 1) void gemm_q4_wg(
        const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
        const unsigned short* __restrict__ xb, float* __restrict__ y,
        const int m, const int rows, const int cols, const int nblk,
        const int a_is_f16, const int repacked, const int ksplit,
        float* __restrict__ part, unsigned short* __restrict__ dbg) {
    gw_body<128, 64>(scales, quants, xb, y, m, rows, cols, nblk, a_is_f16, repacked, ksplit, part, dbg);
}
// BN=64 twin: at decode widths the BN=128 grid is only rows/128 blocks (48 for Qwen gate|up)
// against 114 SMs — HALF THE GPU IDLE. Narrower tiles double the block count; measured
// +21% at M=64 while losing at M=2000, so the host dispatches by width.
extern "C" __global__ __launch_bounds__(256, 1) void gemm_q4_wg64(
        const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
        const unsigned short* __restrict__ xb, float* __restrict__ y,
        const int m, const int rows, const int cols, const int nblk,
        const int a_is_f16, const int repacked, const int ksplit,
        float* __restrict__ part, unsigned short* __restrict__ dbg) {
    gw_body<64, 32>(scales, quants, xb, y, m, rows, cols, nblk, a_is_f16, repacked, ksplit, part, dbg);
}

extern "C" __global__ __launch_bounds__(128 * FH_QPB, 2) void fa_prefill_hop(
                                      const unsigned short* __restrict__ qkv16,
                                      const unsigned short* __restrict__ vt16,
                                      unsigned short* __restrict__ out,
                                      const int T, const int Tp, const int n_heads, const int n_kv,
                                      const float scale) {
    const int qb = blockIdx.x;
    const int wg = threadIdx.x >> 7;                 // which query head of the pair
    const int h = blockIdx.y * FH_QPB + wg;
    const int per = n_heads + 2 * n_kv;
    // The host only takes this path when n_heads/n_kv is even, which is what makes the two heads
    // of a block land on the SAME kv head — the shared tile below would otherwise be wrong.
    const int kvh = h / (n_heads / n_kv);
    const int tid = threadIdx.x & 127;               // lane within this warpgroup
    const int lane = threadIdx.x & 31, w = tid >> 5;
    const int gid = lane >> 2, t4 = lane & 3;
    const int q0 = qb * FH_BQ;

    extern __shared__ unsigned short fhsh[];
    // descriptors address in 16-byte units, and a core matrix is 128 B; align the base so every
    // brick is a whole number of them.
    unsigned short* base = (unsigned short*)(((unsigned long long)fhsh + 127ull) & ~127ull);
    unsigned short* sQall = base;                           // [FH_QPB][64][128] blocked
    unsigned short* sQ = sQall + wg * (FH_BQ * FH_HD);      // this warpgroup's head
    unsigned short* sKb = sQall + FH_QPB * FH_BQ * FH_HD;   // [FH_NS][32][128] blocked, SHARED
    unsigned short* sVb = sKb + FH_NS * FH_BK * FH_HD;      // [FH_NS][128][32] blocked (V^T)
    // NOTE: FH_SHARED above must match FH_NS.

    // Q, pre-scaled by 1/sqrt(head_dim) once as before, straight into the blocked layout. Each
    // warpgroup stages its own head, so this loop is over 128 threads, not the whole block.
    for (int i = tid; i < FH_BQ * (FH_HD / 8); i += 128) {
        const int r = i / (FH_HD / 8), c8 = i - r * (FH_HD / 8);
        const uint4 g = *(const uint4*)(qkv16 + (long)(q0 + r) * per * FH_HD + (long)h * FH_HD + c8 * 8);
        const unsigned int* gw = (const unsigned int*)&g;
        unsigned int sw[4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const float lo = h2f((unsigned short)(gw[j] & 0xFFFFu)) * scale;
            const float hi = h2f((unsigned short)(gw[j] >> 16)) * scale;
            sw[j] = fa_pack2(lo, hi);
        }
        *(uint4*)(sQ + FH_OFF(r, c8, FH_HD / 8)) = *(const uint4*)sw;
    }

    float o[FH_HD / 8][4];
    #pragma unroll
    for (int nt = 0; nt < FH_HD / 8; ++nt) { o[nt][0] = 0.f; o[nt][1] = 0.f; o[nt][2] = 0.f; o[nt][3] = 0.f; }
    float m0 = FA_MASKV, m1 = FA_MASKV, l0 = 0.f, l1 = 0.f;

    const int row0 = q0 + 16 * w + gid;
    const int row1 = row0 + 8;
    const int kb_max = (q0 + FH_BQ - 1) / FH_BK;
    const int kb_lo = 0, kb_hi = kb_max + 1;
    // Block-uniform kv head. The staging loop below is indexed by threadIdx.x across all 256
    // threads, so it must not use the per-warpgroup `kvh` (identical in value, but only because
    // of the host's even-ratio gate — spell it out rather than rely on that here).
    const int kvh0 = (blockIdx.y * FH_QPB) / (n_heads / n_kv);

    #define FH_STAGE(KB, BUF) do { \
        unsigned short* dK = sKb + (BUF) * (FH_BK * FH_HD); \
        unsigned short* dV = sVb + (BUF) * (FH_HD * FH_BK); \
        const int _k0 = (KB) * FH_BK; \
        for (int i = threadIdx.x; i < FH_BK * (FH_HD / 8); i += blockDim.x) { \
            const int r = i / (FH_HD / 8), c8 = i - r * (FH_HD / 8); \
            fa_cp16(dK + FH_OFF(r, c8, FH_HD / 8), \
                    qkv16 + (long)(_k0 + r) * per * FH_HD + (long)(n_heads + kvh0) * FH_HD + c8 * 8); \
        } \
        for (int i = threadIdx.x; i < FH_HD * (FH_BK / 8); i += blockDim.x) { \
            const int d = i / (FH_BK / 8), c8 = i - d * (FH_BK / 8); \
            fa_cp16(dV + FH_OFF(d, c8, FH_BK / 8), \
                    vt16 + ((long)kvh0 * FH_HD + d) * Tp + _k0 + c8 * 8); \
        } \
    } while (0)

    // Prime the pipeline with FH_NS-1 tiles in flight. Out-of-range tiles still COMMIT, so the
    // wait_group accounting below stays uniform and the last tiles are still guaranteed complete.
    #pragma unroll
    for (int j = 0; j < FH_NS - 1; ++j) {
        if (kb_lo + j < kb_hi) FH_STAGE(kb_lo + j, j);
        FA_CP_COMMIT();
    }
    // Q went in through ordinary stores (generic proxy); GMMA reads shared through the async
    // proxy, so that traffic needs an explicit fence. K and V arrive by cp.async, which is already
    // in the async proxy — cp.async.wait_group plus the barrier below orders those.
    __syncthreads();
    asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");

    const unsigned long long dQ = fh_desc(sQ, 128, (FH_HD / 8) * 128);

    for (int kb = kb_lo; kb < kb_hi; ++kb) {
        const int k0 = kb * FH_BK;
        const int buf = (kb - kb_lo) % FH_NS;
        // Issue the FH_NS-1-ahead tile BEFORE waiting on this one, so the load for kb+3 is in
        // flight across the whole of kb's softmax and both wgmma groups.
        if (kb + FH_NS - 1 < kb_hi) FH_STAGE(kb + FH_NS - 1, (buf + FH_NS - 1) % FH_NS);
        FA_CP_COMMIT();
        asm volatile("cp.async.wait_group 1;\n" ::);   // <= FH_NS-1 outstanding => kb's is done
        __syncthreads();

        const unsigned long long dK =
            fh_desc(sKb + buf * (FH_BK * FH_HD), 128, (FH_HD / 8) * 128);
        const unsigned long long dV =
            fh_desc(sVb + buf * (FH_HD * FH_BK), 128, (FH_BK / 8) * 128);

        // ---- S = Q . K^T : 8 k-steps of 16 head dims, one m64n32k16 each ----
        float sc[FH_BK / 8][4];
        #pragma unroll
        for (int nt = 0; nt < FH_BK / 8; ++nt) { sc[nt][0] = 0.f; sc[nt][1] = 0.f; sc[nt][2] = 0.f; sc[nt][3] = 0.f; }
        FH_FENCE();
        #pragma unroll
        for (int s = 0; s < FH_HD / 16; ++s)
            #pragma unroll
            for (int kg = 0; kg < FH_BK / 32; ++kg)
                FH_MMA_SS(&sc[4 * kg][0], FH_ADV(dQ, s * 256),
                          FH_ADV(dK, s * 256 + kg * 4 * (FH_HD / 8) * 128));
        FH_COMMIT();
        FH_WAIT(0);

        // ---- causal mask + online softmax: unchanged, the fragment layout is the same ----
        const bool nomask = (k0 + FH_BK - 1 <= q0 + 16 * w) && (k0 + FH_BK <= T);
        float t0 = FA_MASKV, t1 = FA_MASKV;
        if (nomask) {
            #pragma unroll
            for (int nt = 0; nt < FH_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) { t0 = fmaxf(t0, sc[nt][j]); t1 = fmaxf(t1, sc[nt][2 + j]); }
            }
        } else {
            #pragma unroll
            for (int nt = 0; nt < FH_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    const int col = k0 + 8 * nt + t4 * 2 + j;
                    float v0 = sc[nt][j], v1 = sc[nt][2 + j];
                    if (col > row0 || col >= T) v0 = FA_MASKV;
                    if (col > row1 || col >= T) v1 = FA_MASKV;
                    sc[nt][j] = v0; sc[nt][2 + j] = v1;
                    t0 = fmaxf(t0, v0); t1 = fmaxf(t1, v1);
                }
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            t0 = fmaxf(t0, __shfl_xor_sync(0xffffffffu, t0, x));
            t1 = fmaxf(t1, __shfl_xor_sync(0xffffffffu, t1, x));
        }
        const float nm0 = fmaxf(m0, t0), nm1 = fmaxf(m1, t1);
        const bool live0 = nm0 > -1e29f, live1 = nm1 > -1e29f;
        const float r0 = __expf(m0 - nm0), r1 = __expf(m1 - nm1);
        float ps0 = 0.f, ps1 = 0.f;
        #pragma unroll
        for (int nt = 0; nt < FH_BK / 8; ++nt) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const float e0 = live0 ? __expf(sc[nt][j] - nm0) : 0.f;
                const float e1 = live1 ? __expf(sc[nt][2 + j] - nm1) : 0.f;
                sc[nt][j] = e0; sc[nt][2 + j] = e1;
                ps0 += e0; ps1 += e1;
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            ps0 += __shfl_xor_sync(0xffffffffu, ps0, x);
            ps1 += __shfl_xor_sync(0xffffffffu, ps1, x);
        }
        l0 = l0 * r0 + ps0; l1 = l1 * r1 + ps1;
        m0 = nm0; m1 = nm1;
        #pragma unroll
        for (int nt = 0; nt < FH_HD / 8; ++nt) {
            o[nt][0] *= r0; o[nt][1] *= r0; o[nt][2] *= r1; o[nt][3] *= r1;
        }

        // ---- O += P . V : P is already in exactly the A-in-registers fragment ----
        FH_FENCE();
        #pragma unroll
        for (int s = 0; s < FH_BK / 16; ++s) {
            const unsigned int a0 = fa_pack2(sc[2 * s][0],     sc[2 * s][1]);
            const unsigned int a1 = fa_pack2(sc[2 * s][2],     sc[2 * s][3]);
            const unsigned int a2 = fa_pack2(sc[2 * s + 1][0], sc[2 * s + 1][1]);
            const unsigned int a3 = fa_pack2(sc[2 * s + 1][2], sc[2 * s + 1][3]);
            #pragma unroll
            for (int g = 0; g < FH_HD / 32; ++g)
                FH_MMA_RS(&o[4 * g][0], a0, a1, a2, a3,
                          FH_ADV(dV, s * 256 + g * 4 * (FH_BK / 8) * 128));
        }
        FH_COMMIT();
        FH_WAIT(0);
        __syncthreads();
    }
    #undef FH_STAGE

    const float i0 = (l0 > 0.f) ? 1.f / l0 : 0.f, i1 = (l1 > 0.f) ? 1.f / l1 : 0.f;
    #pragma unroll
    for (int nt = 0; nt < FH_HD / 8; ++nt) {
        const int d = 8 * nt + t4 * 2;
        if (row0 < T) {
            unsigned short* dst = out + (long)row0 * n_heads * FH_HD + (long)h * FH_HD + d;
            dst[0] = f2bf(o[nt][0] * i0); dst[1] = f2bf(o[nt][1] * i0);
        }
        if (row1 < T) {
            unsigned short* dst = out + (long)row1 * n_heads * FH_HD + (long)h * FH_HD + d;
            dst[0] = f2bf(o[nt][2] * i1); dst[1] = f2bf(o[nt][3] * i1);
        }
    }
}
#endif

// ---------------------------------------------------------------------------------------------
// WARP-SPECIALISED variant (HOP=2) — FA-3's actual shape: a producer warpgroup, setmaxnreg, BK=128.
//
// Why this exists rather than just raising BK on the kernel above. BK=64 is the largest tile that
// still fits TWO blocks per SM at the uniform 128 registers/thread that __launch_bounds__ forces
// (the S accumulator is BQ*BK/128 registers, on top of O's 64). Going to BK=128 the naive way —
// __launch_bounds__(256,1), which does buy 256 registers/thread — measured 7.06 ms against BK=64's
// 6.29, because it halves occupancy from 16 warps/SM to 8. Trading warps for registers loses.
//
// setmaxnreg is the way out, and it is the piece FA-3 is built around: a third warpgroup adds
// WARPS without consuming the register budget, then the budget is redistributed away from it.
// SGLang's flash_fwd_sm90.py carries the table — (mma, producer) = 1WG:(256,56), 2WG:(240,24),
// 3WG:(160,32), and (224,40) when loading with cp.async rather than TMA, which is our case. Two
// MMA warpgroups at 224 and one producer at 40 is 2*128*224 + 128*40 = 62 464 <= 65 536, so it
// fits in one block of 12 warps: more warps than the naive BK=128 above, and the producer takes
// the cp.async staging — 22% of the kernel per the ablation — off the MMA warpgroups entirely.
//
// Producer/consumer handoff is four NAMED barriers rather than mbarriers: FULL(buf) = buf,
// EMPTY(buf) = 2+buf, each with all 384 threads as participants. `bar.arrive` contributes without
// blocking, `bar.sync` contributes and blocks, so the producer can run a tile ahead of the
// consumers instead of moving in lockstep with them.
#if __CUDA_ARCH__ == 900
#define FW_BQ 64
#define FW_BK 128
#define FW_HD 128
#define FW_QPB 2
#define FW_NS 2
#define FW_CONS 256                       // threads in the two MMA warpgroups
#define FW_ALL  384
#define FW_BAR(id) ((id))
#define FW_ARRIVE(id) asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "r"(FW_ALL) : "memory")
#define FW_WAIT(id)  asm volatile("bar.sync   %0, %1;\n" :: "r"(id), "r"(FW_ALL) : "memory")

// mbarrier producer/consumer handoff. A named-barrier event is a software convergence of 384
// threads (bar.arrive/bar.sync) per tile; an mbarrier flip is a shared-memory phase bit. Bigger:
// the producer used to block on `cp.async.wait_group 0` before it could signal readiness — with
// cp.async.mbarrier.arrive the barrier completes WHEN THE COPIES LAND, so the producer issues the
// next tile immediately and genuinely runs ahead of the consumers.
__device__ __forceinline__ unsigned int fw_smaddr(const void* p2) {
    unsigned int a;
    asm("{ .reg .u64 q; cvta.to.shared.u64 q, %1; cvt.u32.u64 %0, q; }" : "=r"(a) : "l"(p2));
    return a;
}
#define FW_MB_INIT(mb, cnt) \
    asm volatile("mbarrier.init.shared.b64 [%0], %1;\n" :: "r"(fw_smaddr(mb)), "r"(cnt))
#define FW_MB_ARRIVE(mb) do { unsigned long long _st; \
    asm volatile("mbarrier.arrive.shared.b64 %0, [%1];\n" : "=l"(_st) : "r"(fw_smaddr(mb))); \
} while (0)
#define FW_MB_CPARRIVE(mb) \
    asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];\n" :: "r"(fw_smaddr(mb)))
#define FW_MB_WAIT(mb, phase) \
    asm volatile("{ .reg .pred P1; LAB%=: mbarrier.try_wait.parity.shared.b64 P1, [%0], %1; " \
                 "@!P1 bra LAB%=; }\n" :: "r"(fw_smaddr(mb)), "r"(phase) : "memory")

// PERSISTENT BLOCKS + a work queue, biggest tiles first (FA-3's tile scheduler, minus the fancy).
//
// The launched grid was (T/64 query blocks, head-pairs) = 256 blocks on 114 SMs — 2.25 waves of
// blocks whose work varies 1:16 with causal depth. The efficiency-vs-T curve had already measured
// the cost: 71 TF/s at T=2k against 93 at T=8k (where more waves self-pack), i.e. ~24% of the
// kernel at 2k is wave-tail waste. So: one block per SM, an atomic ticket counter, and tickets
// ordered by DESCENDING query block so the deep diagonal tiles start first and the shallow ones
// fill the tail. Named-barrier bookkeeping per ticket must return to zero — every bar.arrive
// needs a matching bar.sync before the next ticket reuses the ids — hence the producer's drain
// loop below.
extern "C" __global__ __launch_bounds__(FW_ALL, 1) void fa_prefill_ws(
                                      const unsigned short* __restrict__ qkv16,
                                      const unsigned short* __restrict__ vt16,
                                      unsigned short* __restrict__ out,
                                      const int T, const int Tp, const int n_heads, const int n_kv,
                                      const float scale, unsigned int* __restrict__ ticket) {
    const int wgid = threadIdx.x >> 7;              // 0,1 = MMA consumers; 2 = producer
    const int per = n_heads + 2 * n_kv;
    const int nqb = (T + FW_BQ - 1) / FW_BQ;
    const int npairs = n_heads / FW_QPB;
    const int ntick = nqb * npairs;
    __shared__ int stick;

    extern __shared__ unsigned short fwsh[];
    unsigned short* base = (unsigned short*)(((unsigned long long)fwsh + 127ull) & ~127ull);
    // 4 mbarriers in the first 32 bytes; tiles start on the next 128-byte line.
    unsigned long long* mb_ready = (unsigned long long*)base;          // [FW_NS]
    unsigned long long* mb_freed = (unsigned long long*)base + FW_NS;  // [FW_NS]
    unsigned short* sQall = base + 64;                        // [FW_QPB][64][128]
    unsigned short* sKb = sQall + FW_QPB * FW_BQ * FW_HD;     // [FW_NS][128][128]
    unsigned short* sVb = sKb + FW_NS * FW_BK * FW_HD;        // [FW_NS][128][128] (V^T)
    if (threadIdx.x == 0) {
        FW_MB_INIT(&mb_ready[0], 128); FW_MB_INIT(&mb_ready[1], 128);  // producer cp completions
        FW_MB_INIT(&mb_freed[0], 256); FW_MB_INIT(&mb_freed[1], 256);  // consumer arrivals
    }
    // Rolling parity per buffer, per side. Counts balance every ticket, so the phase simply keeps
    // toggling across tickets — nothing is ever reset.
    unsigned int ph_ready0 = 0, ph_ready1 = 0, ph_freed0 = 0, ph_freed1 = 0;
    __syncthreads();

    for (;;) {
    if (threadIdx.x == 0) stick = (int)atomicAdd(ticket, 1u);
    __syncthreads();
    const int t_ = stick;
    if (t_ >= ntick) return;
    const int qb = nqb - 1 - t_ / npairs;           // deepest causal blocks first
    const int pair = t_ - (t_ / npairs) * npairs;
    const int q0 = qb * FW_BQ;
    const int kvh0 = (pair * FW_QPB) / (n_heads / n_kv);

    const int kb_max = (q0 + FW_BQ - 1) / FW_BK;
    const int kb_lo = 0, kb_hi = kb_max + 1;

    if (wgid == 2) {
        // ---------------- producer: nothing but staging ----------------
        asm volatile("setmaxnreg.dec.sync.aligned.u32 40;\n" ::: "memory");
        const int tid = threadIdx.x & 127;
        for (int kb = kb_lo; kb < kb_hi; ++kb) {
            const int buf = (kb - kb_lo) % FW_NS;
            const int k0 = kb * FW_BK;
            if (kb - kb_lo >= FW_NS) {                             // consumers released it
                FW_MB_WAIT(&mb_freed[buf], buf ? ph_freed1 : ph_freed0);
                if (buf) ph_freed1 ^= 1; else ph_freed0 ^= 1;
            }
            unsigned short* dK = sKb + buf * (FW_BK * FW_HD);
            unsigned short* dV = sVb + buf * (FW_HD * FW_BK);
            // BK=128 tiles overhang the 64-row padding of qkv16/vt16 on short prompts, so any
            // row/column beyond Tp is a zero-fill, not a read. Scores there are masked anyway
            // (col >= T), but the LOAD must not happen.
            for (int i = tid; i < FW_BK * (FW_HD / 8); i += 128) {
                const int r = i / (FW_HD / 8), c8 = i - r * (FW_HD / 8);
                fa_cp16z(dK + FH_OFF(r, c8, FW_HD / 8),
                         qkv16 + (long)(k0 + r) * per * FW_HD + (long)(n_heads + kvh0) * FW_HD + c8 * 8,
                         (k0 + r < Tp) ? 16u : 0u);
            }
            for (int i = tid; i < FW_HD * (FW_BK / 8); i += 128) {
                const int d = i / (FW_BK / 8), c8 = i - d * (FW_BK / 8);
                fa_cp16z(dV + FH_OFF(d, c8, FW_BK / 8),
                         vt16 + ((long)kvh0 * FW_HD + d) * Tp + k0 + c8 * 8,
                         (k0 + c8 * 8 + 7 < Tp) ? 16u : 0u);
            }
            // Completes when THIS thread's copies land — no wait_group, no proxy fence here
            // (the consumer fences before reading); the producer moves straight on.
            FW_MB_CPARRIVE(&mb_ready[buf]);
        }
        // Drain the releases the loop above never waited for, so parity stays aligned for the
        // next ticket (a skipped wait would leave this side one phase behind forever).
        const int ntiles = kb_hi - kb_lo;
        for (int j = (ntiles > FW_NS ? ntiles - FW_NS : 0); j < ntiles; ++j) {
            const int b = j % FW_NS;
            FW_MB_WAIT(&mb_freed[b], b ? ph_freed1 : ph_freed0);
            if (b) ph_freed1 ^= 1; else ph_freed0 ^= 1;
        }
        continue;
    }

    // ---------------- consumers ----------------
    asm volatile("setmaxnreg.inc.sync.aligned.u32 224;\n" ::: "memory");
    const int wg = wgid;                             // which query head of the pair
    const int h = pair * FW_QPB + wg;
    const int tid = threadIdx.x & 127;
    const int lane = threadIdx.x & 31, w = tid >> 5;
    const int gid = lane >> 2, t4 = lane & 3;
    unsigned short* sQ = sQall + wg * (FW_BQ * FW_HD);

    for (int i = tid; i < FW_BQ * (FW_HD / 8); i += 128) {
        const int r = i / (FW_HD / 8), c8 = i - r * (FW_HD / 8);
        const uint4 g = *(const uint4*)(qkv16 + (long)(q0 + r) * per * FW_HD + (long)h * FW_HD + c8 * 8);
        const unsigned int* gw = (const unsigned int*)&g;
        unsigned int sw[4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const float lo = h2f((unsigned short)(gw[j] & 0xFFFFu)) * scale;
            const float hi = h2f((unsigned short)(gw[j] >> 16)) * scale;
            sw[j] = fa_pack2(lo, hi);
        }
        *(uint4*)(sQ + FH_OFF(r, c8, FW_HD / 8)) = *(const uint4*)sw;
    }
    // Q is written by ordinary stores; GMMA reads through the async proxy.
    asm volatile("bar.sync 4, %0;\n" :: "r"(FW_CONS) : "memory");
    asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");

    float o[FW_HD / 8][4];
    #pragma unroll
    for (int nt = 0; nt < FW_HD / 8; ++nt) { o[nt][0] = 0.f; o[nt][1] = 0.f; o[nt][2] = 0.f; o[nt][3] = 0.f; }
    float m0 = FA_MASKV, m1 = FA_MASKV, l0 = 0.f, l1 = 0.f;
    const int row0 = q0 + 16 * w + gid, row1 = row0 + 8;
    const unsigned long long dQ = fh_desc(sQ, 128, (FW_HD / 8) * 128);

    // THE FA-3 MAINLOOP SHAPE — S(k+1) in flight underneath softmax(k) and PV(k).
    //
    // ncu on the sequential version: tensor pipe 22.5% active, no spills, nothing saturated, one
    // block per SM (which is intrinsic to setmaxnreg designs — the fat register pool leaves no
    // room for a second block, so there is no other block to hide latency behind). The
    // wait-S -> softmax -> PV -> wait chain simply exposed every latency in sequence. The earlier
    // intra-warpgroup-overlap experiment that measured FLAT was run in the 2-blocks/SM regime,
    // where the second block was already covering the gaps; at 1 block/SM the overlap is
    // load-bearing, not optional.
    //
    // wgmma groups RETIRE IN COMMIT ORDER, which is what makes the schedule sound: per iteration
    // the commits are PV(k) then S(k+1), so `wait_group 1` retires exactly PV(k) — the buffer
    // release the producer is waiting on — while S(k+1) keeps flying. The accumulators never
    // alias: S(k+1) accumulates into `sc` only after P(k) was packed out to `pr`, and `o` is
    // rescaled only after the previous PV retired.
    #define FW_ISSUE_S(BUF) do { \
        const unsigned long long _dk = \
            fh_desc(sKb + (BUF) * (FW_BK * FW_HD), 128, (FW_HD / 8) * 128); \
        _Pragma("unroll") \
        for (int nt = 0; nt < FW_BK / 8; ++nt) { sc[nt][0] = 0.f; sc[nt][1] = 0.f; \
                                                 sc[nt][2] = 0.f; sc[nt][3] = 0.f; } \
        FH_FENCE(); \
        _Pragma("unroll") \
        for (int s2 = 0; s2 < FW_HD / 16; ++s2) \
            FH_MMA128_SS(&sc[0][0], FH_ADV(dQ, s2 * 256), FH_ADV(_dk, s2 * 256)); \
        FH_COMMIT(); \
    } while (0)

    float sc[FW_BK / 8][4];
    FW_MB_WAIT(&mb_ready[0], ph_ready0); ph_ready0 ^= 1;          // first tile has landed
    asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
    FW_ISSUE_S(0);

    for (int kb = kb_lo; kb < kb_hi; ++kb) {
        const int k0 = kb * FW_BK;
        const int buf = (kb - kb_lo) % FW_NS;
        const unsigned long long dV = fh_desc(sVb + buf * (FW_HD * FW_BK), 128, (FW_BK / 8) * 128);
        // Entering this iteration the outstanding groups are [PV(k-1) already retired, S(k)]:
        // retire S(k) so the softmax below may read `sc`.
        FH_WAIT(0);

        const bool nomask = (k0 + FW_BK - 1 <= q0 + 16 * w) && (k0 + FW_BK <= T);
        float t0 = FA_MASKV, t1 = FA_MASKV;
        if (nomask) {
            #pragma unroll
            for (int nt = 0; nt < FW_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) { t0 = fmaxf(t0, sc[nt][j]); t1 = fmaxf(t1, sc[nt][2 + j]); }
            }
        } else {
            #pragma unroll
            for (int nt = 0; nt < FW_BK / 8; ++nt) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    const int col = k0 + 8 * nt + t4 * 2 + j;
                    float v0 = sc[nt][j], v1 = sc[nt][2 + j];
                    if (col > row0 || col >= T) v0 = FA_MASKV;
                    if (col > row1 || col >= T) v1 = FA_MASKV;
                    sc[nt][j] = v0; sc[nt][2 + j] = v1;
                    t0 = fmaxf(t0, v0); t1 = fmaxf(t1, v1);
                }
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            t0 = fmaxf(t0, __shfl_xor_sync(0xffffffffu, t0, x));
            t1 = fmaxf(t1, __shfl_xor_sync(0xffffffffu, t1, x));
        }
        const float nm0 = fmaxf(m0, t0), nm1 = fmaxf(m1, t1);
        const bool live0 = nm0 > -1e29f, live1 = nm1 > -1e29f;
        const float r0 = __expf(m0 - nm0), r1 = __expf(m1 - nm1);
        float ps0 = 0.f, ps1 = 0.f;
        #pragma unroll
        for (int nt = 0; nt < FW_BK / 8; ++nt) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const float e0 = live0 ? __expf(sc[nt][j] - nm0) : 0.f;
                const float e1 = live1 ? __expf(sc[nt][2 + j] - nm1) : 0.f;
                sc[nt][j] = e0; sc[nt][2 + j] = e1;
                ps0 += e0; ps1 += e1;
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            ps0 += __shfl_xor_sync(0xffffffffu, ps0, x);
            ps1 += __shfl_xor_sync(0xffffffffu, ps1, x);
        }
        l0 = l0 * r0 + ps0; l1 = l1 * r1 + ps1;
        m0 = nm0; m1 = nm1;
        #pragma unroll
        for (int nt = 0; nt < FW_HD / 8; ++nt) {
            o[nt][0] *= r0; o[nt][1] *= r0; o[nt][2] *= r1; o[nt][3] *= r1;
        }

        // P out of `sc` into dedicated registers, so S(k+1) can reuse `sc` as its accumulator.
        unsigned int pr[FW_BK / 16][4];
        #pragma unroll
        for (int s = 0; s < FW_BK / 16; ++s) {
            pr[s][0] = fa_pack2(sc[2 * s][0],     sc[2 * s][1]);
            pr[s][1] = fa_pack2(sc[2 * s][2],     sc[2 * s][3]);
            pr[s][2] = fa_pack2(sc[2 * s + 1][0], sc[2 * s + 1][1]);
            pr[s][3] = fa_pack2(sc[2 * s + 1][2], sc[2 * s + 1][3]);
        }
        FH_FENCE();
        #pragma unroll
        for (int s = 0; s < FW_BK / 16; ++s)
            FH_MMA128_RS(&o[0][0], pr[s][0], pr[s][1], pr[s][2], pr[s][3],
                         FH_ADV(dV, s * 256));
        FH_COMMIT();                                              // group: PV(k)
        // Issue S(k+1) into the OTHER buffer while PV(k) executes. Committed after PV(k), so the
        // wait_group 1 below retires PV(k) and leaves S(k+1) flying under the barrier work and
        // the next tile's staging.
        if (kb + 1 < kb_hi) {
            // tile k+1 has landed
            FW_MB_WAIT(&mb_ready[buf ^ 1], (buf ^ 1) ? ph_ready1 : ph_ready0);
            if (buf ^ 1) ph_ready1 ^= 1; else ph_ready0 ^= 1;
            asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
            FW_ISSUE_S(buf ^ 1);
            FH_WAIT(1);                                           // PV(k) retired; S(k+1) flying
        } else {
            FH_WAIT(0);                                           // last tile: drain everything
        }
        FW_MB_ARRIVE(&mb_freed[buf]);                             // buffer k is free
    }
    #undef FW_ISSUE_S

    const float i0 = (l0 > 0.f) ? 1.f / l0 : 0.f, i1 = (l1 > 0.f) ? 1.f / l1 : 0.f;
    #pragma unroll
    for (int nt = 0; nt < FW_HD / 8; ++nt) {
        const int d = 8 * nt + t4 * 2;
        if (row0 < T) {
            unsigned short* dst = out + (long)row0 * n_heads * FW_HD + (long)h * FW_HD + d;
            dst[0] = f2bf(o[nt][0] * i0); dst[1] = f2bf(o[nt][1] * i0);
        }
        if (row1 < T) {
            unsigned short* dst = out + (long)row1 * n_heads * FW_HD + (long)h * FW_HD + d;
            dst[0] = f2bf(o[nt][2] * i1); dst[1] = f2bf(o[nt][3] * i1);
        }
    }
    }   // ticket loop
}
#endif

#if __CUDA_ARCH__ != 900
extern "C" __global__ void fa_prefill_ws(const unsigned short* __restrict__ qkv16,
                                         const unsigned short* __restrict__ vt16,
                                         unsigned short* __restrict__ out, const int T,
                                         const int Tp, const int n_heads, const int n_kv,
                                         const float scale, unsigned int* __restrict__ ticket) {}
#endif

#if __CUDA_ARCH__ != 900
// The symbol must exist for module loading on every other arch; `hop` gates the launch.
extern "C" __global__ void fa_prefill_hop(const unsigned short* __restrict__ qkv16,
                                          const unsigned short* __restrict__ vt16,
                                          unsigned short* __restrict__ out, const int T,
                                          const int Tp, const int n_heads, const int n_kv,
                                          const float scale) {}
#endif

// ---------------------------------------------------------------------------------------------
// q4_0 x bf16 GEMM on tensor cores — weights read as q4, dequantised in-register into mma
// fragments, never materialised as bf16 anywhere.
//
// Why this kernel exists (three problems, one cause). Prefill and batched decode currently
// multiply *resident bf16 mirrors* of the q4 weights:
//   * they stream 2.82 GB of weights per prefill against 0.79 GB for q4 — our entire 4.5-bit
//     structural advantage, discarded exactly where the traffic is largest;
//   * mirrors for a 128-expert MoE layer are 805 MB, 38 GB across 48 layers, which is the real
//     reason batched/prefill MoE is still refused — not the routing;
//   * they double resident weight memory, which is what caps the server's KV slots.
//
// Shapes match `gemm_pre`: C[m, rows] = X[m, cols] . W[rows, cols]^T, W in q4_0 (32-wide blocks
// along `cols`, f16 scale + 4 packed u32). W's layout is already [n][k] with k contiguous, which is
// exactly the B-operand layout mma .row.col wants — no transpose anywhere.
//
// The dequant is the same bit trick as gemv_q4 v5 (measured 1829 GB/s, 91% of HBM peak): OR the
// nibble into the mantissa of 1024.0h and let one fma.rn.f16x2 fold out both the 1024 and q4_0's
// -8 zero point. Here it feeds shared memory rather than an accumulator.
#define QG_BM 64
#define QG_BN 64
#define QG_BK 32            // one q4_0 block per output row per iteration
#define QG_AS (QG_BK + 8)   // padded: unpadded strides put all 8 fragment rows on 4 banks
#define QG_BS (QG_BK + 8)

#if __CUDA_ARCH__ >= 800
extern "C" __global__ __launch_bounds__(128, 4) void gemm_q4_tc(
    const unsigned short* __restrict__ scales,   /* [rows, nblk] f16 */
    const uint4* __restrict__ quants,            /* [rows, nblk] packed */
    const unsigned short* __restrict__ X,        /* [m, cols] bf16 */
    float* __restrict__ C,                       /* [m, rows] f32, ldc = rows */
    const int m, const int rows, const int cols, const int nblk)
{
    const int tile_m = blockIdx.x * QG_BM, tile_n = blockIdx.y * QG_BN;
    const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
    const int gid = lane >> 2, t4 = lane & 3;

    __shared__ unsigned short sA[QG_BM * QG_AS];   // activations, [BM][BK]
    __shared__ unsigned short sB[QG_BN * QG_BS];   // dequantised weights, [BN][BK]

    // Each warp owns 16 rows of C and all 64 columns: 8 n-tiles x 4 f32 = 32 accumulators.
    float acc[QG_BN / 8][4];
    #pragma unroll
    for (int nt = 0; nt < QG_BN / 8; ++nt) {
        acc[nt][0] = 0.f; acc[nt][1] = 0.f; acc[nt][2] = 0.f; acc[nt][3] = 0.f;
    }

    for (int k0 = 0; k0 < cols; k0 += QG_BK) {
        __syncthreads();
        // activations: bf16 -> f16 (the mma operand type). Values are post-rmsnorm, O(1), so the
        // narrower exponent is not a concern; the reverse would not be true for raw residuals.
        for (int i = threadIdx.x; i < QG_BM * QG_BK; i += blockDim.x) {
            const int r = i / QG_BK, c = i - r * QG_BK;
            const int gr = tile_m + r;
            float v = 0.f;
            if (gr < m && (k0 + c) < cols) v = bf2f(X[(long)gr * cols + k0 + c]);
            sA[r * QG_AS + c] = f2h(v);
        }
        // weights: one q4_0 block per output row, dequantised straight into shared
        const int blk = k0 >> 5;
        for (int i = threadIdx.x; i < QG_BN * 4; i += blockDim.x) {
            const int r = i >> 2, part = i & 3;        // 4 threads per row, 8 values each
            const int gn = tile_n + r;
            if (gn >= rows || blk >= nblk) {
                #pragma unroll
                for (int j = 0; j < 8; ++j) sB[r * QG_BS + part * 8 + j] = 0;
                continue;
            }
            const long o = (long)gn * nblk + blk;
            const uint4 q = quants[o];
            const unsigned int wq[4] = { q.x, q.y, q.z, q.w };
            const unsigned int sc2 = ((unsigned int)scales[o] << 16) | (unsigned int)scales[o];
            // quantize_q4_0's packing: for word i, the LOW nibble of byte b is element (4i + b) and
            // the HIGH nibble is element (4i + b + 16). The 0x000F000F mask pulls the nibbles of
            // bytes 0 and 2 in one go, so a single half2 carries elements that are TWO apart — they
            // must be stored at +0 and +2, not adjacently.
            const int wi = part & 1, hi = part >> 1;   // part 0,1 -> low nibbles; 2,3 -> the +16 half
            #pragma unroll
            for (int b = 0; b < 2; ++b) {
                const int i4 = wi * 2 + b;
                const unsigned int v = wq[i4];
                #pragma unroll
                for (int r2 = 0; r2 < 2; ++r2) {
                    const unsigned int mant = ((v >> (8 * r2 + 4 * hi)) & 0x000F000Fu) | 0x64006400u;
                    // -1032 removes the 1024 bias AND q4_0's -8 zero point in one fma, then scale
                    const unsigned int dq = hmul2(hadd2(mant, 0xE408E408u), sc2);
                    const int e0 = i4 * 4 + hi * 16 + r2;
                    sB[r * QG_BS + e0]     = (unsigned short)(dq & 0xFFFFu);
                    sB[r * QG_BS + e0 + 2] = (unsigned short)(dq >> 16);
                }
            }
        }
        __syncthreads();

        #pragma unroll
        for (int ks = 0; ks < QG_BK / 16; ++ks) {
            const unsigned short* ap = sA + (long)(16 * w + gid) * QG_AS + ks * 16 + t4 * 2;
            const unsigned int a0 = *(const unsigned int*)(ap);
            const unsigned int a1 = *(const unsigned int*)(ap + 8 * QG_AS);
            const unsigned int a2 = *(const unsigned int*)(ap + 8);
            const unsigned int a3 = *(const unsigned int*)(ap + 8 * QG_AS + 8);
            #pragma unroll
            for (int nt = 0; nt < QG_BN / 8; ++nt) {
                const unsigned short* bp = sB + (long)(8 * nt + gid) * QG_BS + ks * 16 + t4 * 2;
                const unsigned int b0 = *(const unsigned int*)(bp);
                const unsigned int b1 = *(const unsigned int*)(bp + 8);
                FA_MMA(acc[nt][0], acc[nt][1], acc[nt][2], acc[nt][3], a0, a1, a2, a3, b0, b1);
            }
        }
    }

    // C is [m, rows] row-major with ldc = rows, matching gemm_pre's cuBLAS output exactly
    #pragma unroll
    for (int nt = 0; nt < QG_BN / 8; ++nt) {
        const int n0 = tile_n + 8 * nt + t4 * 2;
        const int r0 = tile_m + 16 * w + gid, r1 = r0 + 8;
        if (r0 < m) {
            if (n0 + 0 < rows) C[(long)r0 * rows + n0 + 0] = acc[nt][0];
            if (n0 + 1 < rows) C[(long)r0 * rows + n0 + 1] = acc[nt][1];
        }
        if (r1 < m) {
            if (n0 + 0 < rows) C[(long)r1 * rows + n0 + 0] = acc[nt][2];
            if (n0 + 1 < rows) C[(long)r1 * rows + n0 + 1] = acc[nt][3];
        }
    }
}
#else
extern "C" __global__ void gemm_q4_tc(const unsigned short* scales, const uint4* quants,
                                      const unsigned short* X, float* C,
                                      const int m, const int rows, const int cols, const int nblk) {}
#endif

// ---------------------------------------------------------------------------------------------
// q4_0 GEMM v2 — BARRIER-FREE, for decode widths.
//
// v1 (gemm_q4_tc) is correct but 10x slower than cuBLAS on bf16 mirrors at every width. ncu-level
// reasoning, not guesswork: BK=32 gave 64 k-iterations x 2 __syncthreads, BM=64 wasted 8/9 of the
// tile at M=8, and 96 blocks under-filled 132 SMs with the k-loop serial. It was barrier-bound at
// 47 GB/s where cuBLAS moves its (3.55x larger) weights at 1.68 TB/s.
//
// The fix is to stop staging weights at all. One q4_0 block is 32 k-values = exactly TWO mma
// k-steps, and a lane's B fragment needs only four of them (k = t4*2, +1, +8, +9). So each lane
// loads one uint4 + its scale straight from global and dequantises into b0/b1 in registers. The four
// lanes sharing `gid` read the same 16 bytes, which coalesces into one transaction, so the block
// still reads each weight row exactly once. No shared memory, no barriers.
//
// Conveniently, the four elements a lane needs for ks=0 are all LOW nibbles and for ks=1 all HIGH
// nibbles, and e / e+1 always fall in the same packed word, so each half2 is one shift-mask pair.
//
// Parallelism comes from k (split across blockIdx.z), because at decode widths the output tile is
// far too small to fill the GPU: M=8, N=6144 is 24 blocks of 4 warps. Partials go to a [ks][m][n]
// buffer that qg2_reduce sums — deliberately NOT f32 atomicAdd, which reorders, and a shifted last
// bit can flip an argmax and break the greedy-id gate that has caught every real bug here.
// 4 warps, each taking ONE 16-row group and ALL 64 columns of the tile.
//
// The first shape had BM=16 with warps splitting the columns, which meant M=64 needed four M-tiles
// that EACH re-read every weight — 28 MB, worse than the 25 MB bf16 mirror it was supposed to beat.
// With BM=64 the block reads each weight row exactly once for any M up to 64; the four warps issue
// the same weight loads, and L1 serves three of the four. At low M the redundant mma is free because
// the kernel is bandwidth-bound there, not tensor-bound.
#define QG2_BN 256
#define QG2_BM 64

// element `e` of a q4_0 block: e<16 is the low nibble of byte (e&3) in word (e>>2), e>=16 the high
__device__ __forceinline__ unsigned int qg2_pair(const unsigned int* w, int e, unsigned int sc2) {
    const int hi = e >> 4, e2 = e & 15;
    const int wi = e2 >> 2, sh = 8 * (e2 & 3) + 4 * hi;
    const unsigned int v = w[wi];
    // nibbles at `sh` and `sh+8` become the two halves of a half2 (16 bits apart)
    const unsigned int mant = ((v >> sh) & 0xFu) | (((v >> (sh + 8)) & 0xFu) << 16) | 0x64006400u;
    return hmul2(hadd2(mant, 0xE408E408u), sc2);   // -1032 kills the 1024 bias and q4_0's -8
}

// ncu drove every shape here. On the previous one: DRAM read exactly 7.19 MB — the weights, once,
// optimal — at 5.3 % of peak bandwidth, ALU 51 %, L1 50 %, tensor 8 %, and 3.34 M global load sectors
// for those 7.19 MB. Never bandwidth-bound; drowning in redundant load INSTRUCTIONS.
//
// Accounting for those sectors showed the culprit is the ACTIVATIONS, not the weights. With one
// column-block per 64 outputs there are 96 of them, and each re-reads the whole of A: at M=64 that is
// 96 x 64 x 2048 x 2 B = 25 MB against 7.1 MB of weights — and all four warps re-read it again, so
// ~100 MB of L1 traffic. Two fixes, both here:
//
//   * A is staged in SHARED once per block per k-chunk (one barrier per FOUR q4 blocks, not per
//     block), killing the 4x warp duplication and making the loads fully coalesced.
//   * QG2_BN doubles to 128, halving the number of column-blocks and so halving how often A is
//     re-read. Each warp then owns four n-tiles and all four row groups, so a weight load still
//     feeds 2 k-steps x 4 row groups = 8 mma.
// q4 blocks per shared-A chunk. Eight instead of four halves the barrier count: at M=8 each block
// walks 16 q4 blocks, so this is 2 chunks rather than 4, and with ~6 us of the 15 us being traffic
// the per-chunk barrier plus global-load latency is a measurable share of what remains.
#define QG2_KC 8
#define QG2_AS (QG2_KC * 32 + 8)       // padded: unpadded rows land on the same 4 banks

// bf16x2 arithmetic and a bf16 mma, so NOTHING is converted in the inner loop.
//
// A arrives as bf16 and was being converted to f16 on every fragment read — two bf2f plus two
// cvt.rn.f16.f32 per pair, ~96 ops per lane per k-block, the largest ALU term left. mma has a bf16
// variant on sm_80+, so if B dequantises to bf16 instead the conversions disappear entirely.
//
// The f16 dequant trick uses 1024.0h because a 4-bit nibble lands exactly in f16's 10-bit mantissa.
// bf16 has only 7 mantissa bits, so 1024 would round — but 128.0 (0x4300) works: 128*(1+q/128) =
// 128+q is exact for q in 0..127. Adding bf16 -136 (0xC308) then removes both the bias and q4_0's -8.
__device__ __forceinline__ float h2f_inv(unsigned short h) { return h2f(h); }

__device__ __forceinline__ unsigned int badd2(unsigned int a, unsigned int b) {
    unsigned int d; asm("add.rn.bf16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); return d;
}
__device__ __forceinline__ unsigned int bmul2(unsigned int a, unsigned int b) {
    unsigned int d; asm("mul.rn.bf16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); return d;
}
// two f32 -> one packed bf16x2 (low half = `lo`), same rounding as f2bf
__device__ __forceinline__ unsigned int pack_b2(float lo, float hi) {
    return ((unsigned int)f2bf(hi) << 16) | (unsigned int)f2bf(lo);
}

#define FA_MMA_BF(d0,d1,d2,d3, a0,a1,a2,a3, b0,b1) \
    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " \
                 "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" \
                 : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) \
                 : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1))

// ldmatrix: one instruction loads a whole 16x16 A fragment for the warp, already in mma register
// order, and is conflict-free by construction. It replaces FOUR scalar 32-bit shared loads per
// (row group, k-step) — with four row groups and two k-steps that is 32 LDS instructions per lane
// per k-block, which is what remains after traffic, redundancy and conversion were each eliminated.
//
// Address contract: lane l supplies the address of row (l & 15) at column offset ((l >> 4) * 8), and
// the four returned registers are a0..a3 in the order mma.m16n8k16 wants. Each lane's address must be
// 16-byte aligned: QG2_AS = 136 halves gives a 272-byte row stride (17*16) and every column offset is
// a multiple of 8 halves, so both hold.
__device__ __forceinline__ void ldm4(unsigned int* d, const unsigned short* smem) {
    unsigned int sp;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(sp) : "l"(smem));
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
                 : "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3]) : "r"(sp));
}

// bf16 -> e4m3 with a per-tensor scale, and the absmax that produces that scale.
//
// NVRTC here has no include paths, so `cuda_fp8.h` is unavailable — but the hardware instruction is
// reachable directly: `cvt.rn.satfinite.e4m3x2.f32` converts TWO floats to two e4m3 in one op, with
// round-to-nearest and saturation to +-448 for free.
__device__ __forceinline__ unsigned short f2e4m3x2(float lo, float hi) {
    unsigned short d;   // .b16 destination: two e4m3 packed, low byte = `lo`
    asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(d) : "f"(hi), "f"(lo));
    return d;
}

// The reverse, for the fp8 KV cache: two packed e4m3 -> two f32, via the hardware's
// e4m3x2 -> f16x2 conversion (exact: every e4m3 value is representable in f16).
__device__ __forceinline__ float2 e4m3x2_2f(unsigned short s) {
    unsigned int h2;
    asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(h2) : "h"(s));
    return make_float2(h2f((unsigned short)(h2 & 0xFFFFu)), h2f((unsigned short)(h2 >> 16)));
}

// Single f32 -> e4m3 byte (for strided writers whose adjacent elements live in other lanes).
__device__ __forceinline__ unsigned char f2e4m3(float v) {
    return (unsigned char)(f2e4m3x2(v, 0.f) & 0xFFu);
}

// Stage 1 of a two-stage absmax: 64 blocks each reduce a slice into one partial.
//
// The single-block version is the binding cost of the FP8 arm — it reduces m*hidden (131 072
// elements at M=64) once per layer, 28x per decode step, while the FP8 GEMM it feeds is small.
// Widening it from 256 to 1024 threads was already worth 37-64 %; parallelising across blocks is
// the rest of that win, and unlike delayed scaling it keeps the scale exact and per-step.
extern "C" __global__ void absmax_bf_part(const unsigned short* __restrict__ x,
                                          float* __restrict__ part, const long n) {
    float m = 0.f;
    const long stride = (long)gridDim.x * blockDim.x * 4;
    const long n4 = n & ~3L;
    for (long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 4; i < n4; i += stride) {
        const uint2 v = *reinterpret_cast<const uint2*>(x + i);
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.x & 0xFFFFu))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.x >> 16))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.y & 0xFFFFu))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.y >> 16))));
    }
    if (blockIdx.x == 0) {
        for (long i = n4 + (long)threadIdx.x; i < n; i += blockDim.x)
            m = fmaxf(m, fabsf(bf2f(x[i])));
    }
    __shared__ float r[256];
    r[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) r[threadIdx.x] = fmaxf(r[threadIdx.x], r[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) part[blockIdx.x] = r[0];
}

// Stage 2: reduce the 64 partials and emit BOTH the max and max/448 (what the Lt scale pointer
// wants), so the conversion and the matmul stay consistent without a second launch.
extern "C" __global__ void absmax_finish(const float* __restrict__ part,
                                         float* __restrict__ amax, float* __restrict__ scale,
                                         const int g) {
    float m = 0.f;
    for (int i = threadIdx.x; i < g; i += blockDim.x) m = fmaxf(m, part[i]);
    __shared__ float r[64];
    r[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) r[threadIdx.x] = fmaxf(r[threadIdx.x], r[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) { const float v = fmaxf(r[0], 1e-12f); amax[0] = v; scale[0] = v / 448.0f; }
}

// scale/448 — the value cuBLASLt's A/B scale pointers want, derived on-device so nothing
// round-trips to the host inside the decode loop.
extern "C" __global__ void mk_fp8_scale(const float* __restrict__ amax, float* __restrict__ out) {
    if (threadIdx.x == 0) out[0] = fmaxf(amax[0], 1e-12f) / 448.0f;
}

// PRODUCER-SIDE ABSMAX. The decode fp8 path used to run a 3-launch chain per GEMM
// (absmax_bf_part -> absmax_finish -> bf16_to_fp8): tiny serialized kernels whose GPU time —
// not launch overhead — summed to ~3 ms of an 8 ms M=64 step across 4 GEMMs x 32 layers. The
// producers below already hold every output value in registers, so they fold the block max in
// (warp tree + ONE float-bits atomicMax per block; valid for non-negatives) and the chain
// shrinks to amax_to_scale + bf16_to_fp8. amax_to_scale READS AND RESETS the accumulator, so
// one buffer serves all four GEMMs in stream order, graph-replay safe.
__device__ __forceinline__ void amax_commit(float v, float* redax, float* amax) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, o));
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    if (lane == 0) redax[wid] = v;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < (int)((blockDim.x + 31) >> 5)) ? redax[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t = fmaxf(t, __shfl_down_sync(0xffffffffu, t, o));
        if (lane == 0) atomicMax((unsigned int*)amax, __float_as_uint(t));
    }
}

// DELAYED-SCALE FP8 PRODUCERS (TransformerEngine-style). The per-GEMM chain
// (amax_to_scale + bf16_to_fp8) still cost ~1.3 ms/step at m=64 in launch+pass overhead; these
// variants write e4m3 DIRECTLY using the PREVIOUS step's amax (slot layout per (layer, site):
// [0]=amax, [1]=amax/448 for the Lt B-scale pointer, [2]=next-step amax accumulator), and
// fp8_scale_rotate flips [2] into [0]/[1] once per step. Stale-by-one amax clips gracefully on
// e4m3 saturation; the first step seeds amax=64.
__device__ __forceinline__ void amax_commit3(float v, float* redax, float* slot) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, o));
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    if (lane == 0) redax[wid] = v;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < (int)((blockDim.x + 31) >> 5)) ? redax[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t = fmaxf(t, __shfl_down_sync(0xffffffffu, t, o));
        if (lane == 0) atomicMax((unsigned int*)(slot + 2), __float_as_uint(t));
    }
}

extern "C" __global__ void fp8_scale_rotate(float* __restrict__ sl, const int nslot) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < nslot) {
        float* p = sl + (long)i * 3;
        const float a = fmaxf(p[2], 1e-6f);
        p[0] = a;
        p[1] = a / 448.0f;
        p[2] = 0.f;
    }
}

extern "C" __global__ void rmsnorm_add_fp8(float* __restrict__ x, const float* __restrict__ h,
                                           const float* __restrict__ w,
                                           unsigned char* __restrict__ out8,
                                           const int n, const float eps,
                                           float* __restrict__ slot) {
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    unsigned char* om = out8 + (long)m * n;
    __shared__ float red[32];
    __shared__ float redax[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] + hm[i]; xm[i] = v; s = fmaf(v, v, s); }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float inv = rsqrtf(red[0] / (float)n + eps);
    const float qinv = 448.0f / slot[0];
    float mv = 0.f;
    for (int i = 2 * threadIdx.x; i < n; i += 2 * blockDim.x) {
        const float a = xm[i] * inv * w[i];
        const float b = xm[i + 1] * inv * w[i + 1];
        mv = fmaxf(mv, fmaxf(fabsf(a), fabsf(b)));
        const unsigned short pq = f2e4m3x2(a * qinv, b * qinv);
        om[i] = (unsigned char)(pq & 0xFFu);
        om[i + 1] = (unsigned char)((pq >> 8) & 0xFFu);
    }
    amax_commit3(mv, redax, slot);
}

extern "C" __global__ void silu_mul_fp8(const unsigned short* __restrict__ gu,
                                        unsigned char* __restrict__ out8,
                                        const int mi, const int n8, const int act,
                                        float* __restrict__ slot) {
    __shared__ float redax[32];
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    float mv = 0.f;
    if (i < n8) {
        const int per = mi >> 3;
        const int t = i / per, j = i - t * per;
        const long base = (long)t * 2 * mi + (long)j * 8;
        const uint4 gv = *(const uint4*)(gu + base);
        const uint4 uv = *(const uint4*)(gu + base + mi);
        const unsigned int* gw = (const unsigned int*)&gv;
        const unsigned int* uw = (const unsigned int*)&uv;
        const float qinv = 448.0f / slot[0];
        unsigned char ob[8];
        #pragma unroll
        for (int k = 0; k < 4; ++k) {
            const float g0 = bf2f((unsigned short)(gw[k] & 0xFFFFu)), u0 = bf2f((unsigned short)(uw[k] & 0xFFFFu));
            const float g1 = bf2f((unsigned short)(gw[k] >> 16)),     u1 = bf2f((unsigned short)(uw[k] >> 16));
            const float v0 = ffn_act(g0, act) * u0;
            const float v1 = ffn_act(g1, act) * u1;
            mv = fmaxf(mv, fmaxf(fabsf(v0), fabsf(v1)));
            const unsigned short pq = f2e4m3x2(v0 * qinv, v1 * qinv);
            ob[2 * k] = (unsigned char)(pq & 0xFFu);
            ob[2 * k + 1] = (unsigned char)((pq >> 8) & 0xFFu);
        }
        *(uint2*)(out8 + (long)t * mi + (long)j * 8) = *(const uint2*)ob;
    }
    amax_commit3(mv, redax, slot);
}

extern "C" __global__ void attn_reduce_fp8(const float* __restrict__ pm, const float* __restrict__ pl,
                                           const float* __restrict__ pacc,
                                           unsigned char* __restrict__ out8,
                                           const int hd, const int nsplit, const int n_heads,
                                           const int* __restrict__ pos_p, const int win,
                                           float* __restrict__ slot) {
    const int gh = blockIdx.x;
    const int T = pos_p[gh / n_heads] + 1;
    const int used = min(nsplit, (T - 1 + ACHUNK) / ACHUNK);
    const int first = (win > 0 && T > win) ? (T - win) / ACHUNK : 0;
    __shared__ float mg, lg;
    __shared__ float redax[32];
    if (threadIdx.x == 0) {
        float mv2 = -1e30f;
        for (int s2 = first; s2 < used; ++s2) mv2 = fmaxf(mv2, pm[gh*nsplit+s2]);
        float l = 0.f;
        for (int s2 = first; s2 < used; ++s2) l += pl[gh*nsplit+s2] * __expf(pm[gh*nsplit+s2] - mv2);
        mg = mv2; lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    const float qinv = 448.0f / slot[0];
    float mv = 0.f;
    for (int d = 2 * (int)threadIdx.x; d < hd; d += 2 * blockDim.x) {
        float a = 0.f, b = 0.f;
        for (int s2 = first; s2 < used; ++s2) {
            const float e = __expf(pm[gh*nsplit+s2] - mg);
            a = fmaf(e, pacc[((long)gh*nsplit+s2)*hd + d], a);
            b = fmaf(e, pacc[((long)gh*nsplit+s2)*hd + d + 1], b);
        }
        a /= lg; b /= lg;
        mv = fmaxf(mv, fmaxf(fabsf(a), fabsf(b)));
        const unsigned short pq = f2e4m3x2(a * qinv, b * qinv);
        out8[(long)gh * hd + d] = (unsigned char)(pq & 0xFFu);
        out8[(long)gh * hd + d + 1] = (unsigned char)((pq >> 8) & 0xFFu);
    }
    amax_commit3(mv, redax, slot);
}

extern "C" __global__ void amax_to_scale(float* __restrict__ amax,
                                         float* __restrict__ sxmax, float* __restrict__ sx) {
    if (threadIdx.x == 0) {
        const float v = fmaxf(amax[0], 1e-12f);
        sxmax[0] = v; sx[0] = v / 448.0f;
        amax[0] = 0.f;
    }
}

extern "C" __global__ void rmsnorm_add_bf_amax(float* __restrict__ x, const float* __restrict__ h,
                                               const float* __restrict__ w,
                                               unsigned short* __restrict__ out,
                                               const int n, const float eps,
                                               float* __restrict__ amax) {
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    unsigned short* om = out + (long)m * n;
    __shared__ float red[32];
    __shared__ float redax[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] + hm[i]; xm[i] = v; s = fmaf(v, v, s); }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float inv = rsqrtf(red[0] / (float)n + eps);
    float mv = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        const float o2 = xm[i] * inv * w[i];
        om[i] = f2bf(o2);
        mv = fmaxf(mv, fabsf(o2));
    }
    amax_commit(mv, redax, amax);
}

extern "C" __global__ void silu_mul_bf_amax(const unsigned short* __restrict__ gu,
                                            unsigned short* __restrict__ out,
                                            const int mi, const int n8, const int act,
                                            float* __restrict__ amax) {
    __shared__ float redax[32];
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    float mv = 0.f;
    if (i < n8) {
        const int per = mi >> 3;
        const int t = i / per, j = i - t * per;
        const long base = (long)t * 2 * mi + (long)j * 8;
        const uint4 gv = *(const uint4*)(gu + base);
        const uint4 uv = *(const uint4*)(gu + base + mi);
        const unsigned int* gw = (const unsigned int*)&gv;
        const unsigned int* uw = (const unsigned int*)&uv;
        unsigned int ow[4];
        #pragma unroll
        for (int k = 0; k < 4; ++k) {
            const float g0 = bf2f((unsigned short)(gw[k] & 0xFFFFu)), u0 = bf2f((unsigned short)(uw[k] & 0xFFFFu));
            const float g1 = bf2f((unsigned short)(gw[k] >> 16)),     u1 = bf2f((unsigned short)(uw[k] >> 16));
            const float v0 = ffn_act(g0, act) * u0;
            const float v1 = ffn_act(g1, act) * u1;
            mv = fmaxf(mv, fmaxf(fabsf(v0), fabsf(v1)));
            ow[k] = (unsigned int)f2bf(v0) | ((unsigned int)f2bf(v1) << 16);
        }
        *(uint4*)(out + (long)t * mi + (long)j * 8) = *(const uint4*)ow;
    }
    amax_commit(mv, redax, amax);
}

extern "C" __global__ void attn_reduce_b16_amax(const float* __restrict__ pm, const float* __restrict__ pl,
                                                const float* __restrict__ pacc,
                                                unsigned short* __restrict__ out,
                                                const int hd, const int nsplit, const int n_heads,
                                                const int* __restrict__ pos_p, const int win,
                                                float* __restrict__ amax) {
    const int gh = blockIdx.x;
    const int T = pos_p[gh / n_heads] + 1;
    const int used = min(nsplit, (T - 1 + ACHUNK) / ACHUNK);
    const int first = (win > 0 && T > win) ? (T - win) / ACHUNK : 0;
    __shared__ float mg, lg;
    __shared__ float redax[32];
    if (threadIdx.x == 0) {
        float mv2 = -1e30f;
        for (int s2 = first; s2 < used; ++s2) mv2 = fmaxf(mv2, pm[gh*nsplit+s2]);
        float l = 0.f;
        for (int s2 = first; s2 < used; ++s2) l += pl[gh*nsplit+s2] * __expf(pm[gh*nsplit+s2] - mv2);
        mg = mv2; lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    float mv = 0.f;
    for (int d = threadIdx.x; d < hd; d += blockDim.x) {
        float a = 0.f;
        for (int s2 = first; s2 < used; ++s2)
            a = fmaf(__expf(pm[gh*nsplit+s2] - mg), pacc[((long)gh*nsplit+s2)*hd + d], a);
        const float o2 = a / lg;
        out[(long)gh * hd + d] = f2bf(o2);
        mv = fmaxf(mv, fabsf(o2));
    }
    amax_commit(mv, redax, amax);
}

extern "C" __global__ void bf16_to_fp8(const unsigned short* __restrict__ src,
                                       unsigned char* __restrict__ dst,
                                       const float* __restrict__ scale, const long n) {
    const long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 2;
    if (i >= n) return;
    const float inv = 448.0f / fmaxf(scale[0], 1e-12f);
    const float a = bf2f(src[i]) * inv;
    const float b = (i + 1 < n) ? bf2f(src[i + 1]) * inv : 0.f;
    const unsigned short p = f2e4m3x2(a, b);
    dst[i] = (unsigned char)(p & 0xFFu);
    if (i + 1 < n) dst[i + 1] = (unsigned char)((p >> 8) & 0xFFu);
}

// f32 flavours of the two-stage absmax + conversion — the SCORE-PATH attention output
// (p.attn) is f32, and converting it to bf16 first just to measure it would cost the extra
// pass the fp8 arm exists to avoid.
extern "C" __global__ void absmax_f_part(const float* __restrict__ x,
                                         float* __restrict__ part, const long n) {
    float m = 0.f;
    const long stride = (long)gridDim.x * blockDim.x * 4;
    const long n4 = n & ~3L;
    for (long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 4; i < n4; i += stride) {
        const float4 v = *reinterpret_cast<const float4*>(x + i);
        m = fmaxf(m, fmaxf(fmaxf(fabsf(v.x), fabsf(v.y)), fmaxf(fabsf(v.z), fabsf(v.w))));
    }
    if (blockIdx.x == 0) {
        for (long i = n4 + (long)threadIdx.x; i < n; i += blockDim.x)
            m = fmaxf(m, fabsf(x[i]));
    }
    __shared__ float r[256];
    r[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) r[threadIdx.x] = fmaxf(r[threadIdx.x], r[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) part[blockIdx.x] = r[0];
}

extern "C" __global__ void f32_to_fp8(const float* __restrict__ src,
                                      unsigned char* __restrict__ dst,
                                      const float* __restrict__ scale, const long n) {
    const long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 2;
    if (i >= n) return;
    const float inv = 448.0f / fmaxf(scale[0], 1e-12f);
    const float a = src[i] * inv;
    const float b = (i + 1 < n) ? src[i + 1] * inv : 0.f;
    const unsigned short p = f2e4m3x2(a, b);
    dst[i] = (unsigned char)(p & 0xFFu);
    if (i + 1 < n) dst[i + 1] = (unsigned char)((p >> 8) & 0xFFu);
}

// max|x| over a bf16 buffer — one block, grid-stride, so the scale never round-trips to the host.
extern "C" __global__ void absmax_bf(const unsigned short* __restrict__ x,
                                     float* __restrict__ out, const long n) {
    // 1024 threads, and each reads FOUR bf16 per step as one 8-byte load. At M=64 this reduces
    // 131 072 elements and runs 28 times per decode step, so the width is not cosmetic — at 256
    // scalar-loading threads it cost more than the FP8 GEMM it exists to serve.
    float m = 0.f;
    const long n4 = n & ~3L;
    for (long i = (long)threadIdx.x * 4; i < n4; i += (long)blockDim.x * 4) {
        const uint2 v = *reinterpret_cast<const uint2*>(x + i);
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.x & 0xFFFFu))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.x >> 16))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.y & 0xFFFFu))));
        m = fmaxf(m, fabsf(bf2f((unsigned short)(v.y >> 16))));
    }
    for (long i = n4 + (long)threadIdx.x; i < n; i += blockDim.x) m = fmaxf(m, fabsf(bf2f(x[i])));
    __shared__ float r[1024];
    r[threadIdx.x] = m;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) r[threadIdx.x] = fmaxf(r[threadIdx.x], r[threadIdx.x + o]);
        __syncthreads();
    }
    if (threadIdx.x == 0) out[0] = r[0];
}

// Repack a q4_0 block into the layout this GEMM's B fragment wants.
//
// The dequant was 9 ops per two weights because quantize_q4_0 puts the elements a fragment needs
// 8 bits apart (adjacent bytes), so building a half2 costs shift/mask/shift/shift/or/or. Put the two
// elements of each PAIR 16 bits apart in the same word instead and it collapses to shift/mask/or:
// one `& 0x000F000F` yields the half2 directly, exactly as the gemv_q4 v5 trick does.
//
// It fits q4_0's four words exactly: a 32-bit word holds nibbles at offsets 0,4,8,12,16,20,24,28,
// which is four pairs 16 bits apart, and 4 words x 4 pairs = the block's 16 pairs.
// Pair p (elements 2p, 2p+1) -> word p>>2, offsets 4*(p&3) and 4*(p&3)+16.
// A lane then reads word (2*ks) for b0 and (2*ks+1) for b1, both shifted by 4*t4. No change to the
// CPU quantiser and no change to the GEMV kernels, which keep the original layout.
extern "C" __global__ void q4_repack(const uint4* __restrict__ src, uint4* __restrict__ dst,
                                     const long nblocks) {
    const long b = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= nblocks) return;
    const uint4 q = src[b];
    const unsigned int w[4] = { q.x, q.y, q.z, q.w };
    unsigned int o[4] = { 0u, 0u, 0u, 0u };
    #pragma unroll
    for (int e = 0; e < 32; ++e) {
        const int hi = e >> 4, e2 = e & 15;
        const int wi = e2 >> 2, sh = 8 * (e2 & 3) + 4 * hi;
        const unsigned int v = (w[wi] >> sh) & 0xFu;      // still biased by +8, as q4_0 stores it
        const int pr = e >> 1;
        o[pr >> 2] |= v << (4 * (pr & 3) + 16 * (e & 1));
    }
    dst[b] = make_uint4(o[0], o[1], o[2], o[3]);
}

// As qg2_pair, but the caller already holds the packed word. `bo` is the byte offset within it and
// `hi` selects the nibble half (which is exactly the mma k-step).
__device__ __forceinline__ unsigned int qg2_pair_w(unsigned int v, int bo, int hi, unsigned int sc2) {
    const int sh = 8 * bo + 4 * hi;
    const unsigned int mant = ((v >> sh) & 0xFu) | (((v >> (sh + 8)) & 0xFu) << 16) | 0x64006400u;
    return hmul2(hadd2(mant, 0xE408E408u), sc2);
}

// EIGHT warps of 32 columns each, TEMPLATED on the number of 16-row groups.
//
// Activation traffic is (N/QG2_BN) * M * K * 2 — every column-block re-reads all of A — so relative
// to the weights it scales as 3.56*M/BN. At M=64 with BN=128 that was 12.6 MB of A against 7.1 MB of
// weights: the 3.55x weight-byte advantage was being spent on activations. BN=256 brings A down to
// about the weights.
//
// NRG exists because a fixed 64-row tile computes four row groups even when the batch is 8 or 16
// wide, so three quarters of both the mma AND the shared A-fragment reads produce zeros. The
// launcher picks the variant matching the batch width; weights are still read once per block, since
// re-reading only starts when M exceeds the tile.
// NL is n-tiles per warp, so the block covers BN = 8 warps * 8 * NL columns.
//
// BN has to shrink with the batch, for a reason the traffic arithmetic makes plain. Total bytes are
//   weights (fixed)  +  A: (N/BN)*M*K*2  +  reduction: nks*M*rows*4
// At M=64 the A term dominates and BN must be large. At M=8 A is tiny (0.8 MB at BN=256) while the
// REDUCTION is 6.3 MB — half the kernel — because a small tile count forces nks up to get enough
// blocks. Shrinking BN there buys blocks directly, so nks can fall: at M=8, BN=64 gives 96
// column-blocks and nks=6, i.e. 3.1 MB of A and 1.2 MB of partials against 0.8 + 6.3.
template<int NRG, int NL>
__device__ __forceinline__ void qg2_body(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const unsigned short* __restrict__ X, float* __restrict__ P,
    const int m, const int rows, const int cols, const int nblk, const int nks,
    unsigned short* sA)
{
    const int tile_m = blockIdx.x * (NRG * 16);
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    // BN = warps * 8 * NL. Deriving it from blockDim rather than assuming 8 warps: the M<=64
    // variant runs 4 warps, and hardcoding 8 made n0 stride 256 while the launcher tiled by 128.
    const int n0 = blockIdx.y * (int)((blockDim.x >> 5) * 8 * NL);
    const int gid = lane >> 2, t4 = lane & 3;
    const int sp = blockIdx.z;

    const int chunk = (nblk + nks - 1) / nks;
    const int b_lo = sp * chunk, b_hi = min(nblk, b_lo + chunk);

    float acc[NL][NRG][4];
    #pragma unroll
    for (int nl = 0; nl < NL; ++nl)
        #pragma unroll
        for (int rg = 0; rg < NRG; ++rg) {
            acc[nl][rg][0] = 0.f; acc[nl][rg][1] = 0.f;
            acc[nl][rg][2] = 0.f; acc[nl][rg][3] = 0.f;
        }

    for (int cb = b_lo; cb < b_hi; cb += QG2_KC) {
        const int nb = min(QG2_KC, b_hi - cb);
        const int kbase = cb << 5;
        __syncthreads();
        for (int i = threadIdx.x * 8; i < (NRG * 16) * (QG2_KC * 32); i += blockDim.x * 8) {
            const int r = i / (QG2_KC * 32), c = i - r * (QG2_KC * 32);
            const int gr = tile_m + r;
            uint4 v = make_uint4(0u, 0u, 0u, 0u);
            if (gr < m && (kbase + c + 8) <= cols && c < nb * 32) {
                v = *(const uint4*)(X + (long)gr * cols + kbase + c);
            }
            *(uint4*)(sA + r * QG2_AS + c) = v;
        }
        __syncthreads();

        #pragma unroll
        for (int j = 0; j < QG2_KC; ++j) {
            if (j >= nb) break;
            unsigned int a[NRG][2][4];
            const int lrow = lane & 15, lcol = (lane >> 4) << 3;
            #pragma unroll
            for (int rg = 0; rg < NRG; ++rg) {
                #pragma unroll
                for (int ks = 0; ks < 2; ++ks) {
                    ldm4(a[rg][ks], sA + (16 * rg + lrow) * QG2_AS + j * 32 + ks * 16 + lcol);
                }
            }
            #pragma unroll
            for (int nl = 0; nl < NL; ++nl) {
                const int row = n0 + (8 * NL) * warp + 8 * nl + gid;
                if (row >= rows) continue;
                const long o = (long)row * nblk + cb + j;
                const uint4 q = quants[o];
                const unsigned int wq[4] = { q.x, q.y, q.z, q.w };
                // one f16 scale -> bf16, once per weight block, amortised over 8 mma
                const unsigned short sbf = f2bf(h2f_inv(scales[o]));
                const unsigned int sc2 = ((unsigned int)sbf << 16) | (unsigned int)sbf;
                const int rsh = 4 * t4;
                #pragma unroll
                for (int ks = 0; ks < 2; ++ks) {
                    // PERMUTED layout: word (2ks) carries this lane's k,k+1 and (2ks+1) its k+8,k+9,
                    // each pair 16 bits apart, so one shift+mask is the whole unpack. 0x4300 is
                    // bf16 128.0; 0xC308 is bf16 -136, removing the bias and q4_0's -8 together.
                    const unsigned int m0 = ((wq[2 * ks]     >> rsh) & 0x000F000Fu) | 0x43004300u;
                    const unsigned int m1 = ((wq[2 * ks + 1] >> rsh) & 0x000F000Fu) | 0x43004300u;
                    const unsigned int b0 = bmul2(badd2(m0, 0xC308C308u), sc2);
                    const unsigned int b1 = bmul2(badd2(m1, 0xC308C308u), sc2);
                    #pragma unroll
                    for (int rg = 0; rg < NRG; ++rg) {
                        FA_MMA_BF(acc[nl][rg][0], acc[nl][rg][1], acc[nl][rg][2], acc[nl][rg][3],
                                  a[rg][ks][0], a[rg][ks][1], a[rg][ks][2], a[rg][ks][3], b0, b1);
                    }
                }
            }
        }
    }

    float* Pk = P + (long)sp * m * rows;
    #pragma unroll
    for (int nl = 0; nl < NL; ++nl) {
        const int c0 = n0 + (8 * NL) * warp + 8 * nl + t4 * 2;
        #pragma unroll
        for (int rg = 0; rg < NRG; ++rg) {
            const int rr0 = tile_m + 16 * rg + gid, rr1 = rr0 + 8;
            if (rr0 < m) {
                if (c0 < rows)     Pk[(long)rr0 * rows + c0]     = acc[nl][rg][0];
                if (c0 + 1 < rows) Pk[(long)rr0 * rows + c0 + 1] = acc[nl][rg][1];
            }
            if (rr1 < m) {
                if (c0 < rows)     Pk[(long)rr1 * rows + c0]     = acc[nl][rg][2];
                if (c0 + 1 < rows) Pk[(long)rr1 * rows + c0 + 1] = acc[nl][rg][3];
            }
        }
    }
}

// GROUPED TENSOR-CORE MoE GEMM (MOET=1). The per-pair GEMV reads an expert's weights once per
// (row, expert) PAIR and is latency-bound scalar work: the first batched ladder measured
// 319/644/905/995 at M=2/8/32/64 against vLLM's 943/3108/5610 — the deficit grows with M
// exactly as a per-pair kernel must.
//
// Shape: grid (n-tile, expert). A block owns one expert and 128 output columns; it walks that
// expert's pair list (the counting sort already produced mplist/moff/mcnt) in tiles of 16
// pairs, GATHERING those pairs' activation rows into shared, then runs the same bf16-mma
// dequant-in-register inner loop qg2_body proved — so each expert's weights are read ONCE per
// (column-tile, 16-pair tile) instead of once per pair, and the arithmetic runs on tensor
// cores instead of scalar FFMA.
//
// x is [rows, x_stride] f32 (the MoE FFN's normed activations); pair p reads row p / topk.
// Output y is [pairs, rows_per_expert] f32, exactly the per-pair layout silu_mul_moe and
// moe_combine already consume.
#define MG_KC 4                      // q4 blocks (128 k) staged per iteration
#define MG_AS (MG_KC * 32 + 8)       // padded activation row stride
#if __CUDA_ARCH__ >= 800
template<int NL, int MT> __device__ __forceinline__ void moe_t_body(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{
    const int nw = *nwork;
    for (int wi = blockIdx.y; wi < nw; wi += (int)gridDim.y) {
    const unsigned int witem = work[wi];
    const int e = (int)(witem >> 16);
    const int cnt = ecnt[e];
    if (cnt == 0) continue;
    const int off = eoff[e];
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
    const int gid = lane >> 2, t4 = lane & 3;
    const int lrow = lane & 15, lcol = (lane >> 4) << 3;
    // 8 warps x NL n-tiles x 8 columns, so the block covers 64*NL columns.
    const int n0 = blockIdx.x * (8 * NL * 8);

    __shared__ unsigned short sA[16 * MT * MG_AS];   // MT sub-tiles of 16 gathered pair-rows
    __shared__ unsigned int spair[16 * MT];          // their pair ids (for the scatter)

    // PAIR TILES ARE A GRID DIMENSION, not a loop: with identical prompts only ~8 experts
    // activate, so (column-tile, expert) alone was 96 blocks on 132 SMs while each walked all
    // of its expert's pair tiles serially. blockIdx.z spreads those tiles across the machine.
    {
        const int p0 = (int)(witem & 0xFFFFu) * (16 * MT);
        if (p0 >= cnt) continue;
        const int pb = min(16 * MT, cnt - p0);
        float acc[MT][NL][4];
        #pragma unroll
        for (int mt = 0; mt < MT; ++mt)
            #pragma unroll
            for (int nl = 0; nl < NL; ++nl) { acc[mt][nl][0] = 0.f; acc[mt][nl][1] = 0.f; acc[mt][nl][2] = 0.f; acc[mt][nl][3] = 0.f; }
        if ((int)threadIdx.x < 16 * MT) {
            spair[threadIdx.x] = ((int)threadIdx.x < pb) ? plist[off + p0 + threadIdx.x] : 0xFFFFFFFFu;
        }
        __syncthreads();

        for (int cb = 0; cb < nblk; cb += MG_KC) {
            const int nb = min(MG_KC, nblk - cb);
            const int kbase = cb << 5;
            __syncthreads();
            // gather: 16 pair-rows x (MG_KC*32) k-elements, f32 -> bf16, 8 at a time
            for (int i = threadIdx.x * 8; i < (16 * MT) * (MG_KC * 32); i += blockDim.x * 8) {
                const int r = i / (MG_KC * 32), c = i - r * (MG_KC * 32);
                unsigned int w4[4] = { 0u, 0u, 0u, 0u };
                if (r < pb && c < nb * 32 && (kbase + c + 8) <= x_stride) {
                    const unsigned int pid = spair[r];
                    const float* xr = x + (long)(pid / (unsigned int)x_div) * x_stride + kbase + c;
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) w4[j] = pack_b2(xr[2 * j], xr[2 * j + 1]);
                }
                *(uint4*)(sA + r * MG_AS + c) = *(const uint4*)w4;
            }
            __syncthreads();

            #pragma unroll
            for (int j = 0; j < MG_KC; ++j) {
                if (j >= nb) break;
                // A fragments for every row sub-tile this block owns, held live across the
                // weight loop below. That is the point of MT: one 16-byte quantised weight quad
                // now feeds MT mma instead of one, so arithmetic intensity scales with tile
                // height. At MT=1 a q4 quad buys ~16 FLOP/byte against an H100 balance point
                // near 300, which is why wide batches ran at about 1% of tensor peak.
                unsigned int a[MT][2][4];
                #pragma unroll
                for (int mt = 0; mt < MT; ++mt) {
                    #pragma unroll
                    for (int ks = 0; ks < 2; ++ks) {
                        ldm4(a[mt][ks], sA + (mt * 16 + lrow) * MG_AS + j * 32 + ks * 16 + lcol);
                    }
                }
                #pragma unroll
                for (int nl = 0; nl < NL; ++nl) {
                    const int row = n0 + (8 * NL) * warp + 8 * nl + gid;
                    if (row >= rows_per_expert) continue;
                    const long o = ((long)e * rows_per_expert + row) * (long)nblk + cb + j;
                    const uint4 q = quants[o];
                    const unsigned int wq[4] = { q.x, q.y, q.z, q.w };
                    const unsigned short sbf = f2bf(h2f_inv(scales[o]));
                    const unsigned int sc2 = ((unsigned int)sbf << 16) | (unsigned int)sbf;
                    const int rsh = 4 * t4;
                    #pragma unroll
                    for (int ks = 0; ks < 2; ++ks) {
                        const unsigned int m0 = ((wq[2 * ks]     >> rsh) & 0x000F000Fu) | 0x43004300u;
                        const unsigned int m1 = ((wq[2 * ks + 1] >> rsh) & 0x000F000Fu) | 0x43004300u;
                        const unsigned int b0 = bmul2(badd2(m0, 0xC308C308u), sc2);
                        const unsigned int b1 = bmul2(badd2(m1, 0xC308C308u), sc2);
                        #pragma unroll
                        for (int mt = 0; mt < MT; ++mt) {
                            FA_MMA_BF(acc[mt][nl][0], acc[mt][nl][1], acc[mt][nl][2], acc[mt][nl][3],
                                      a[mt][ks][0], a[mt][ks][1], a[mt][ks][2], a[mt][ks][3], b0, b1);
                        }
                    }
                }
            }
        }
        // scatter: fragment rows gid / gid+8 are pair-tile rows; columns are this warp's tile
        #pragma unroll
        for (int mt = 0; mt < MT; ++mt) {
            const int ra = mt * 16 + gid, rb = ra + 8;
            #pragma unroll
            for (int nl = 0; nl < NL; ++nl) {
                const int c0 = n0 + (8 * NL) * warp + 8 * nl + t4 * 2;
                if (c0 >= rows_per_expert) continue;
                if (ra < pb) {
                    float* yr = y + (long)spair[ra] * rows_per_expert;
                    yr[c0] = acc[mt][nl][0];
                    if (c0 + 1 < rows_per_expert) yr[c0 + 1] = acc[mt][nl][1];
                }
                if (rb < pb) {
                    float* yr = y + (long)spair[rb] * rows_per_expert;
                    yr[c0] = acc[mt][nl][2];
                    if (c0 + 1 < rows_per_expert) yr[c0 + 1] = acc[mt][nl][3];
                }
            }
        }
        __syncthreads();
    }
    }
}

// Two widths, because the right one depends on how many work items the routing sort emitted and
// that changes with batch width. The grid is (column tiles) x (work items): at M=8 there are only
// a few dozen work items, so halving the tile width to double the column tiles is what keeps 132
// SMs fed (+10.1%), while at M=32/64 there are already plenty and the narrow tile just reloads
// weights more often (-5 to -8%). Widening further is not the other end of a trend -- 512-wide
// tiles cost 26-47% everywhere, because the grid collapses below the SM count. Gather traffic is
// never the binding constraint here; block count is.
extern "C" __global__ __launch_bounds__(256, 2) void gemm_q4_moe_t(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<2, 1>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

extern "C" __global__ __launch_bounds__(256, 4) void gemm_q4_moe_t64(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<1, 1>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

extern "C" __global__ __launch_bounds__(256, 3) void gemm_q4_moe_t64h2(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<1, 2>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

extern "C" __global__ __launch_bounds__(256, 2) void gemm_q4_moe_t64h4(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<1, 4>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

extern "C" __global__ __launch_bounds__(256, 2) void gemm_q4_moe_th2(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<2, 2>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

extern "C" __global__ __launch_bounds__(256, 1) void gemm_q4_moe_th4(
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,
    const float* __restrict__ x, float* __restrict__ y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* __restrict__ plist, const int* __restrict__ eoff,
    const int* __restrict__ ecnt, const int x_stride, const int x_div,
    const unsigned int* __restrict__ work, const int* __restrict__ nwork)
{ moe_t_body<2, 4>(scales, quants, x, y, nblk, rows_per_expert, cols, plist, eoff, ecnt,
                x_stride, x_div, work, nwork); }

#else
extern "C" __global__ void gemm_q4_moe_t(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
extern "C" __global__ void gemm_q4_moe_t64(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
extern "C" __global__ void gemm_q4_moe_t64h2(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
extern "C" __global__ void gemm_q4_moe_t64h4(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
extern "C" __global__ void gemm_q4_moe_th2(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
extern "C" __global__ void gemm_q4_moe_th4(
    const unsigned short* scales, const uint4* quants, const float* x, float* y,
    const int nblk, const int rows_per_expert, const int cols,
    const unsigned int* plist, const int* eoff, const int* ecnt,
    const int x_stride, const int x_div, const unsigned int* work, const int* nwork) {}
#endif

// MINB is the resident-blocks hint. The narrow variants are cheap — rg1 is 61 registers and 4.3 KB
// of shared against rg4's 118 and 17.4 KB — so they can keep four blocks per SM where rg4 fits two,
// and narrow batches are exactly where parallelism is scarce and latency needs hiding.
// BN = (THREADS/32) * 8 * NL. Warp count is a knob because it trades the TWO redundancies against
// each other: warps that differ only in columns all read the same A fragments (8 warps => every
// ldmatrix issued 8x), while warps that differ in rows would duplicate the weight loads instead.
// Fewer warps with more columns each raises mma-per-A-read without changing BN or the tile count.
#define QG2_ENTRY(NAME, NRG, NL, THREADS, MINB)                                                 \
extern "C" __global__ __launch_bounds__(THREADS, MINB) void NAME(                                      \
    const unsigned short* __restrict__ scales, const uint4* __restrict__ quants,                \
    const unsigned short* __restrict__ X, float* __restrict__ P,                                \
    const int m, const int rows, const int cols, const int nblk, const int nks) {                \
    __shared__ unsigned short sA[(NRG) * 16 * QG2_AS];                                          \
    qg2_body<NRG, NL>(scales, quants, X, P, m, rows, cols, nblk, nks, sA);                        \
}

#if __CUDA_ARCH__ >= 800
// MEASURED: pushing rg1 to six blocks made ptxas spill (40 registers + 4 bytes of stack) and cost
// more than the residency bought — 0.019 -> 0.020 ms. Four is the most it holds without spilling.
QG2_ENTRY(gemm_q4_v2_rg1,   1, 1, 256, 4)   // 8 warps, BN = 64,  M <= 16
QG2_ENTRY(gemm_q4_v2_rg2,   2, 2, 256, 3)   // 8 warps, BN = 128, M <= 32
QG2_ENTRY(gemm_q4_v2_rg4n2, 4, 2, 256, 2)   // 8 warps, BN = 128, M <= 64  — see below
QG2_ENTRY(gemm_q4_v2_rg4,   4, 4, 256, 2)   // 8 warps, BN = 256, M > 64
// Why M<=64 gets its own variant, in two steps. First: at BN=256 there are only 24 column-tiles and
// the reduction budget caps nks at 4, so the launch was 96 blocks — 48 of 132 SMs. Blocks, not
// traffic, were the limit (0.035 -> 0.028 ms).
//
// Then ncu said the remainder is SHARED-MEMORY bandwidth: L1TEX 51 %, DRAM only 8 %. With 8 warps
// splitting columns, all 8 read the SAME A fragments, so every ldmatrix is issued 8x — ~98 MB of
// shared reads per GEMM.
//
// MEASURED NEGATIVE, kept so it is not retried: dropping to 4 warps while doubling NL keeps BN and
// the tile count at 128 but gives each warp 4 n-tiles instead of 2, halving that redundancy. It was
// SLOWER — 0.028 -> 0.038 ms. Halving the warps also halves the warps in flight (1536 -> 768), and
// the lost latency hiding cost more than the saved shared reads. So the shared-read redundancy is
// real but it is NOT the binding constraint; occupancy is. Any fix has to raise mma-per-A-read
// WITHOUT removing warps — i.e. more columns per warp at constant warp count, which needs a larger
// BN and therefore a way to keep the block count up that is not split-K.
#else
extern "C" __global__ void gemm_q4_v2_rg1(const unsigned short* s, const uint4* q, const unsigned short* X, float* P, const int m, const int r, const int c, const int nb, const int nk) {}
extern "C" __global__ void gemm_q4_v2_rg2(const unsigned short* s, const uint4* q, const unsigned short* X, float* P, const int m, const int r, const int c, const int nb, const int nk) {}
extern "C" __global__ void gemm_q4_v2_rg4(const unsigned short* s, const uint4* q, const unsigned short* X, float* P, const int m, const int r, const int c, const int nb, const int nk) {}
#endif

// Deterministic split-K reduction: fixed summation order, so the result is bit-reproducible.
extern "C" __global__ void qg2_reduce(const float* __restrict__ P, float* __restrict__ C,
                                      const int n, const int nks) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float s = 0.f;
    for (int k = 0; k < nks; ++k) s += P[(long)k * n + i];
    C[i] = s;
}

extern "C" __global__ void dequant_q4_bf16(const unsigned short* __restrict__ scales,
                                          const uint4* __restrict__ quants,
                                          unsigned short* __restrict__ out,
                                          const int nblk, const int rows, const long ooff) {
    const long blk = (long)blockIdx.x * blockDim.x + threadIdx.x;   // one q-block per thread
    const long total = (long)rows * nblk;
    if (blk >= total) return;
    out += ooff;
    const uint4 q = quants[blk];
    const float sc = h2f(scales[blk]);
    unsigned short* o = out + blk * 32;
    unsigned int w[4] = { q.x, q.y, q.z, q.w };
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        const unsigned int v = w[i];
        #pragma unroll
        for (int byte = 0; byte < 4; ++byte) {
            const int lo = (int)((v >> (8 * byte)) & 0xFu) - 8;
            const int hi = (int)((v >> (8 * byte + 4)) & 0xFu) - 8;
            // element (4i + byte) and (4i + byte + 16) — the quantize_q4_0 packing
            o[i * 4 + byte] = f2bf(sc * (float)lo);
            o[i * 4 + byte + 16] = f2bf(sc * (float)hi);
        }
    }
}

// Per-ROW scaled f32 -> f16. The unscaled conversion loses precision badly on real activations
// (silu reaches ~1.5e4 by layer 2, where f16's ulp is 16), and it made the batched/prefill paths
// return different tokens than the q4_0 GEMV path. Divide each row by its own max, convert, and let
// `scale_rows` put the factor back on the f32 output — exactly the trick the GEMV already uses.
extern "C" __global__ void f32_to_bf16(const float* __restrict__ x,
                                       unsigned short* __restrict__ y, const int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] = f2bf(x[i]);
}

// y[m, r] *= max|x_m| — undoes the per-row scaling above.
extern "C" __global__ void scale_rows(float* __restrict__ y, const float* __restrict__ xs,
                                      const int rows, const int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    const int m = i / rows;
    y[i] *= xs[2 * m];
}

// f32 activations -> f16 for the cuBLAS operand
extern "C" __global__ void f32_to_f16(const float* __restrict__ x, unsigned short* __restrict__ y,
                                      const int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    unsigned short h;
    const float v = x[i];
    asm("{ .reg .f16 t; cvt.rn.f16.f32 t, %1; mov.b16 %0, t; }" : "=h"(h) : "f"(v));
    y[i] = h;
}

// ---------------- prefill (compute-bound: T tokens at once, tensor cores) ----------------
// Sequential prefill measured 144 tok/s against vLLM's 84 407 on the same model — it was the decode
// path run once per prompt token. Here T tokens go through together: projections are one GEMM each,
// and attention is two cuBLAS strided-batched GEMMs per layer (Q*K^T then P*V) so the T^2 work also
// lands on tensor cores instead of an FMA loop.
extern "C" __global__ void embed_prefill(const float* __restrict__ e, float* __restrict__ x,
                                         float* __restrict__ h,
                                         const unsigned int* __restrict__ tokens,
                                         const int n, const int T, const float esc) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= T * n) return;
    const int t = i / n, d = i - t * n;
    x[i] = e[(long)tokens[t] * n + d] * esc;
    h[i] = 0.f;
}

// q/k norm + RoPE at each token's OWN position, then append to the single-sequence KV cache.
extern "C" __global__ void rope_prefill(float* __restrict__ qkv, unsigned short* __restrict__ qkv16,
                                        unsigned short* __restrict__ kc,
                                        unsigned short* __restrict__ vc,
                                        const float* __restrict__ qw, const float* __restrict__ kw,
                                        const int n_heads, const int n_kv, const int hd,
                                        const int T, const float* __restrict__ invf, const int qkn,
                                        const float eps, const int pos0, const int* __restrict__ kv_slot_off_p,
                                        const int seg_t,     // wave prefill: rows are B segments of seg_t; seg_t==T => one segment
                                        const int max_seq) { // head-major slab: per-(slot,head) plane stride
    // ONE WARP PER HEAD, four heads per block. The previous shape was one 128-thread block per
    // head doing a 7-step __syncthreads reduction over 128 elements — 64k blocks and 3.3 ms of the
    // 36 ms prefill for ~80 MB of traffic, i.e. pure sync latency. Warp shuffles need no barrier.
    // SLOT-INDEPENDENT PREFILL GRAPH: the KV write offset is read from DEVICE memory, so ONE
    // captured graph per prompt length serves every slot. Keyed per (length, slot) it cost a
    // ~30 ms capture per slot — at 64 slots that was ~half the M=64 wall clock (STEP_PROF:
    // admit 581 ms per 64 steps against sync 265 ms).
    const int per = n_heads + 2 * n_kv;
    const int lane = threadIdx.x & 31;
    const int gidx = blockIdx.x * (int)(blockDim.x >> 5) + (int)(threadIdx.x >> 5);
    if (gidx >= T * per) return;
    const int t = gidx / per, hidx = gidx - t * per;
    // wave prefill: row t belongs to segment t/seg_t — its own position clock and KV slab
    const int kv_slot_off = kv_slot_off_p[t / seg_t];
    const int pos = pos0 + (t % seg_t);
    const int kvdim = n_kv * hd;
    float* row = qkv + (long)t * per * hd;
    // The f16 copy the attention path reads is written HERE, from values already in registers,
    // instead of by a separate pass that re-read 32 MB of f32 per layer to write 16 MB of f16.
    unsigned short* row16 = qkv16 + (long)t * per * hd + (long)hidx * hd;
    if (hidx >= n_heads + n_kv) {
        const int vh = hidx - n_heads - n_kv;
        const float* src = row + (long)(n_heads + n_kv + vh) * hd;
        const long off = ((long)kv_slot_off * n_kv + (long)vh * max_seq + pos) * hd;
        if (hd == 128) {          // vectorized: one float4 per lane instead of four strided f32
            const float4 v4 = *(const float4*)(src + lane * 4);
            const int i = lane * 4;
            vc[off+i] = f2bf(v4.x); vc[off+i+1] = f2bf(v4.y); vc[off+i+2] = f2bf(v4.z); vc[off+i+3] = f2bf(v4.w);
            unsigned short h4[4] = { f2h(v4.x), f2h(v4.y), f2h(v4.z), f2h(v4.w) };
            *(uint2*)(row16 + i) = *(const uint2*)h4;
            return;
        }
        for (int i = lane; i < hd; i += 32) {
            const float v = src[i];
            vc[off + i] = f2bf(v);
            row16[i] = f2h(v);
        }
        return;
    }
    const bool is_q = hidx < n_heads;
    float* base = row + (long)hidx * hd;
    const float* nw = is_q ? qw : kw;
    // REGISTER-RESIDENT for the (universal) hd=128 case: each lane holds elements
    // {lane, lane+32, lane+64, lane+96}, so the rope pairs (i, i+64) and (i+32, i+96) are BOTH in
    // the same lane's registers. norm + rotate + both emits then need exactly ONE read of the f32
    // row and NO f32 write-back at all. The old path re-read and re-wrote f32 through L2 at every
    // stage — norm write 24.6 MB, rotated write-back 24.6 MB per layer (the write-back was needed
    // only so the k-head loop below could re-read it; q heads got it for free, as pure waste) —
    // and that traffic was the whole cost: this stage measured 1.43 ms against a ~0.85 ms floor.
    if (hd == 128) {
        float v0 = base[lane], v1 = base[lane + 32], v2 = base[lane + 64], v3 = base[lane + 96];
        if (qkn) {
            float s2 = fmaf(v0, v0, fmaf(v1, v1, fmaf(v2, v2, v3 * v3)));
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) s2 += __shfl_xor_sync(0xffffffffu, s2, o);
            const float inv = rsqrtf(s2 * (1.f / 128.f) + eps);
            v0 *= inv * nw[lane];      v1 *= inv * nw[lane + 32];
            v2 *= inv * nw[lane + 64]; v3 *= inv * nw[lane + 96];
        }
        const float a0 = (float)pos * invf[lane], a1 = (float)pos * invf[lane + 32];
        const float c0 = __cosf(a0), s0 = __sinf(a0), c1 = __cosf(a1), s1 = __sinf(a1);
        const float r0 = v0 * c0 - v2 * s0, r2 = v0 * s0 + v2 * c0;   // pair (i, i+64)
        const float r1 = v1 * c1 - v3 * s1, r3 = v1 * s1 + v3 * c1;   // pair (i+32, i+96)
        row16[lane] = f2h(r0);      row16[lane + 32] = f2h(r1);
        row16[lane + 64] = f2h(r2); row16[lane + 96] = f2h(r3);
        if (!is_q) {
            const int kh = hidx - n_heads;
            const long off = ((long)kv_slot_off * n_kv + (long)kh * max_seq + pos) * hd;
            kc[off + lane] = f2bf(r0);      kc[off + lane + 32] = f2bf(r1);
            kc[off + lane + 64] = f2bf(r2); kc[off + lane + 96] = f2bf(r3);
        }
        return;
    }
    // generic head_dim fallback: the original staged path
    if (qkn) {
        float s2 = 0.f;
        for (int i = lane; i < hd; i += 32) s2 = fmaf(base[i], base[i], s2);
        for (int o = 16; o > 0; o >>= 1) s2 += __shfl_xor_sync(0xffffffffu, s2, o);
        const float inv = rsqrtf(s2 / (float)hd + eps);
        for (int i = lane; i < hd; i += 32) base[i] = base[i] * inv * nw[i];
        __syncwarp();
    }
    const int half = hd >> 1;
    for (int i = lane; i < half; i += 32) {
        const float ang = (float)pos * invf[i];
        const float c = __cosf(ang), sn = __sinf(ang);
        const float a = base[i], b2 = base[i + half];
        const float r0 = a * c - b2 * sn, r1 = a * sn + b2 * c;
        row16[i] = f2h(r0);
        row16[i + half] = f2h(r1);
        base[i] = r0;                 // k heads still feed the KV cache below
        base[i + half] = r1;
    }
    __syncwarp();
    if (!is_q) {
        const int kh = hidx - n_heads;
        const long off = ((long)kv_slot_off * n_kv + (long)kh * max_seq + pos) * hd;
        for (int i = lane; i < hd; i += 32) {
            kc[off + i] = f2bf(base[i]);
        }
    }
}

// fp8 (e4m3) KV twin of rope_prefill: same math and the same f16 qkv16 output for the FA path;
// only the slab writes change to e4m3 bytes. K's hd=128 register layout holds strided elements
// (lane, lane+32, +64, +96), so K writes single bytes — 32 consecutive lanes still coalesce.
extern "C" __global__ void rope_prefill_f8(float* __restrict__ qkv, unsigned short* __restrict__ qkv16,
                                           unsigned char* __restrict__ kc,
                                           unsigned char* __restrict__ vc,
                                           const float* __restrict__ qw, const float* __restrict__ kw,
                                           const int n_heads, const int n_kv, const int hd,
                                           const int T, const float* __restrict__ invf, const int qkn,
                                           const float eps, const int pos0, const int* __restrict__ kv_slot_off_p,
                                           const int seg_t, const int max_seq) {
    const int per = n_heads + 2 * n_kv;
    const int lane = threadIdx.x & 31;
    const int gidx = blockIdx.x * (int)(blockDim.x >> 5) + (int)(threadIdx.x >> 5);
    if (gidx >= T * per) return;
    const int t = gidx / per, hidx = gidx - t * per;
    const int kv_slot_off = kv_slot_off_p[t / seg_t];
    const int pos = pos0 + (t % seg_t);
    const int kvdim = n_kv * hd;
    float* row = qkv + (long)t * per * hd;
    unsigned short* row16 = qkv16 + (long)t * per * hd + (long)hidx * hd;
    if (hidx >= n_heads + n_kv) {
        const int vh = hidx - n_heads - n_kv;
        const float* src = row + (long)(n_heads + n_kv + vh) * hd;
        const long off = ((long)kv_slot_off * n_kv + (long)vh * max_seq + pos) * hd;
        if (hd == 128) {
            const float4 v4 = *(const float4*)(src + lane * 4);
            const int i = lane * 4;
            *reinterpret_cast<unsigned short*>(vc + off + i) = f2e4m3x2(v4.x, v4.y);
            *reinterpret_cast<unsigned short*>(vc + off + i + 2) = f2e4m3x2(v4.z, v4.w);
            unsigned short h4[4] = { f2h(v4.x), f2h(v4.y), f2h(v4.z), f2h(v4.w) };
            *(uint2*)(row16 + i) = *(const uint2*)h4;
            return;
        }
        for (int i = lane; i < hd; i += 32) {
            const float v = src[i];
            vc[off + i] = f2e4m3(v);
            row16[i] = f2h(v);
        }
        return;
    }
    const bool is_q = hidx < n_heads;
    float* base = row + (long)hidx * hd;
    const float* nw = is_q ? qw : kw;
    if (hd == 128) {
        float v0 = base[lane], v1 = base[lane + 32], v2 = base[lane + 64], v3 = base[lane + 96];
        if (qkn) {
            float s2 = fmaf(v0, v0, fmaf(v1, v1, fmaf(v2, v2, v3 * v3)));
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) s2 += __shfl_xor_sync(0xffffffffu, s2, o);
            const float inv = rsqrtf(s2 * (1.f / 128.f) + eps);
            v0 *= inv * nw[lane];      v1 *= inv * nw[lane + 32];
            v2 *= inv * nw[lane + 64]; v3 *= inv * nw[lane + 96];
        }
        const float a0 = (float)pos * invf[lane], a1 = (float)pos * invf[lane + 32];
        const float c0 = __cosf(a0), s0 = __sinf(a0), c1 = __cosf(a1), s1 = __sinf(a1);
        const float r0 = v0 * c0 - v2 * s0, r2 = v0 * s0 + v2 * c0;
        const float r1 = v1 * c1 - v3 * s1, r3 = v1 * s1 + v3 * c1;
        row16[lane] = f2h(r0);      row16[lane + 32] = f2h(r1);
        row16[lane + 64] = f2h(r2); row16[lane + 96] = f2h(r3);
        if (!is_q) {
            const int kh = hidx - n_heads;
            const long off = ((long)kv_slot_off * n_kv + (long)kh * max_seq + pos) * hd;
            kc[off + lane] = f2e4m3(r0);      kc[off + lane + 32] = f2e4m3(r1);
            kc[off + lane + 64] = f2e4m3(r2); kc[off + lane + 96] = f2e4m3(r3);
        }
        return;
    }
    if (qkn) {
        float s2 = 0.f;
        for (int i = lane; i < hd; i += 32) s2 = fmaf(base[i], base[i], s2);
        for (int o = 16; o > 0; o >>= 1) s2 += __shfl_xor_sync(0xffffffffu, s2, o);
        const float inv = rsqrtf(s2 / (float)hd + eps);
        for (int i = lane; i < hd; i += 32) base[i] = base[i] * inv * nw[i];
        __syncwarp();
    }
    const int half = hd >> 1;
    for (int i = lane; i < half; i += 32) {
        const float ang = (float)pos * invf[i];
        const float c = __cosf(ang), sn = __sinf(ang);
        const float a = base[i], b2 = base[i + half];
        const float r0 = a * c - b2 * sn, r1 = a * sn + b2 * c;
        row16[i] = f2h(r0);
        row16[i + half] = f2h(r1);
        base[i] = r0;
        base[i + half] = r1;
    }
    __syncwarp();
    if (!is_q) {
        const int kh = hidx - n_heads;
        const long off = ((long)kv_slot_off * n_kv + (long)kh * max_seq + pos) * hd;
        for (int i = lane; i < hd; i += 32) {
            kc[off + i] = f2e4m3(base[i]);
        }
    }
}

// WAVE-PREFILL ATTENTION: B same-length segments of seg_t tokens, one block per (segment, head).
// The admission wave's prefills used to replay one graph per request — 63 serial ~3.5 ms
// prefills put every first token at ~235 ms. Batched, the GEMMs amortise across the wave and
// attention is this kernel: seg_t <= 64 so Q,K tiles and the score matrix all fit in smem, and
// the whole wave's attention is ~1k blocks of ~70k MACs — microseconds. Longer prompts keep the
// serial path where their own T makes the per-request graph efficient.
// qkv16 layout as rope_prefill wrote it: [S rows][per*hd] f16, row = seg*seg_t + i.
extern "C" __global__ void attn_wave(const unsigned short* __restrict__ qkv16,
                                     unsigned short* __restrict__ out,   // [S, nh*hd] BF16
                                     const int seg_t, const int n_heads, const int n_kv,
                                     const int hd, const float scale) {
    const int seg = blockIdx.x, h = blockIdx.y;
    const int per = n_heads + 2 * n_kv;
    const int kvh = h / (n_heads / n_kv);
    const long base = (long)seg * seg_t * per * hd;
    extern __shared__ float aw[];
    float* sq = aw;                       // [seg_t][hd]  (f32 for the dot products)
    float* sk = aw + seg_t * hd;          // [seg_t][hd]
    float* sc = aw + 2 * seg_t * hd;      // [seg_t][seg_t]
    const int tid = threadIdx.x;
    for (int i = tid; i < seg_t * hd; i += blockDim.x) {
        const int r = i / hd, d = i - r * hd;
        sq[i] = h2f(qkv16[base + (long)r * per * hd + (long)h * hd + d]);
        sk[i] = h2f(qkv16[base + (long)r * per * hd + (long)(n_heads + kvh) * hd + d]);
    }
    __syncthreads();
    // scores, causal
    for (int i = tid; i < seg_t * seg_t; i += blockDim.x) {
        const int r = i / seg_t, c2 = i - r * seg_t;
        float d = 0.f;
        if (c2 <= r) {
            const float* q = sq + r * hd;
            const float* k = sk + c2 * hd;
            for (int x = 0; x < hd; ++x) d = fmaf(q[x], k[x], d);
            d *= scale;
        } else {
            d = -1e30f;
        }
        sc[i] = d;
    }
    __syncthreads();
    // per-row softmax (rows striped over threads; seg_t <= 64 so every thread owns whole rows)
    for (int r = tid; r < seg_t; r += blockDim.x) {
        float mx = -1e30f;
        for (int c2 = 0; c2 <= r; ++c2) mx = fmaxf(mx, sc[r * seg_t + c2]);
        float sum = 0.f;
        for (int c2 = 0; c2 <= r; ++c2) {
            const float e = __expf(sc[r * seg_t + c2] - mx);
            sc[r * seg_t + c2] = e;
            sum += e;
        }
        const float inv = 1.f / sum;
        for (int c2 = 0; c2 <= r; ++c2) sc[r * seg_t + c2] *= inv;
    }
    __syncthreads();
    // V tile REUSES sq's storage — Q is dead once the scores exist, and a fresh sv tile at
    // seg_t=64/hd=128 would blow past the 84 KB smem opt-in. Reading V from qkv16 per (r,d)
    // was 68k scattered global loads per block: 14.7 ms of the b=61 wave against 2.5 staged.
    for (int i = tid; i < seg_t * hd; i += blockDim.x) {
        const int r = i / hd, d = i - r * hd;
        sq[i] = h2f(qkv16[base + (long)r * per * hd + (long)(n_heads + n_kv + kvh) * hd + d]);
    }
    __syncthreads();
    // out[r] = P[r,:] · V — one thread per (row, dim) stripe, V smem-resident
    for (int i = tid; i < seg_t * hd; i += blockDim.x) {
        const int r = i / hd, d = i - r * hd;
        float a = 0.f;
        for (int c2 = 0; c2 <= r; ++c2) {
            a = fmaf(sc[r * seg_t + c2], sq[c2 * hd + d], a);
        }
        // attn16 is BF16 — gemm_pre consumes CUDA_R_16BF (f2h here = garbage o-proj every layer)
        out[((long)seg * seg_t + r) * (long)(n_heads * hd) + (long)h * hd + d] = f2bf(a);
    }
}

// Causal softmax over the [nh, T, T] score block, in place, and emit the f16 copy the second GEMM
// consumes. One block per (head, query row).
// Causal softmax over f16 scores, TWO passes instead of three.
//
// The third pass existed only to divide by the row sum. That is a full read+write of a 128 MB
// score block per layer (~7 GB over 28 layers) to apply a per-row scalar — so instead the row sums
// are emitted to `rowsum` and folded into the attention output afterwards by `scale_out_rows`,
// which touches T*nh*hd (16 MB) rather than nh*T*T.
extern "C" __global__ void softmax_causal(unsigned short* __restrict__ sc16,
                                          float* __restrict__ rowsum,
                                          const int T, const float scale, const int win,
                                          const int seg_t) {
    const int h = blockIdx.y, i = blockIdx.x;
    unsigned short* row16 = sc16 + ((long)h * T + i) * T;
    const int n = i + 1;                      // causal: attend to 0..i
    // WAVE: rows are B independent segments of seg_t — row i may only attend within its own
    // segment, so the floor rises to the segment start. seg_t == T => one segment, unchanged.
    // The masked P entries are written as hard zeros, so the full-T PV GEMM is per-segment
    // correct with no further masking.
    const int seg0 = (i / seg_t) * seg_t;
    // sliding window: attend only to the last `win` positions (win = 0 means unbounded).
    // In-segment offsets are < seg_t <= 64 << win, so the window never bites inside a segment.
    int lo = (win > 0 && n > win) ? (n - win) : 0;
    lo = max(lo, seg0);
    __shared__ float red[256];
    float mx = -1e30f;
    for (int j = lo + threadIdx.x; j < n; j += blockDim.x) mx = fmaxf(mx, h2f(row16[j]) * scale);
    red[threadIdx.x] = mx; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + o]);
        __syncthreads();
    }
    const float m = red[0]; __syncthreads();
    float sum = 0.f;
    for (int j = threadIdx.x; j < T; j += blockDim.x) {
        // unnormalised exp for the in-window part, hard zero for the masked tail AND the
        // beyond-window head (the GEMM computed the full row, so both must be cleared or they
        // would leak into PV)
        float e2 = 0.f;
        if (j >= lo && j < n) { e2 = __expf(h2f(row16[j]) * scale - m); sum += e2; }
        unsigned short hh;
        asm("{ .reg .f16 t; cvt.rn.f16.f32 t, %1; mov.b16 %0, t; }" : "=h"(hh) : "f"(e2));
        row16[j] = hh;
    }
    red[threadIdx.x] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o];
        __syncthreads();
    }
    if (threadIdx.x == 0) rowsum[(long)h * T + i] = (red[0] > 0.f) ? (1.f / red[0]) : 0.f;
}

// out[t, h*hd + d] *= 1/rowsum[h, t] — the normalisation the softmax no longer applies.
extern "C" __global__ void scale_out_rows(float* __restrict__ out, const float* __restrict__ rowsum,
                                          const int T, const int nh, const int hd) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    const int n = T * nh * hd;
    if (i >= n) return;
    const int t = i / (nh * hd);
    const int h = (i - t * nh * hd) / hd;
    out[i] *= rowsum[(long)h * T + t];
}

extern "C" __global__ void embed_batch(const float* __restrict__ e, float* __restrict__ x,
                                       float* __restrict__ h,
                                       const unsigned int* __restrict__ tokens,
                                       const int n, const int M, const float esc) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= M * n) return;
    const int m = i / n, d = i - m * n;
    x[i] = e[(long)tokens[m] * n + d] * esc;
    h[i] = 0.f;
}

extern "C" __global__ void rmsnorm_add_b(float* __restrict__ x, const float* __restrict__ h,
                                         const float* __restrict__ w, float* __restrict__ out,
                                         float* __restrict__ xs, const int n, const float eps) {
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    float* om = out + (long)m * n;
    __shared__ float red[256];
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] + hm[i]; xm[i] = v; s = fmaf(v, v, s); }
    red[threadIdx.x] = s; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o]; __syncthreads(); }
    const float inv = rsqrtf(red[0] / (float)n + eps);
    __syncthreads();
    float mx = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] * inv * w[i]; om[i] = v; mx = fmaxf(mx, fabsf(v)); }
    red[threadIdx.x] = mx; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + o]); __syncthreads(); }
    if (threadIdx.x == 0) { const float v = red[0]; xs[2*m] = v; xs[2*m+1] = (v > 0.f) ? (1.f/v) : 0.f; }
}

// rmsnorm(x + h) * w emitted as bf16 for the tensor-core GEMMs.
//
// Two things go away versus rmsnorm_add_b: the f32 output that the next GEMM only re-read to
// convert (16 MB read + 8 MB written per call at T=2000), and the absmax reduction, which exists
// solely for the q4 GEMV path and is dead weight on the prefill/batched GEMM path.
extern "C" __global__ void rmsnorm_add_hf(float* __restrict__ x, const float* __restrict__ h,
                                          const float* __restrict__ w,
                                          unsigned short* __restrict__ out,
                                          const int n, const float eps) {
    // f16 twin of rmsnorm_add_bf: the Marlin q4 GEMM multiplies f16 operands, so the decode
    // arm's activations are produced in f16 directly instead of converted per tile.
    // Same warp-shuffle sum as the bf16 twin. It MUST be shuffle-based or red[] sized to the
    // launch: this kernel is launched with the 1024-thread config, and the old red[256] smem
    // tree wrote out of bounds for every thread above 255.
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    unsigned short* om = out + (long)m * n;
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] + hm[i]; xm[i] = v; s = fmaf(v, v, s); }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float inv = rsqrtf(red[0] / (float)n + eps);
    for (int i = threadIdx.x; i < n; i += blockDim.x) om[i] = f2h(xm[i] * inv * w[i]);
}

extern "C" __global__ void attn_reduce_h16(const float* __restrict__ pm, const float* __restrict__ pl,
                                           const float* __restrict__ pacc,
                                           unsigned short* __restrict__ out,
                                           const int hd, const int nsplit, const int n_heads,
                                           const int* __restrict__ pos_p, const int win) {
    const int gh = blockIdx.x;
    const int T = pos_p[gh / n_heads] + 1;
    const int used = min(nsplit, (T - 1 + ACHUNK) / ACHUNK);
    const int first = (win > 0 && T > win) ? (T - win) / ACHUNK : 0;
    __shared__ float mg, lg;
    if (threadIdx.x == 0) {
        float mv = -1e30f;
        for (int s2 = first; s2 < used; ++s2) mv = fmaxf(mv, pm[gh*nsplit+s2]);
        float l = 0.f;
        for (int s2 = first; s2 < used; ++s2) l += pl[gh*nsplit+s2] * __expf(pm[gh*nsplit+s2] - mv);
        mg = mv; lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    for (int d = threadIdx.x; d < hd; d += blockDim.x) {
        float a = 0.f;
        for (int s2 = first; s2 < used; ++s2)
            a = fmaf(__expf(pm[gh*nsplit+s2] - mg), pacc[((long)gh*nsplit+s2)*hd + d], a);
        out[(long)gh * hd + d] = f2h(a / lg);
    }
}

extern "C" __global__ void silu_mul_hf(const float* __restrict__ gu,
                                       unsigned short* __restrict__ out,
                                       const int mi, const int n, const int act) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    const int t = i / mi, j = i - t * mi;
    const long base = (long)t * 2 * mi;
    const float g = gu[base + j], u = gu[base + mi + j];
    out[i] = f2h(ffn_act(g, act) * u);
}

extern "C" __global__ void rmsnorm_add_bf(float* __restrict__ x, const float* __restrict__ h,
                                          const float* __restrict__ w,
                                          unsigned short* __restrict__ out,
                                          const int n, const float eps) {
    // MEASURED: 749 us of a 6.6 ms decode step across its 56 invocations — one block per row
    // with 256 threads walking 2048 elements is barrier/latency bound, not bandwidth bound.
    // The launch now uses 1024 threads; the smem tree that sized red[] to match cost ten
    // __syncthreads per row, so the sum is a warp-shuffle tree instead: one barrier, and red[]
    // shrinks to one float per warp.
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    unsigned short* om = out + (long)m * n;
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float v = xm[i] + hm[i]; xm[i] = v; s = fmaf(v, v, s); }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float inv = rsqrtf(red[0] / (float)n + eps);
    for (int i = threadIdx.x; i < n; i += blockDim.x) om[i] = f2bf(xm[i] * inv * w[i]);
}

// KV cache in bf16, not f32.
//
// At M=64 with ~150 tokens of context this is the DOMINANT traffic of a decode step: attention reads
// 64 slots x 150 positions x 1024 dims x 4 B x 2 (K and V) = 78 MB per layer, ~2.2 GB per step —
// nearly three times the 0.79 GB of weights. vLLM's cache is bf16; ours was f32, so we were paying
// double on the largest term in the step.
//
// LAYOUT IS HEAD-MAJOR: element (slot, kvh, pos, d) lives at ((slot*n_kv + kvh)*max_seq + pos)*hd
// + d. The old pos-major layout handed every attention reader 256 B (one head's row) every
// kvdim*2 B — attn_partial_b measured 0.68 TB/s (34% of peak) on it at 2k context. Head-major
// makes each (slot, head) reader STREAM its plane contiguously, and the query heads of one GQA
// group hit the SAME plane through L2. A slot's block is still max_seq*kvdim elements at
// slot*max_seq*kvdim (so slot-level offsets like d_kvoff and move_slot's block math carry over);
// only the within-slot order changed.
extern "C" __global__ void qk_rope_b(float* __restrict__ qkv, unsigned short* __restrict__ kc,
                                     unsigned short* __restrict__ vc,
                                     const float* __restrict__ qw, const float* __restrict__ kw,
                                     const int n_heads, const int n_kv, const int hd,
                                     const int* __restrict__ pos_p, const float* __restrict__ invf,
                                     const int qkn, const float eps, const int max_seq,
                                     const unsigned int* __restrict__ active, const int rot) {
    const int per = n_heads + 2 * n_kv;
    const int m = blockIdx.x / per, hidx = blockIdx.x - m * per;
    // Inactive rows must not touch the KV cache: with mixed prefill a slot INSIDE the width can
    // be mid-prefill (reserved, inactive), and an unconditional write at its stale pos would
    // corrupt the prompt KV the chunks are laying down. (Their attention/logits are garbage
    // either way and are discarded; the KV write is the only side effect that persists.)
    if (active && !active[m]) return;
    // per-SEQUENCE position: with continuous batching the slots are at different offsets
    const int pos = pos_p[m];
    float* qkvm = qkv + (long)m * per * hd;
    if (hidx >= n_heads + n_kv) {
        const int vh = hidx - n_heads - n_kv;
        const float* src = qkvm + (long)(n_heads + n_kv + vh) * hd;
        unsigned short* dst = vc + (((long)m * n_kv + vh) * max_seq + pos) * hd;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) dst[i] = f2bf(src[i]);
        return;
    }
    const bool is_q = hidx < n_heads;
    float* base = qkvm + (long)hidx * hd;
    const float* nw = is_q ? qw : kw;
    __shared__ float red[128];
    if (qkn) {
        float s = 0.f;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) s = fmaf(base[i], base[i], s);
        red[threadIdx.x] = s; __syncthreads();
        for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o]; __syncthreads(); }
        const float inv = rsqrtf(red[0] / (float)hd + eps);
        for (int i = threadIdx.x; i < hd; i += blockDim.x) base[i] = base[i] * inv * nw[i];
        __syncthreads();
    }
    // PARTIAL ROTARY (rot > 0): only the first `rot` dims rotate, as pairs (i, i+rot/2); dims
    // >= rot pass through untouched. rot == 0 is the ordinary full-width rope over hd.
    const int rd = (rot > 0 && rot < hd) ? rot : hd;
    const int half = rd >> 1;
    for (int i = threadIdx.x; i < half; i += blockDim.x) {
        const float ang = (float)pos * invf[i];
        const float c = __cosf(ang), sn = __sinf(ang);
        const float a = base[i], b = base[i + half];
        base[i] = a * c - b * sn;
        base[i + half] = a * sn + b * c;
    }
    if (!is_q) {
        __syncthreads();
        const int kh = hidx - n_heads;
        unsigned short* dst = kc + (((long)m * n_kv + kh) * max_seq + pos) * hd;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) dst[i] = f2bf(base[i]);
    }
}

// fp8 (e4m3) KV twin of qk_rope_b: identical rope/norm math, K/V APPENDED AS BYTES. The
// element-index arithmetic is unchanged — one e4m3 byte per element where the original had one
// bf16 short. Exists as a separate kernel so the bf16 tier's hot path carries no branch.
extern "C" __global__ void qk_rope_b_f8(float* __restrict__ qkv, unsigned char* __restrict__ kc,
                                        unsigned char* __restrict__ vc,
                                        const float* __restrict__ qw, const float* __restrict__ kw,
                                        const int n_heads, const int n_kv, const int hd,
                                        const int* __restrict__ pos_p, const float* __restrict__ invf,
                                        const int qkn, const float eps, const int max_seq,
                                        const unsigned int* __restrict__ active) {
    const int per = n_heads + 2 * n_kv;
    const int m = blockIdx.x / per, hidx = blockIdx.x - m * per;
    if (active && !active[m]) return;
    const int pos = pos_p[m];
    float* qkvm = qkv + (long)m * per * hd;
    if (hidx >= n_heads + n_kv) {
        const int vh = hidx - n_heads - n_kv;
        const float* src = qkvm + (long)(n_heads + n_kv + vh) * hd;
        unsigned char* dst = vc + (((long)m * n_kv + vh) * max_seq + pos) * hd;
        for (int i = threadIdx.x * 2; i + 1 < hd; i += blockDim.x * 2)
            *reinterpret_cast<unsigned short*>(dst + i) = f2e4m3x2(src[i], src[i + 1]);
        return;
    }
    const bool is_q = hidx < n_heads;
    float* base = qkvm + (long)hidx * hd;
    const float* nw = is_q ? qw : kw;
    __shared__ float red[128];
    if (qkn) {
        float s = 0.f;
        for (int i = threadIdx.x; i < hd; i += blockDim.x) s = fmaf(base[i], base[i], s);
        red[threadIdx.x] = s; __syncthreads();
        for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o]; __syncthreads(); }
        const float inv = rsqrtf(red[0] / (float)hd + eps);
        for (int i = threadIdx.x; i < hd; i += blockDim.x) base[i] = base[i] * inv * nw[i];
        __syncthreads();
    }
    const int half = hd >> 1;
    for (int i = threadIdx.x; i < half; i += blockDim.x) {
        const float ang = (float)pos * invf[i];
        const float c = __cosf(ang), sn = __sinf(ang);
        const float a = base[i], b = base[i + half];
        base[i] = a * c - b * sn;
        base[i + half] = a * sn + b * c;
    }
    if (!is_q) {
        __syncthreads();
        const int kh = hidx - n_heads;
        unsigned char* dst = kc + (((long)m * n_kv + kh) * max_seq + pos) * hd;
        for (int i = threadIdx.x * 2; i + 1 < hd; i += blockDim.x * 2)
            *reinterpret_cast<unsigned short*>(dst + i) = f2e4m3x2(base[i], base[i + 1]);
    }
}

// grid.y comes from the LAUNCH (the kernel strides by gridDim.y), sized per batch width so that
// slot's position. The previous shape put ceil(max_seq/ACHUNK) chunk-blocks per (slot, head) in
// the grid — at max_seq 4096 that is 16 chunks, ~15 of which early-exit on a short context, and
// with 64 slots x 16 heads x 28 layers that was ~430k EMPTY BLOCK DISPATCHES per decode step, at
// ~8 ns each ≈ 2.5 ms/step. Measured: M=64 in-process decode 16.0k tok/s at max_seq=256 against
// 9.8k at max_seq=4096, SAME contexts — the missing 39% was the dispatch of nothing.
//
// But Y must scale with width: a flat Y=4 at M=1 launched 64 blocks onto 114 SMs, each walking 8
// chunks of a 2000-token context SERIALLY — single-stream long-context fell 289 -> 199 tok/s.
// The launch picks Y so blocks ≈ a few GPU-fulls: all 16 chunks in parallel at M=1, Y=4 at M>=8.
// The partials layout (nsplit-strided), the reduce and the captured graph are unchanged; each
// width's graph bakes its own Y.
// Async slot bookkeeping: a (1,1) kernel launch is stream-ordered where cudarc's memcpy_htod of
// a pageable host slice SYNCHRONIZES — which is what made a 64-request admission burst serialize
// (each occupy's 3 host copies waited out the slot's whole prefill: ~3.6 ms per request).
// Set one device i32 by VALUE. Exists because a captured pageable-host memcpy records the host
// POINTER — a stack temporary — and every graph replay re-reads that freed stack slot. The value
// here is a kernel parameter, baked into the captured node, which is exactly right for the
// prefill graphs (keyed on T).
// ============================ GATED DELTANET (linear attention) ============================
// The recurrent primitive of the Qwen3.5/3.6 hybrid layers (30 of 40 in Qwen3.6-35B-A3B, and
// the same shape in Ornith-35B). Ported from the pinned CPU spec in `deltanet.rs`, which is
// itself a verbatim mirror of transformers `torch_recurrent_gated_delta_rule`. Per token, per
// value head, with persistent state S in R^{dk x dv}:
//
//   q,k <- l2norm(q,k)  (eps INSIDE the sqrt, FLA convention);  q <- q/sqrt(dk)
//   beta = sigmoid(b);  g = -exp(A_log) * softplus(a + dt_bias)
//   S <- S * exp(g);  kv_mem = k^T S;  delta = (v - kv_mem) * beta
//   S <- S + k (x) delta;  o = q^T S;  core = rmsnorm(o) * w * silu(z)
//
// DEPTHWISE CAUSAL CONV + SiLU over the fused qkv projection, one thread per channel, ring
// state of the last `kn` inputs per (slot, channel). The ring shifts by one each token, which
// is `kn` reads and `kn` writes of 4 B — trivial next to the recurrence, so it stays scalar.
extern "C" __global__ void dn_conv(float* __restrict__ mixed,      // [M, conv_dim] in/out
                                   float* __restrict__ ring,       // [M, conv_dim, kn]
                                   const float* __restrict__ cw,   // [conv_dim, kn]
                                   const int conv_dim, const int kn, const int M) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= M * conv_dim) return;
    const int m = idx / conv_dim, c = idx - m * conv_dim;
    float* rg = ring + ((long)m * conv_dim + c) * kn;
    const float* w = cw + (long)c * kn;
    float acc = 0.f;
    #pragma unroll 4
    for (int t = 0; t < kn - 1; ++t) { const float v = rg[t + 1]; rg[t] = v; acc = fmaf(v, w[t], acc); }
    const float nv = mixed[(long)m * conv_dim + c];
    rg[kn - 1] = nv;
    acc = fmaf(nv, w[kn - 1], acc);
    mixed[(long)m * conv_dim + c] = acc / (1.f + __expf(-acc));   // SiLU
}

// THE RECURRENCE. One block per (slot, value head); one thread per output dim j, so every
// thread owns column j of S for the whole step and the two passes need no cross-thread traffic
// beyond the two norms. S is staged in shared ONCE (dk*dv*4 B = 64 KB at 128x128) and written
// back once: the naive global version reads and writes it twice, and this kernel is pure
// state bandwidth. Threads walk i with a contiguous [dv]-wide row per step, so every access is
// coalesced.
extern "C" __global__ void dn_step(float* __restrict__ S,          // [M, nv, dk, dv]
                                   const float* __restrict__ mixed,// [M, conv_dim] post-conv
                                   const float* __restrict__ zin,  // [M, nv*dv]
                                   const float* __restrict__ bin,  // [M, nv]
                                   const float* __restrict__ ain,  // [M, nv]
                                   const float* __restrict__ a_log,// [nv]
                                   const float* __restrict__ dtb,  // [nv]
                                   const float* __restrict__ nw,   // [dv]
                                   float* __restrict__ core,       // [M, nv*dv]
                                   const int nk, const int nv, const int dk, const int dv,
                                   const float eps) {
    const int h = blockIdx.x, m = blockIdx.y;
    const int j = threadIdx.x;                 // one thread per value dim
    const int rep = nv / nk, kh = h / rep;     // repeat_interleave: v-heads share a k-head
    const int conv_dim = 2 * nk * dk + nv * dv;
    const float* mx = mixed + (long)m * conv_dim;
    const float* qs = mx + (long)kh * dk;
    const float* ks = mx + (long)nk * dk + (long)kh * dk;
    const float* vs = mx + (long)2 * nk * dk + (long)h * dv;

    extern __shared__ float dsh[];
    float* sS = dsh;                 // [dk][dv]
    float* sq = dsh + dk * dv;       // [dk]
    float* sk = sq + dk;             // [dk]
    float* red = sk + dk;            // [32]

    // ---- l2norm(q), l2norm(k); q also takes the 1/sqrt(dk) scale ----
    // eps is INSIDE the sqrt (FLA convention) — rsqrt(sum + eps), not rsqrt(sum) with a guard.
    float qv = (j < dk) ? qs[j] : 0.f, kv = (j < dk) ? ks[j] : 0.f;
    float sq2 = qv * qv, sk2 = kv * kv;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sq2 += __shfl_xor_sync(0xffffffffu, sq2, o);
        sk2 += __shfl_xor_sync(0xffffffffu, sk2, o);
    }
    const int lane = j & 31, wid = j >> 5, nwarp = blockDim.x >> 5;
    if (lane == 0) { red[wid] = sq2; red[nwarp + wid] = sk2; }
    __syncthreads();
    if (j == 0) {
        float a2 = 0.f, b2 = 0.f;
        for (int w2 = 0; w2 < nwarp; ++w2) { a2 += red[w2]; b2 += red[nwarp + w2]; }
        red[0] = rsqrtf(a2 + 1e-6f);
        red[1] = rsqrtf(b2 + 1e-6f);
    }
    __syncthreads();
    const float qscale = red[0] * rsqrtf((float)dk), kscale = red[1];
    if (j < dk) { sq[j] = qv * qscale; sk[j] = kv * kscale; }

    // ---- per-head gates ----
    const float beta = 1.f / (1.f + __expf(-bin[(long)m * nv + h]));
    const float ax = ain[(long)m * nv + h] + dtb[h];
    // softplus, overflow-safe exactly as the reference's f32 path
    const float sp = (ax > 20.f) ? ax : __logf(1.f + __expf(ax));
    const float decay = __expf(-__expf(a_log[h]) * sp);
    __syncthreads();

    // ---- stage S ----
    float* Sg = S + ((long)m * nv + h) * dk * dv;
    for (int i = j; i < dk * dv; i += blockDim.x) sS[i] = Sg[i];
    __syncthreads();

    // ---- pass 1: decay in place, kv_mem[j] = sum_i k[i] * S[i][j] ----
    float kv_mem = 0.f;
    for (int i = 0; i < dk; ++i) {
        const float sv = sS[i * dv + j] * decay;
        sS[i * dv + j] = sv;
        kv_mem = fmaf(sk[i], sv, kv_mem);
    }
    const float delta = (vs[j] - kv_mem) * beta;

    // ---- pass 2: rank-1 update, o[j] = sum_i q[i] * S[i][j] ----
    float o = 0.f;
    for (int i = 0; i < dk; ++i) {
        const float sv = fmaf(sk[i], delta, sS[i * dv + j]);
        sS[i * dv + j] = sv;
        o = fmaf(sq[i], sv, o);
    }
    __syncthreads();
    for (int i = j; i < dk * dv; i += blockDim.x) Sg[i] = sS[i];

    // ---- gated RMSNorm per head: rmsnorm(o)*w * silu(z) ----
    float ss = o * o;
    #pragma unroll
    for (int o2 = 16; o2 > 0; o2 >>= 1) ss += __shfl_xor_sync(0xffffffffu, ss, o2);
    if (lane == 0) red[wid] = ss;
    __syncthreads();
    if (j == 0) {
        float t = 0.f;
        for (int w2 = 0; w2 < nwarp; ++w2) t += red[w2];
        red[0] = rsqrtf(t / (float)dv + eps);
    }
    __syncthreads();
    const float z = zin[(long)m * nv * dv + (long)h * dv + j];
    core[(long)m * nv * dv + (long)h * dv + j] =
        o * red[0] * nw[j] * (z / (1.f + __expf(-z)));
}

// ATTENTION OUTPUT GATE (Qwen3.5/3.6): out *= sigmoid(gate), where the gate rows live at the
// tail of the fused qkv buffer (see the loader). Applied to the bf16 attention output in place,
// between the reduce and o_proj, so neither of those kernels changes.
// SHARED EXPERT (Qwen3.5/3.6): every token runs one dense SwiGLU MLP in addition to its routed
// experts, scaled by sigmoid of a [1, hidden] projection of the layer input. One thread per
// (slot, hidden) element; the per-slot gate is a dot product folded into the same kernel so the
// scalar never round-trips through global memory.
extern "C" __global__ void shared_expert_add(float* __restrict__ sw_out,
                                             const float* __restrict__ x,
                                             const float* __restrict__ gate_w,
                                             const int hidden, const int M) {
    const int m = blockIdx.x;
    __shared__ float g;
    // one block per slot: warp 0 reduces the gate dot product, then everyone applies it
    float acc = 0.f;
    for (int i = threadIdx.x; i < hidden; i += blockDim.x)
        acc = fmaf(gate_w[i], x[(long)m * hidden + i], acc);
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) acc += __shfl_xor_sync(0xffffffffu, acc, o);
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    if (lane == 0) red[wid] = acc;
    __syncthreads();
    if (threadIdx.x == 0) {
        float t = 0.f;
        for (int w2 = 0; w2 < nwarp; ++w2) t += red[w2];
        g = 1.f / (1.f + __expf(-t));
    }
    __syncthreads();
    for (int i = threadIdx.x; i < hidden; i += blockDim.x)
        sw_out[(long)m * hidden + i] *= g;
}

extern "C" __global__ void attn_out_gate_b(const float* __restrict__ qkv,
                                           unsigned short* __restrict__ attn16,
                                           const int gate_off, const int nhhd,
                                           const int per_hd, const int M) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= M * nhhd) return;
    const int m = i / nhhd, j = i - m * nhhd;
    const float g = qkv[(long)m * per_hd + gate_off + j];
    const float o = bf2f(attn16[i]);
    attn16[i] = f2bf(o / (1.f + __expf(-g)));
}

extern "C" __global__ void set_i32(int* __restrict__ dst, const int v) { *dst = v; }

extern "C" __global__ void slot_ctl(unsigned int* __restrict__ tok, int* __restrict__ pos,
                                    unsigned int* __restrict__ active,
                                    const int slot, const unsigned int token, const int posv,
                                    const int mode) {   // 0: occupy (active+pos), 1: patch token
    if (mode == 0) { active[slot] = 1u; pos[slot] = posv; }
    else tok[slot] = token;
}

// STRIP ATTENTION (chunked prefill). Query rows are ONE CHUNK of a prompt — absolute positions
// [pos0, pos0+n) — and they attend to slab positions [0, pos0+r] (causal, r = row's offset in
// the chunk). Both K and V come from the slot's KV slab, because rope_prefill has already
// written this chunk's K/V there at the right positions; nothing reads qkv16's K/V here.
// One block per (32-row query tile, head), online softmax over 64-key tiles: the running
// max/sum/accumulator never materialize a [n, T] score matrix.
// PREFIX GATHER for the FA-based chunked prefill: copy K and V for positions [0, pos0) out of
// the slot's bf16 KV slab into the prefill buffer's f16 K/V slots, so the EXISTING flash kernel
// (f_fa, mma-based and our fastest attention) can run over rows [0, pos0+n) with no new
// attention kernel at all. The chunk's own rows are already in place from rope_prefill.
// Cost: one pass over the prefix per layer (~0.15 ms at pos0=1536 across 32 layers) plus FA
// over the full triangle instead of the strip — 1.25x the attention work at 2 chunks, 1.875x
// at 4, against attention being ~10% of prefill. The scalar strip kernel below is 9x.
extern "C" __global__ void kv_gather(const unsigned short* __restrict__ kc,
                                     const unsigned short* __restrict__ vc,
                                     unsigned short* __restrict__ qkv16,
                                     const int pos0, const int n_heads, const int n_kv,
                                     const int hd, const int kv_off, const int max_seq) {
    const int per = n_heads + 2 * n_kv;
    const int kvdim = n_kv * hd;
    const long total = (long)pos0 * kvdim;          // one element per (position, kv dim)
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total;
         i += (long)gridDim.x * blockDim.x) {
        const int t = (int)(i / kvdim), c = (int)(i - (long)t * kvdim);
        const int kvh = c / hd, d = c - kvh * hd;
        // head-major slab: the slot block starts at kv_off*n_kv planes of max_seq rows
        const long src = ((long)kv_off * n_kv + (long)kvh * max_seq + t) * hd + d;
        unsigned short* row = qkv16 + (long)t * per * hd;
        row[(long)(n_heads + kvh) * hd + d]         = f2h(bf2f(kc[src]));
        row[(long)(n_heads + n_kv + kvh) * hd + d]  = f2h(bf2f(vc[src]));
    }
}

// STATUS: the SCALAR fallback arm (CHUNKFA=0). Correct but 9x slower than a one-shot prefill;
// the shipped chunk arm gathers the prefix K/V and reuses f_fa instead (1.45x). Kept because it
// needs no prefill-buffer growth — a future memory-tight path may want it.
// Chunked output is byte-identical to the one-shot prefill at every chunk size tested
// (2, 4, 128, 512, 4096) on both an exact-reference short prompt and a 279-token prose prompt.
// Two bugs got it there, and neither was in the attention math:
//   1. the causal limit used the row's index WITHIN the query tile (pos0 + r) instead of its
//      absolute position (pos0 + r0 + r), so every tile past the first missed r0 keys;
//   2. the tail never set d_pos, so the SEQUENTIAL decode head resumed from a stale position.
//      Symptom: the first sampled token was right and everything after it was wrong — which
//      reads exactly like a broken attention kernel and is not one. A layer-0 numeric diff
//      against f_fa (row sums agreeing to bf16 rounding: -0.44344 vs -0.44238) is what proved
//      the kernel innocent and sent the search downstream.
// WHAT IS LEFT — SPEED. Llama-8B 2k prompt: TTFT 667 ms against the one-shot path's 72.8 (9x).
// The staged K/V tiles below removed the 32x re-read of every key row, but the inner product is
// still SCALAR: one lane computes a whole (row, key) dot product over 128 dims. f_fa does the
// same work with mma tiles, and that is the gap. The next version is f_fa's structure with K/V
// read from the slab instead of qkv16 — not another iteration on this shape.
extern "C" __global__ void attn_strip(const unsigned short* __restrict__ qkv16,
                                      const unsigned short* __restrict__ kc,
                                      const unsigned short* __restrict__ vc,
                                      unsigned short* __restrict__ out,
                                      const int n, const int pos0, const int n_heads,
                                      const int n_kv, const int hd, const int kv_off,
                                      const float scale, const int max_seq) {
    // One block per (32-row query tile, head). Online softmax over 64-key tiles: scores land in
    // shared memory ONCE (the first cut recomputed them for the PV pass and used one atomicAdd
    // per (key, dim) — 10x slower than a one-shot prefill), then the PV pass gives each lane
    // four output dims and walks the key tile in registers.
    const int qt = blockIdx.x, h = blockIdx.y;
    const int per = n_heads + 2 * n_kv;
    const int kvh = h / (n_heads / n_kv);
    const int kvdim = n_kv * hd;
    const int r0 = qt * 32;
    const int nq = min(32, n - r0);
    if (nq <= 0) return;
    extern __shared__ float ss[];
    float* qs  = ss;                       // [32][hd]
    float* acc = qs + 32 * hd;             // [32][hd]
    float* sc  = acc + 32 * hd;            // [32][64] scores for the current key tile
    float* rm  = sc + 32 * 64;             // [32] running max
    float* rl  = rm + 32;                  // [32] running sum
    float* rc  = rl + 32;                  // [32] this tile's rescale exp(m_old - m_new)
    // K and V tiles staged ONCE per key tile: without this every one of the 32 query rows in
    // the block re-reads the same K/V rows from global (32x the traffic), which is what made
    // the first version 7x slower than a one-shot prefill.
    unsigned short* ks = (unsigned short*)(rc + 32);   // [64][hd] bf16
    unsigned short* vs = ks + 64 * hd;                 // [64][hd] bf16
    const int tid = threadIdx.x, nthr = blockDim.x;
    for (int i = tid; i < nq * hd; i += nthr) {
        const int r = i / hd, d = i - r * hd;
        qs[i] = h2f(qkv16[(long)(r0 + r) * per * hd + (long)h * hd + d]);
        acc[i] = 0.f;
    }
    for (int i = tid; i < 32; i += nthr) { rm[i] = -1e30f; rl[i] = 0.f; }
    __syncthreads();
    const int warp = tid >> 5, lane = tid & 31, nwarp = nthr >> 5;
    const int kmax = pos0 + r0 + nq;
    for (int k0 = 0; k0 < kmax; k0 += 64) {
        const int kn = min(64, kmax - k0);
        __syncthreads();
        for (int i = tid; i < kn * (hd / 8); i += nthr) {
            const int t = i / (hd / 8), c = i - t * (hd / 8);
            const long g0 = ((long)kv_off * n_kv + (long)kvh * max_seq + (k0 + t)) * hd + c * 8;
            *(uint4*)(ks + (long)t * hd + c * 8) = *(const uint4*)(kc + g0);
            *(uint4*)(vs + (long)t * hd + c * 8) = *(const uint4*)(vc + g0);
        }
        __syncthreads();
        // ---- scores into smem, one warp per row, lanes over keys ----
        for (int r = warp; r < nq; r += nwarp) {
            const int lim = pos0 + r0 + r;   // ABSOLUTE position of this row; using pos0 + r
            const float* qr = qs + r * hd;   // made every tile past the first miss r0 keys
            float mloc = -1e30f;
            for (int t = lane; t < kn; t += 32) {
                const int kpos = k0 + t;
                float d2 = -1e30f;
                if (kpos <= lim) {
                    const unsigned short* kt = ks + (long)t * hd;
                    float a = 0.f;
                    for (int i = 0; i < hd; i += 4) {
                        const uint2 kk = *reinterpret_cast<const uint2*>(kt + i);
                        a = fmaf(qr[i+0], bf2f((unsigned short)(kk.x & 0xFFFFu)), a);
                        a = fmaf(qr[i+1], bf2f((unsigned short)(kk.x >> 16)), a);
                        a = fmaf(qr[i+2], bf2f((unsigned short)(kk.y & 0xFFFFu)), a);
                        a = fmaf(qr[i+3], bf2f((unsigned short)(kk.y >> 16)), a);
                    }
                    d2 = a * scale;
                    mloc = fmaxf(mloc, d2);
                }
                sc[r * 64 + t] = d2;
            }
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) mloc = fmaxf(mloc, __shfl_xor_sync(0xffffffffu, mloc, o));
            const float m_new = fmaxf(rm[r], mloc);
            float lsum = 0.f;
            for (int t = lane; t < kn; t += 32) {
                const float v = sc[r * 64 + t];
                const float e2 = (v > -1e29f) ? __expf(v - m_new) : 0.f;
                sc[r * 64 + t] = e2;
                lsum += e2;
            }
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) lsum += __shfl_xor_sync(0xffffffffu, lsum, o);
            if (lane == 0) {
                const float corr = (rm[r] > -1e29f) ? __expf(rm[r] - m_new) : 0.f;
                rl[r] = rl[r] * corr + lsum;
                rm[r] = m_new;
                rc[r] = corr;                // the PV pass rescales acc by this before adding
            }
        }
        __syncthreads();
        // ---- PV: each lane owns hd/32 dims of one row and walks the whole key tile ----
        for (int r = warp; r < nq; r += nwarp) {
            const int lim = pos0 + r0 + r;
            const float corr = rc[r];
            for (int dblk = lane * 4; dblk < hd; dblk += 128) {
                float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
                for (int t = 0; t < kn; ++t) {
                    const int kpos = k0 + t;
                    if (kpos > lim) break;
                    const float pw = sc[r * 64 + t];
                    const unsigned short* vt = vs + (long)t * hd + dblk;
                    const uint2 vv = *reinterpret_cast<const uint2*>(vt);
                    a0 = fmaf(pw, bf2f((unsigned short)(vv.x & 0xFFFFu)), a0);
                    a1 = fmaf(pw, bf2f((unsigned short)(vv.x >> 16)), a1);
                    a2 = fmaf(pw, bf2f((unsigned short)(vv.y & 0xFFFFu)), a2);
                    a3 = fmaf(pw, bf2f((unsigned short)(vv.y >> 16)), a3);
                }
                float* ar = acc + r * hd + dblk;
                ar[0] = ar[0] * corr + a0; ar[1] = ar[1] * corr + a1;
                ar[2] = ar[2] * corr + a2; ar[3] = ar[3] * corr + a3;
            }
        }
        __syncthreads();
    }
    for (int i = tid; i < nq * hd; i += nthr) {
        const int r = i / hd, d = i - r * hd;
        const float inv = (rl[r] > 0.f) ? (1.f / rl[r]) : 0.f;
        out[(long)(r0 + r) * (long)(n_heads * hd) + (long)h * hd + d] = f2bf(acc[i] * inv);
    }
}

// FLASH-DECODING SHAPE: one block per (slot, kv head, chunk); the 64-key K and V tiles are
// staged in SHARED MEMORY ONCE (16 KB each, bf16 as loaded), then every query head of the GQA
// group reads them from smem. That banks the 4x K/V traffic saving WITHOUT the two things that
// sank attn_partial_g: the grid keeps its blocks (grid.y compensates) and no warp performs
// another head's reduction — each (head, key-slice) is a warp's private work, using the same
// 4-rows-per-warp / 8-lanes-per-row layout the base kernel proved.
extern "C" __global__ void attn_partial_f(const float* __restrict__ q,
                                          const unsigned short* __restrict__ kc,
                                          const unsigned short* __restrict__ vc,
                                          float* __restrict__ pm, float* __restrict__ pl,
                                          float* __restrict__ pacc,
                                          const int n_heads, const int n_kv, const int hd,
                                          const int* __restrict__ pos_p, const int nsplit,
                                          const int max_seq, const float scale, const int win) {
    const int gq = n_heads / n_kv;                     // query heads per kv head
    const int m = blockIdx.x / n_kv, kvh = blockIdx.x - m * n_kv;
    const int T = pos_p[m] + 1;
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const long kvbase = ((long)m * n_kv + kvh) * (long)max_seq * hd;
    const int hbase = kvh * gq;
    extern __shared__ float shf[];
    unsigned short* ks = (unsigned short*)shf;             // [ACHUNK][hd] bf16
    unsigned short* vs = ks + ACHUNK * hd;                 // [ACHUNK][hd] bf16
    float* qs = (float*)(vs + ACHUNK * hd);                // [gq][hd]
    float* sc = qs + gq * hd;                              // [gq][ACHUNK]
    float* vred = sc + gq * ACHUNK;                        // [2][gq][hd] warp-pair reduction
    const float* qm = q + (long)m * (n_heads + 2 * n_kv) * hd;
    if ((int)blockIdx.y * ACHUNK >= T) return;
    for (int i = threadIdx.x; i < gq * hd; i += blockDim.x) {
        const int g = i / hd, d = i - g * hd;
        qs[i] = qm[(long)(hbase + g) * hd + d];
    }
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarp = blockDim.x >> 5;
    const int wpg = nwarp / gq;                            // warps per query head (>= 1)
    const int myh = warp / wpg, myw = warp - myh * wpg;    // this warp's head + slice index
    for (int sp = blockIdx.y; sp * ACHUNK < T; sp += (int)gridDim.y) {
        if ((sp + 1) * ACHUNK <= wlo) continue;
        const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
        const int start = max(start0, wlo);
        const int n = end - start;
        __syncthreads();
        // ---- stage K and V for this chunk: 16-byte loads, fully coalesced ----
        for (int i = threadIdx.x; i < n * (hd / 8); i += blockDim.x) {
            const int t = i / (hd / 8), c = i - t * (hd / 8);
            const long g0 = kvbase + (long)(start + t) * hd + c * 8;
            *(uint4*)(ks + (long)t * hd + c * 8) = *(const uint4*)(kc + g0);
            *(uint4*)(vs + (long)t * hd + c * 8) = *(const uint4*)(vc + g0);
        }
        __syncthreads();
        // ---- scores: this warp owns (head myh, keys myw::wpg), 4 keys per iteration ----
        if (myh < gq) {
            const int sub = lane >> 3, sl = lane & 7;
            for (int t = myw * 4; t < n; t += wpg * 4) {
                const int tt = t + sub;
                float d = 0.f;
                if (tt < n) {
                    const unsigned short* kt = ks + (long)tt * hd + sl * 8;
                    const uint4 k0 = *(const uint4*)kt;
                    const uint4 k1 = *(const uint4*)(kt + 64);
                    const float* q0 = qs + myh * hd + sl * 8;
                    const float* q1 = q0 + 64;
                    const unsigned int* w0 = (const unsigned int*)&k0;
                    const unsigned int* w1 = (const unsigned int*)&k1;
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        d = fmaf(q0[2*j],   bf2f((unsigned short)(w0[j] & 0xFFFFu)), d);
                        d = fmaf(q0[2*j+1], bf2f((unsigned short)(w0[j] >> 16)), d);
                        d = fmaf(q1[2*j],   bf2f((unsigned short)(w1[j] & 0xFFFFu)), d);
                        d = fmaf(q1[2*j+1], bf2f((unsigned short)(w1[j] >> 16)), d);
                    }
                }
                #pragma unroll
                for (int o = 4; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
                if (sl == 0 && tt < n) sc[myh * ACHUNK + tt] = d * scale;
            }
        }
        __syncthreads();
        // ---- per-head softmax (one thread per head; ACHUNK is small) ----
        for (int g = threadIdx.x; g < gq; g += blockDim.x) {
            float* scg = sc + g * ACHUNK;
            float mx = -1e30f;
            for (int i = 0; i < n; ++i) mx = fmaxf(mx, scg[i]);
            float sum = 0.f;
            for (int i = 0; i < n; ++i) { const float e2 = __expf(scg[i] - mx); scg[i] = e2; sum += e2; }
            const int gh = m * n_heads + hbase + g;
            pm[(long)gh * nsplit + sp] = mx;
            pl[(long)gh * nsplit + sp] = sum;
        }
        __syncthreads();
        // ---- PV from smem: warp owns (head myh, keys myw::wpg); 32 lanes x 4 dims ----
        if (myh < gq) {
            float4 a4 = make_float4(0.f, 0.f, 0.f, 0.f);
            for (int t = myw; t < n; t += wpg) {
                const float pw = sc[myh * ACHUNK + t];
                const uint2 vv = *(const uint2*)(vs + (long)t * hd + lane * 4);
                a4.x = fmaf(pw, bf2f((unsigned short)(vv.x & 0xFFFFu)), a4.x);
                a4.y = fmaf(pw, bf2f((unsigned short)(vv.x >> 16)), a4.y);
                a4.z = fmaf(pw, bf2f((unsigned short)(vv.y & 0xFFFFu)), a4.z);
                a4.w = fmaf(pw, bf2f((unsigned short)(vv.y >> 16)), a4.w);
            }
            *(float4*)(vred + (myw * gq + myh) * hd + lane * 4) = a4;
        }
        __syncthreads();
        for (int i = threadIdx.x; i < gq * hd; i += blockDim.x) {
            const int g = i / hd, d = i - g * hd;
            float a2 = 0.f;
            for (int w2 = 0; w2 < wpg; ++w2) a2 += vred[(w2 * gq + g) * hd + d];
            const int gh = m * n_heads + hbase + g;
            pacc[((long)gh * nsplit + sp) * hd + d] = a2;
        }
    }
}

// GQA-GROUPED DECODE ATTENTION. attn_partial_b runs one block per (slot, QUERY head), so with
// Llama's 32 q-heads over 8 kv-heads every K/V row is read FOUR times — measured 790 GB/s
// against ~2 TB/s of HBM, the largest non-GEMM term of the step. This kernel runs one block per
// (slot, KV head) and serves all gq = n_heads/n_kv query heads of that group from ONE K/V read.
// Layout of pm/pl/pacc is unchanged (indexed by gh = m*n_heads + h), so the reduce kernels and
// every caller are untouched.
extern "C" __global__ void attn_partial_g(const float* __restrict__ q,
                                          const unsigned short* __restrict__ kc,
                                          const unsigned short* __restrict__ vc,
                                          float* __restrict__ pm, float* __restrict__ pl,
                                          float* __restrict__ pacc,
                                          const int n_heads, const int n_kv, const int hd,
                                          const int* __restrict__ pos_p, const int nsplit,
                                          const int max_seq, const float scale, const int win,
                                          const int hpb) {
    // hpb = query heads served per block (must divide n_heads/n_kv). hpb = n_heads/n_kv shares
    // each K/V row across the whole GQA group but shrinks the grid by the same factor; hpb = 2
    // keeps half the blocks and still halves K/V traffic. Measured: hpb = 4 (full group) LOST
    // 20% on Llama — grid 2048 -> 512 blocks starved the SMs.
    const int gq = hpb;
    const int ngrp = n_heads / hpb;                // blocks per slot
    const int m = blockIdx.x / ngrp, grp = blockIdx.x - m * ngrp;
    const int kvh = (grp * hpb) / (n_heads / n_kv);
    const int hbase = grp * hpb;
    const int T = pos_p[m] + 1;
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const long kvbase = ((long)m * n_kv + kvh) * (long)max_seq * hd;
    extern __shared__ float sh[];
    float* qs = sh;                                // [gq][hd]
    float* sc = qs + gq * hd;                      // [gq][ACHUNK]
    float* vred = sc + gq * ACHUNK;                // [8][gq][hd]
    const float* qm = q + (long)m * (n_heads + 2 * n_kv) * hd;
    if ((int)blockIdx.y * ACHUNK >= T) return;
    for (int i = threadIdx.x; i < gq * hd; i += blockDim.x) {
        const int g = i / hd, d = i - g * hd;
        qs[i] = qm[(long)(hbase + g) * hd + d];
    }
    __syncthreads();
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarp = blockDim.x >> 5;
    for (int sp = blockIdx.y; sp * ACHUNK < T; sp += (int)gridDim.y) {
        if ((sp + 1) * ACHUNK <= wlo) continue;
        const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
        const int start = max(start0, wlo);
        // ---- scores: one K row read serves gq heads ----
        for (int t = start + warp; t < end; t += nwarp) {
            const unsigned short* kt = kc + kvbase + (long)t * hd;
            const uint4 k0 = *reinterpret_cast<const uint4*>(kt + lane * 8);
            float kv[8];
            const unsigned int* kw = (const unsigned int*)&k0;
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                kv[2*j]   = bf2f((unsigned short)(kw[j] & 0xFFFFu));
                kv[2*j+1] = bf2f((unsigned short)(kw[j] >> 16));
            }
            #pragma unroll 1
            for (int g = 0; g < gq; ++g) {
                const float* qg = qs + g * hd + lane * 8;
                float d = 0.f;
                #pragma unroll
                for (int j = 0; j < 8; ++j) d = fmaf(qg[j], kv[j], d);
                #pragma unroll
                for (int o = 16; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
                if (lane == 0) sc[g * ACHUNK + (t - start)] = d * scale;
            }
        }
        __syncthreads();
        // ---- per-head softmax over this chunk ----
        const int n = end - start;
        for (int g = threadIdx.x; g < gq; g += blockDim.x) {
            float* scg = sc + g * ACHUNK;
            float mx = -1e30f;
            for (int i = 0; i < n; ++i) mx = fmaxf(mx, scg[i]);
            float sum = 0.f;
            for (int i = 0; i < n; ++i) { const float e2 = __expf(scg[i] - mx); scg[i] = e2; sum += e2; }
            const int gh = m * n_heads + hbase + g;
            pm[(long)gh * nsplit + sp] = mx;
            pl[(long)gh * nsplit + sp] = sum;
        }
        __syncthreads();
        // ---- PV: one V row read serves gq heads ----
        {
            float4 acc[8];
            #pragma unroll
            for (int g = 0; g < 8; ++g) acc[g] = make_float4(0.f, 0.f, 0.f, 0.f);
            for (int t = start + warp; t < end; t += nwarp) {
                const uint2 vv = *reinterpret_cast<const uint2*>(
                    vc + kvbase + (long)t * hd + lane * 4);
                const float v0 = bf2f((unsigned short)(vv.x & 0xFFFFu));
                const float v1 = bf2f((unsigned short)(vv.x >> 16));
                const float v2 = bf2f((unsigned short)(vv.y & 0xFFFFu));
                const float v3 = bf2f((unsigned short)(vv.y >> 16));
                #pragma unroll 1
                for (int g = 0; g < gq; ++g) {
                    const float pw = sc[g * ACHUNK + (t - start)];
                    acc[g].x = fmaf(pw, v0, acc[g].x);
                    acc[g].y = fmaf(pw, v1, acc[g].y);
                    acc[g].z = fmaf(pw, v2, acc[g].z);
                    acc[g].w = fmaf(pw, v3, acc[g].w);
                }
            }
            for (int g = 0; g < gq; ++g)
                *reinterpret_cast<float4*>(vred + (warp * gq + g) * hd + lane * 4) = acc[g];
            __syncthreads();
            for (int i = threadIdx.x; i < gq * hd; i += blockDim.x) {
                const int g = i / hd, d = i - g * hd;
                float a2 = 0.f;
                for (int w2 = 0; w2 < nwarp; ++w2) a2 += vred[(w2 * gq + g) * hd + d];
                const int gh = m * n_heads + hbase + g;
                pacc[((long)gh * nsplit + sp) * hd + d] = a2;
            }
            __syncthreads();
        }
    }
}

extern "C" __global__ void attn_partial_b(const float* __restrict__ q,
                                          const unsigned short* __restrict__ kc,
                                          const unsigned short* __restrict__ vc,
                                          float* __restrict__ pm, float* __restrict__ pl,
                                          float* __restrict__ pacc,
                                          const int n_heads, const int n_kv, const int hd,
                                          const int* __restrict__ pos_p, const int nsplit,
                                          const int max_seq, const float scale, const int win) {
    const int m = blockIdx.x / n_heads, h = blockIdx.x - m * n_heads;
    const int T = pos_p[m] + 1;          // each slot has its own context length
    // sliding window: keys below `wlo` are invisible. Chunks entirely below it are skipped;
    // the reduce kernels start from the same first-chunk index, so their stale partials are
    // never merged.
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const int kvh = h / (n_heads / n_kv);
    const long kvbase = ((long)m * n_kv + kvh) * (long)max_seq * hd;
    const int gh = m * n_heads + h;

    extern __shared__ float sh[];
    float* qs = sh; float* sc = sh + hd;
    // The qkv scratch rows are (n_heads + 2*n_kv)*hd wide — q heads first. Striding by
    // n_heads*hd here made every slot m >= 1 read its Q out of slot m-1's K/V region: fluent,
    // WRONG output at every M >= 2, unseen for weeks because the HF gate only ever ran
    // sequential requests and fluency was mistaken for correctness at M > 1.
    const float* qm = q + (long)m * (n_heads + 2 * n_kv) * hd;
    // q staged once, reused for every chunk this block owns — and only if there is any.
    if ((int)blockIdx.y * ACHUNK >= T) return;
    for (int i = threadIdx.x; i < hd; i += blockDim.x) qs[i] = qm[(long)h * hd + i];
    __syncthreads();

    for (int sp = blockIdx.y; sp * ACHUNK < T; sp += (int)gridDim.y) {
    if ((sp + 1) * ACHUNK <= wlo) continue;      // whole chunk is beyond the window
    const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
    const int start = max(start0, wlo);
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarp = blockDim.x >> 5;
    if (hd == 128) {
        // MEASURED: this kernel was 498 us of a 6.6 ms M=64 step — the largest non-GEMM term,
        // and only ~38% of HBM peak. One warp per key row issued a single 8-byte load and then
        // stalled on a 5-step shuffle reduction, eight times over: one load in flight per lane
        // against a ten-instruction dependent chain. Four rows per warp-iteration with 16-byte
        // loads puts two loads in flight and cuts the reduction chains 8 -> 2.
        const int sub = lane >> 3;                 // which of the 4 rows this lane serves
        const int sl = lane & 7;                   // lane's slice within that row
        for (int t = start + warp * 4; t < end; t += nwarp * 4) {
            const int tt = t + sub;
            float d = 0.f;
            if (tt < end) {
                const unsigned short* kt = kc + kvbase + (long)tt * hd + sl * 8;
                const uint4 k0 = *reinterpret_cast<const uint4*>(kt);
                const uint4 k1 = *reinterpret_cast<const uint4*>(kt + 64);
                const float* q0 = qs + sl * 8;
                const float* q1 = q0 + 64;
                d = fmaf(q0[0], bf2f((unsigned short)(k0.x & 0xFFFFu)), d);
                d = fmaf(q0[1], bf2f((unsigned short)(k0.x >> 16)), d);
                d = fmaf(q0[2], bf2f((unsigned short)(k0.y & 0xFFFFu)), d);
                d = fmaf(q0[3], bf2f((unsigned short)(k0.y >> 16)), d);
                d = fmaf(q0[4], bf2f((unsigned short)(k0.z & 0xFFFFu)), d);
                d = fmaf(q0[5], bf2f((unsigned short)(k0.z >> 16)), d);
                d = fmaf(q0[6], bf2f((unsigned short)(k0.w & 0xFFFFu)), d);
                d = fmaf(q0[7], bf2f((unsigned short)(k0.w >> 16)), d);
                d = fmaf(q1[0], bf2f((unsigned short)(k1.x & 0xFFFFu)), d);
                d = fmaf(q1[1], bf2f((unsigned short)(k1.x >> 16)), d);
                d = fmaf(q1[2], bf2f((unsigned short)(k1.y & 0xFFFFu)), d);
                d = fmaf(q1[3], bf2f((unsigned short)(k1.y >> 16)), d);
                d = fmaf(q1[4], bf2f((unsigned short)(k1.z & 0xFFFFu)), d);
                d = fmaf(q1[5], bf2f((unsigned short)(k1.z >> 16)), d);
                d = fmaf(q1[6], bf2f((unsigned short)(k1.w & 0xFFFFu)), d);
                d = fmaf(q1[7], bf2f((unsigned short)(k1.w >> 16)), d);
            }
            // 8-lane groups are warp-aligned, so lane sl==0 of each group ends up with its own
            // eight partials only — the cross-group pulls land in lanes whose result is dropped.
            #pragma unroll
            for (int o = 4; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
            if (sl == 0 && tt < end) sc[tt - start] = d * scale;
        }
    } else {
    for (int t = start + warp; t < end; t += nwarp) {
        const unsigned short* kt = kc + kvbase + (long)t * hd;
        float d = 0.f;
        for (int i = lane * 4; i < hd; i += 128) {
            // four bf16 = 8 bytes, half the bytes the f32 float4 moved
            const uint2 kk = *reinterpret_cast<const uint2*>(kt + i);
            d = fmaf(qs[i+0], bf2f((unsigned short)(kk.x & 0xFFFFu)), d);
            d = fmaf(qs[i+1], bf2f((unsigned short)(kk.x >> 16)), d);
            d = fmaf(qs[i+2], bf2f((unsigned short)(kk.y & 0xFFFFu)), d);
            d = fmaf(qs[i+3], bf2f((unsigned short)(kk.y >> 16)), d);
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
        if (lane == 0) sc[t - start] = d * scale;
    }
    }
    __syncthreads();
    __shared__ float red[256];
    const int n = end - start;
    float mxv = -1e30f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) mxv = fmaxf(mxv, sc[i]);
    red[threadIdx.x] = mxv; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x+o]); __syncthreads(); }
    const float mg = red[0]; __syncthreads();
    float sum = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float e2 = __expf(sc[i] - mg); sc[i] = e2; sum += e2; }
    red[threadIdx.x] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x+o]; __syncthreads(); }
    const float ls = red[0]; __syncthreads();
    if (hd == 128) {
        __shared__ float vred[8 * 128];
        const int g = threadIdx.x >> 5;
        float4 a4 = make_float4(0.f,0.f,0.f,0.f);
        for (int t = start + g; t < end; t += 8) {
            const float pw = sc[t - start];
            const uint2 vv = *reinterpret_cast<const uint2*>(vc + kvbase + (long)t*hd + lane*4);
            a4.x = fmaf(pw, bf2f((unsigned short)(vv.x & 0xFFFFu)), a4.x);
            a4.y = fmaf(pw, bf2f((unsigned short)(vv.x >> 16)), a4.y);
            a4.z = fmaf(pw, bf2f((unsigned short)(vv.y & 0xFFFFu)), a4.z);
            a4.w = fmaf(pw, bf2f((unsigned short)(vv.y >> 16)), a4.w);
        }
        *reinterpret_cast<float4*>(vred + g*128 + lane*4) = a4;
        __syncthreads();
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float a2 = 0.f;
            #pragma unroll
            for (int gg = 0; gg < 8; ++gg) a2 += vred[gg*128 + d];
            pacc[((long)gh * nsplit + sp) * hd + d] = a2;
        }
    } else {
        // generic head_dim (Gemma4: 256 sliding / 512 global): each thread owns output dims
        // d, d+256, ... and walks the chunk serially. Bandwidth-fine; the fancy path stays 128.
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float a2 = 0.f;
            for (int t = start; t < end; ++t)
                a2 = fmaf(sc[t - start],
                          bf2f(vc[kvbase + (long)t*hd + d]), a2);
            pacc[((long)gh * nsplit + sp) * hd + d] = a2;
        }
    }
    if (threadIdx.x == 0) { pm[gh * nsplit + sp] = mg; pl[gh * nsplit + sp] = ls; }
    __syncthreads();     // sc/red are reused by the next chunk this block strides to
    }
}

// bf16 twins of the fa mma helpers, for kernels whose operands come from the bf16 KV slab.
#define FA_MMAB(d0,d1,d2,d3, a0,a1,a2,a3, b0,b1) \
    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " \
                 "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" \
                 : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) \
                 : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1))

__device__ __forceinline__ unsigned int fa_pack2bf(float a, float b) {
    unsigned int r;
    asm("{ .reg .b16 lo, hi; cvt.rn.bf16.f32 lo, %1; cvt.rn.bf16.f32 hi, %2; mov.b32 %0, {lo,hi}; }"
        : "=r"(r) : "f"(a), "f"(b));
    return r;
}

// ldmatrix with per-8x8 transpose: the B operand for O = P.V read straight from a ROW-MAJOR
// [key][dim] tile (the decode path has no pre-transposed V; the prefill path built vt16 instead).
__device__ __forceinline__ void fa_ldm4t(unsigned int& d0, unsigned int& d1, unsigned int& d2,
                                         unsigned int& d3, const unsigned short* smem) {
    unsigned int s;
    asm("{ .reg .u64 p; cvta.to.shared.u64 p, %1; cvt.u32.u64 %0, p; }" : "=r"(s) : "l"(smem));
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) : "r"(s));
}

// TENSOR-CORE split-K decode attention (ATTNT=1). ncu on attn_partial_b: SM 58%, DRAM 25%,
// L2 hit ~49%, 525 GB/s effective — the kernel is INSTRUCTION-ISSUE bound (~2.2 scalar
// dequant+FFMA instructions per bf16 element ~= the whole 3.5 ms). mma fragments cut issue
// count ~8x; the head-major slab planes are exactly the contiguous tiles cp.async wants.
// Shape: one WARP owns one 64-key ACHUNK end to end (stage K/V, S = Q.K^T, single-tile
// softmax — no online rescale, the chunk IS the tile — P.V, scatter per-chunk partials).
// The pm/pl/pacc contract and the reduce are unchanged. hd==128, bf16 slab only.
extern "C" __global__ void attn_partial_t(const float* __restrict__ q,
                                          const unsigned short* __restrict__ kc,
                                          const unsigned short* __restrict__ vc,
                                          float* __restrict__ pm, float* __restrict__ pl,
                                          float* __restrict__ pacc,
                                          const int n_heads, const int n_kv, const int hd,
                                          const int* __restrict__ pos_p, const int nsplit,
                                          const int max_seq, const float scale, const int win) {
#if __CUDA_ARCH__ >= 800
    const int m2 = blockIdx.x / n_kv, kvh = blockIdx.x - m2 * n_kv;
    const int T = pos_p[m2] + 1;
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const int gq = n_heads / n_kv;
    const int hbase = kvh * gq;
    const long kvbase = ((long)m2 * n_kv + kvh) * (long)max_seq * hd;
    const int lane = threadIdx.x & 31, w = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    const int gid = lane >> 2, t4 = lane & 3;
    const int lrow = 8 * ((lane >> 3) & 1) + (lane & 7);
    const int lcol = 8 * (lane >> 4);
    const int QS = 136, KS = 136;   // +8 halves of row padding: bank-conflict-free ldmatrix

    extern __shared__ unsigned short tsh[];
    unsigned short* sQ = tsh;                              // [16][QS]
    unsigned short* sK = sQ + 16 * QS + w * (64 * KS);     // this warp's K tile
    unsigned short* sV = sQ + 16 * QS + nwarp * (64 * KS) + w * (64 * KS);

    // Q: gq real rows (bf16, pre-scaled), pad rows zeroed. One block-wide stage.
    for (int i = threadIdx.x; i < 16 * (hd / 8); i += blockDim.x) {
        const int r = i / (hd / 8), c = i - r * (hd / 8);
        unsigned int sw[4];
        if (r < gq) {
            const float* qr = q + (long)m2 * (n_heads + 2 * n_kv) * hd + (long)(hbase + r) * hd + c * 8;
            #pragma unroll
            for (int j = 0; j < 4; ++j) sw[j] = fa_pack2bf(qr[2 * j] * scale, qr[2 * j + 1] * scale);
        } else {
            sw[0] = 0u; sw[1] = 0u; sw[2] = 0u; sw[3] = 0u;
        }
        *(uint4*)(sQ + r * QS + c * 8) = *(const uint4*)sw;
    }
    __syncthreads();

    for (int sp = blockIdx.y * nwarp + w; sp * ACHUNK < T; sp += (int)gridDim.y * nwarp) {
        if ((sp + 1) * ACHUNK <= wlo) continue;
        const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
        const int start = max(start0, wlo);
        const int n = end - start;
        // ---- stage this warp's K and V tiles (16-byte cp.async, zero-filled past n) ----
        for (int i = lane; i < 64 * (hd / 8); i += 32) {
            const int r = i / (hd / 8), c = i - r * (hd / 8);
            const unsigned sz = (r < n) ? 16u : 0u;
            fa_cp16z(sK + r * KS + c * 8, kc + kvbase + (long)(start + r) * hd + c * 8, sz);
            fa_cp16z(sV + r * KS + c * 8, vc + kvbase + (long)(start + r) * hd + c * 8, sz);
        }
        FA_CP_COMMIT();
        asm volatile("cp.async.wait_group 0;\n" ::);
        __syncwarp();

        // ---- S = Q . K^T: 8 n-tiles of 8 keys, hd/16 k-steps (fa's exact K addressing) ----
        float sc[8][4];
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) { sc[nt][0] = 0.f; sc[nt][1] = 0.f; sc[nt][2] = 0.f; sc[nt][3] = 0.f; }
        #pragma unroll
        for (int ks = 0; ks < 8; ++ks) {
            unsigned int a0, a1, a2, a3;
            fa_ldm4(a0, a1, a2, a3, sQ + (long)lrow * QS + ks * 16 + lcol);
            #pragma unroll
            for (int np = 0; np < 4; ++np) {
                unsigned int b0, b1, b2, b3;
                fa_ldm4(b0, b1, b2, b3, sK + (long)(16 * np + lrow) * KS + ks * 16 + lcol);
                FA_MMAB(sc[2 * np][0], sc[2 * np][1], sc[2 * np][2], sc[2 * np][3],
                        a0, a1, a2, a3, b0, b2);
                FA_MMAB(sc[2 * np + 1][0], sc[2 * np + 1][1], sc[2 * np + 1][2], sc[2 * np + 1][3],
                        a0, a1, a2, a3, b1, b3);
            }
        }
        // ---- boundary mask + single-tile softmax (rows = query heads; pad rows are garbage
        //      that is simply never written back) ----
        float mx0 = -1e30f, mx1 = -1e30f;
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int col = 8 * nt + t4 * 2 + j;
                float v0 = sc[nt][j], v1 = sc[nt][2 + j];
                if (col >= n) { v0 = -1e30f; v1 = -1e30f; }
                sc[nt][j] = v0; sc[nt][2 + j] = v1;
                mx0 = fmaxf(mx0, v0); mx1 = fmaxf(mx1, v1);
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, x));
            mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, x));
        }
        float s0 = 0.f, s1 = 0.f;
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const float e0 = (sc[nt][j] > -1e29f) ? __expf(sc[nt][j] - mx0) : 0.f;
                const float e1 = (sc[nt][2 + j] > -1e29f) ? __expf(sc[nt][2 + j] - mx1) : 0.f;
                sc[nt][j] = e0; sc[nt][2 + j] = e1;
                s0 += e0; s1 += e1;
            }
        }
        #pragma unroll
        for (int x = 1; x < 4; x <<= 1) {
            s0 += __shfl_xor_sync(0xffffffffu, s0, x);
            s1 += __shfl_xor_sync(0xffffffffu, s1, x);
        }
        // ---- O = P . V, V read straight from the row-major tile via trans-ldmatrix ----
        float o[16][4];
        #pragma unroll
        for (int nt = 0; nt < 16; ++nt) { o[nt][0] = 0.f; o[nt][1] = 0.f; o[nt][2] = 0.f; o[nt][3] = 0.f; }
        #pragma unroll
        for (int ks = 0; ks < 4; ++ks) {
            const unsigned int a0 = fa_pack2bf(sc[2 * ks][0],     sc[2 * ks][1]);
            const unsigned int a1 = fa_pack2bf(sc[2 * ks][2],     sc[2 * ks][3]);
            const unsigned int a2 = fa_pack2bf(sc[2 * ks + 1][0], sc[2 * ks + 1][1]);
            const unsigned int a3 = fa_pack2bf(sc[2 * ks + 1][2], sc[2 * ks + 1][3]);
            #pragma unroll
            for (int np = 0; np < 8; ++np) {
                unsigned int b0, b1, b2, b3;
                fa_ldm4t(b0, b1, b2, b3, sV + (long)(16 * ks + lrow) * KS + np * 16 + lcol);
                // .trans pairing: source matrices are (k 0-7|k 8-15|k 0-7,d 8-15|k 8-15,d 8-15);
                // transposed, (d0,d1) become dims 0-7 x both k-halves and (d2,d3) dims 8-15 —
                // NOT the plain-load (b0,b2)/(b1,b3) pairing the fa PV uses on pre-transposed V.
                FA_MMAB(o[2 * np][0], o[2 * np][1], o[2 * np][2], o[2 * np][3],
                        a0, a1, a2, a3, b0, b1);
                FA_MMAB(o[2 * np + 1][0], o[2 * np + 1][1], o[2 * np + 1][2], o[2 * np + 1][3],
                        a0, a1, a2, a3, b2, b3);
            }
        }
        // ---- scatter: rows gid / gid+8 that are REAL heads write their per-chunk partials ----
        if (gid < gq) {
            const int gh = m2 * n_heads + hbase + gid;
            #pragma unroll
            for (int nt = 0; nt < 16; ++nt) {
                const int d = 8 * nt + t4 * 2;
                pacc[((long)gh * nsplit + sp) * hd + d]     = o[nt][0];
                pacc[((long)gh * nsplit + sp) * hd + d + 1] = o[nt][1];
            }
            if (t4 == 0) { pm[gh * nsplit + sp] = mx0; pl[gh * nsplit + sp] = s0; }
        }
        if (gid + 8 < gq) {
            const int gh = m2 * n_heads + hbase + gid + 8;
            #pragma unroll
            for (int nt = 0; nt < 16; ++nt) {
                const int d = 8 * nt + t4 * 2;
                pacc[((long)gh * nsplit + sp) * hd + d]     = o[nt][2];
                pacc[((long)gh * nsplit + sp) * hd + d + 1] = o[nt][3];
            }
            if (t4 == 0) { pm[gh * nsplit + sp] = mx1; pl[gh * nsplit + sp] = s1; }
        }
        __syncwarp();   // sK/sV are this warp's own; re-staged next iteration
    }
#endif
}

// WARP-PRIVATE split-K decode attention (ATTNW=1 to enable). attn_partial_b measures 0.68 TB/s
// at m=8/2k and stayed flat through the Y-sweep AND the head-major relayout — the remaining
// suspect is its own structure: each 64-key chunk is load-K -> ~18 block-wide __syncthreads of
// softmax reduction -> load-V -> barrier, so a block spends most of its life with no loads in
// flight. Here each WARP owns whole chunks: scores in a per-warp smem slice, max/sum by warp
// shuffle, V accumulated by the same warp, partials written per (head, chunk) exactly as the
// reduce already expects. ZERO block barriers in the loop; __syncwarp only.
// hd==128 only (the dispatch guards); output is the same pm/pl/pacc contract as attn_partial_b.
extern "C" __global__ void attn_partial_w(const float* __restrict__ q,
                                          const unsigned short* __restrict__ kc,
                                          const unsigned short* __restrict__ vc,
                                          float* __restrict__ pm, float* __restrict__ pl,
                                          float* __restrict__ pacc,
                                          const int n_heads, const int n_kv, const int hd,
                                          const int* __restrict__ pos_p, const int nsplit,
                                          const int max_seq, const float scale, const int win) {
    const int m = blockIdx.x / n_heads, h = blockIdx.x - m * n_heads;
    const int T = pos_p[m] + 1;
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const int kvh = h / (n_heads / n_kv);
    const long kvbase = ((long)m * n_kv + kvh) * (long)max_seq * hd;
    const int gh = m * n_heads + h;
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarp = blockDim.x >> 5;

    extern __shared__ float sh[];
    float* qs = sh;                        // [hd] — the one shared read-only stage
    float* scw = sh + hd + warp * ACHUNK;  // this warp's private score slice
    for (int i = threadIdx.x; i < hd; i += blockDim.x) {
        qs[i] = q[(long)m * (n_heads + 2 * n_kv) * hd + (long)h * hd + i];
    }
    __syncthreads();                       // the only block barrier: qs visibility

    for (int sp = blockIdx.y * nwarp + warp; sp * ACHUNK < T; sp += (int)gridDim.y * nwarp) {
        if ((sp + 1) * ACHUNK <= wlo) continue;
        const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
        const int start = max(start0, wlo);
        // ---- K phase: 4 rows per warp-iteration, 8-lane slices, 16-byte loads ----
        const int sub = lane >> 3, sl = lane & 7;
        for (int t = start; t < end; t += 4) {
            const int tt = t + sub;
            float d = 0.f;
            if (tt < end) {
                const unsigned short* kt = kc + kvbase + (long)tt * hd + sl * 8;
                const uint4 k0 = *reinterpret_cast<const uint4*>(kt);
                const uint4 k1 = *reinterpret_cast<const uint4*>(kt + 64);
                const float* q0 = qs + sl * 8;
                const float* q1 = q0 + 64;
                d = fmaf(q0[0], bf2f((unsigned short)(k0.x & 0xFFFFu)), d);
                d = fmaf(q0[1], bf2f((unsigned short)(k0.x >> 16)), d);
                d = fmaf(q0[2], bf2f((unsigned short)(k0.y & 0xFFFFu)), d);
                d = fmaf(q0[3], bf2f((unsigned short)(k0.y >> 16)), d);
                d = fmaf(q0[4], bf2f((unsigned short)(k0.z & 0xFFFFu)), d);
                d = fmaf(q0[5], bf2f((unsigned short)(k0.z >> 16)), d);
                d = fmaf(q0[6], bf2f((unsigned short)(k0.w & 0xFFFFu)), d);
                d = fmaf(q0[7], bf2f((unsigned short)(k0.w >> 16)), d);
                d = fmaf(q1[0], bf2f((unsigned short)(k1.x & 0xFFFFu)), d);
                d = fmaf(q1[1], bf2f((unsigned short)(k1.x >> 16)), d);
                d = fmaf(q1[2], bf2f((unsigned short)(k1.y & 0xFFFFu)), d);
                d = fmaf(q1[3], bf2f((unsigned short)(k1.y >> 16)), d);
                d = fmaf(q1[4], bf2f((unsigned short)(k1.z & 0xFFFFu)), d);
                d = fmaf(q1[5], bf2f((unsigned short)(k1.z >> 16)), d);
                d = fmaf(q1[6], bf2f((unsigned short)(k1.w & 0xFFFFu)), d);
                d = fmaf(q1[7], bf2f((unsigned short)(k1.w >> 16)), d);
            }
            #pragma unroll
            for (int o = 4; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
            if (sl == 0 && tt < end) scw[tt - start] = d * scale;
        }
        __syncwarp();
        // ---- warp-local softmax over this chunk ----
        const int n = end - start;
        float mx = -1e30f;
        for (int i = lane; i < n; i += 32) mx = fmaxf(mx, scw[i]);
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
        float sum = 0.f;
        for (int i = lane; i < n; i += 32) { const float e = __expf(scw[i] - mx); scw[i] = e; sum += e; }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
        __syncwarp();
        // ---- V phase: this warp walks its own chunk; a lane owns 4 output dims ----
        float4 a4 = make_float4(0.f, 0.f, 0.f, 0.f);
        const unsigned short* vt0 = vc + kvbase + (long)start * hd + lane * 4;
        for (int t = 0; t < n; ++t) {
            const float pw = scw[t];
            const uint2 vv = *reinterpret_cast<const uint2*>(vt0 + (long)t * hd);
            a4.x = fmaf(pw, bf2f((unsigned short)(vv.x & 0xFFFFu)), a4.x);
            a4.y = fmaf(pw, bf2f((unsigned short)(vv.x >> 16)), a4.y);
            a4.z = fmaf(pw, bf2f((unsigned short)(vv.y & 0xFFFFu)), a4.z);
            a4.w = fmaf(pw, bf2f((unsigned short)(vv.y >> 16)), a4.w);
        }
        *reinterpret_cast<float4*>(pacc + ((long)gh * nsplit + sp) * hd + lane * 4) = a4;
        if (lane == 0) {
            pm[gh * nsplit + sp] = mx;
            pl[gh * nsplit + sp] = (sum > 0.f) ? sum : 0.f;
        }
    }
}

// fp8 (e4m3) KV twin of attn_partial_b: same tiling, softmax and partials; K/V loads move HALF
// the bytes and dequantize through the hardware e4m3x2 -> f16x2 conversion. All index arithmetic
// is element-count arithmetic and the slab holds one byte per element, so the formulas carry over
// with u8 pointers.
extern "C" __global__ void attn_partial_b_f8(const float* __restrict__ q,
                                             const unsigned char* __restrict__ kc,
                                             const unsigned char* __restrict__ vc,
                                             float* __restrict__ pm, float* __restrict__ pl,
                                             float* __restrict__ pacc,
                                             const int n_heads, const int n_kv, const int hd,
                                             const int* __restrict__ pos_p, const int nsplit,
                                             const int max_seq, const float scale, const int win) {
    const int m = blockIdx.x / n_heads, h = blockIdx.x - m * n_heads;
    const int T = pos_p[m] + 1;
    const int wlo = (win > 0 && T > win) ? (T - win) : 0;
    const int kvh = h / (n_heads / n_kv);
    const long kvbase = ((long)m * n_kv + kvh) * (long)max_seq * hd;
    const int gh = m * n_heads + h;

    extern __shared__ float sh[];
    float* qs = sh; float* sc = sh + hd;
    const float* qm = q + (long)m * (n_heads + 2 * n_kv) * hd;
    if ((int)blockIdx.y * ACHUNK >= T) return;
    for (int i = threadIdx.x; i < hd; i += blockDim.x) qs[i] = qm[(long)h * hd + i];
    __syncthreads();

    for (int sp = blockIdx.y; sp * ACHUNK < T; sp += (int)gridDim.y) {
    if ((sp + 1) * ACHUNK <= wlo) continue;
    const int start0 = sp * ACHUNK, end = min(start0 + ACHUNK, T);
    const int start = max(start0, wlo);
    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarp = blockDim.x >> 5;
    if (hd == 128) {
        const int sub = lane >> 3;
        const int sl = lane & 7;
        for (int t = start + warp * 4; t < end; t += nwarp * 4) {
            const int tt = t + sub;
            float d = 0.f;
            if (tt < end) {
                const unsigned char* kt = kc + kvbase + (long)tt * hd + sl * 8;
                const uint2 k0 = *reinterpret_cast<const uint2*>(kt);
                const uint2 k1 = *reinterpret_cast<const uint2*>(kt + 64);
                const float* q0 = qs + sl * 8;
                const float* q1 = q0 + 64;
                float2 p0 = e4m3x2_2f((unsigned short)(k0.x & 0xFFFFu));
                float2 p1 = e4m3x2_2f((unsigned short)(k0.x >> 16));
                float2 p2 = e4m3x2_2f((unsigned short)(k0.y & 0xFFFFu));
                float2 p3 = e4m3x2_2f((unsigned short)(k0.y >> 16));
                d = fmaf(q0[0], p0.x, d); d = fmaf(q0[1], p0.y, d);
                d = fmaf(q0[2], p1.x, d); d = fmaf(q0[3], p1.y, d);
                d = fmaf(q0[4], p2.x, d); d = fmaf(q0[5], p2.y, d);
                d = fmaf(q0[6], p3.x, d); d = fmaf(q0[7], p3.y, d);
                p0 = e4m3x2_2f((unsigned short)(k1.x & 0xFFFFu));
                p1 = e4m3x2_2f((unsigned short)(k1.x >> 16));
                p2 = e4m3x2_2f((unsigned short)(k1.y & 0xFFFFu));
                p3 = e4m3x2_2f((unsigned short)(k1.y >> 16));
                d = fmaf(q1[0], p0.x, d); d = fmaf(q1[1], p0.y, d);
                d = fmaf(q1[2], p1.x, d); d = fmaf(q1[3], p1.y, d);
                d = fmaf(q1[4], p2.x, d); d = fmaf(q1[5], p2.y, d);
                d = fmaf(q1[6], p3.x, d); d = fmaf(q1[7], p3.y, d);
            }
            #pragma unroll
            for (int o = 4; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
            if (sl == 0 && tt < end) sc[tt - start] = d * scale;
        }
    } else {
    for (int t = start + warp; t < end; t += nwarp) {
        const unsigned char* kt = kc + kvbase + (long)t * hd;
        float d = 0.f;
        for (int i = lane * 4; i < hd; i += 128) {
            const unsigned int kk = *reinterpret_cast<const unsigned int*>(kt + i);
            const float2 p0 = e4m3x2_2f((unsigned short)(kk & 0xFFFFu));
            const float2 p1 = e4m3x2_2f((unsigned short)(kk >> 16));
            d = fmaf(qs[i+0], p0.x, d);
            d = fmaf(qs[i+1], p0.y, d);
            d = fmaf(qs[i+2], p1.x, d);
            d = fmaf(qs[i+3], p1.y, d);
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) d += __shfl_down_sync(0xffffffffu, d, o);
        if (lane == 0) sc[t - start] = d * scale;
    }
    }
    __syncthreads();
    __shared__ float red[256];
    const int n = end - start;
    float mxv = -1e30f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) mxv = fmaxf(mxv, sc[i]);
    red[threadIdx.x] = mxv; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x+o]); __syncthreads(); }
    const float mg = red[0]; __syncthreads();
    float sum = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) { const float e2 = __expf(sc[i] - mg); sc[i] = e2; sum += e2; }
    red[threadIdx.x] = sum; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) { if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x+o]; __syncthreads(); }
    const float ls = red[0]; __syncthreads();
    if (hd == 128) {
        __shared__ float vred[8 * 128];
        const int g = threadIdx.x >> 5;
        float4 a4 = make_float4(0.f,0.f,0.f,0.f);
        for (int t = start + g; t < end; t += 8) {
            const float pw = sc[t - start];
            const unsigned int vv = *reinterpret_cast<const unsigned int*>(vc + kvbase + (long)t*hd + lane*4);
            const float2 p0 = e4m3x2_2f((unsigned short)(vv & 0xFFFFu));
            const float2 p1 = e4m3x2_2f((unsigned short)(vv >> 16));
            a4.x = fmaf(pw, p0.x, a4.x);
            a4.y = fmaf(pw, p0.y, a4.y);
            a4.z = fmaf(pw, p1.x, a4.z);
            a4.w = fmaf(pw, p1.y, a4.w);
        }
        *reinterpret_cast<float4*>(vred + g*128 + lane*4) = a4;
        __syncthreads();
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float a2 = 0.f;
            #pragma unroll
            for (int gg = 0; gg < 8; ++gg) a2 += vred[gg*128 + d];
            pacc[((long)gh * nsplit + sp) * hd + d] = a2;
        }
    } else {
        for (int d = threadIdx.x; d < hd; d += blockDim.x) {
            float a2 = 0.f;
            for (int t = start; t < end; ++t) {
                const float kv = e4m3x2_2f(vc[kvbase + (long)t*hd + d]).x;
                a2 = fmaf(sc[t - start], kv, a2);
            }
            pacc[((long)gh * nsplit + sp) * hd + d] = a2;
        }
    }
    if (threadIdx.x == 0) { pm[gh * nsplit + sp] = mg; pl[gh * nsplit + sp] = ls; }
    __syncthreads();
    }
}

// As attn_reduce_b but emitting bf16. Both GEMM arms feed the o-projection a bf16 operand, so an
// f32 attention output only gets re-read and converted — one extra pass over m*n_heads*hd per layer,
// plus a kernel launch, on every step.
//
// v2: at 2k context this kernel measured 0.77 ms/step (Llama) — the old shape recomputed
// exp(pm - mg) PER THREAD PER CHUNK (128 x 32 exps) behind a serial thread-0 max/sum head.
// Warp 0 now computes the chunk weights ONCE into shared (parallel max + one exp per chunk),
// and the d-loop is pure coalesced loads + fma. Numerically identical reduction order per d.
extern "C" __global__ void attn_reduce_b16(const float* __restrict__ pm, const float* __restrict__ pl,
                                           const float* __restrict__ pacc,
                                           unsigned short* __restrict__ out,
                                           const int hd, const int nsplit, const int n_heads,
                                           const int* __restrict__ pos_p, const int win) {
    const int gh = blockIdx.x;
    const int T = pos_p[gh / n_heads] + 1;
    const int used = min(nsplit, (T - 1 + ACHUNK) / ACHUNK);
    // chunks wholly below the sliding window were skipped by attn_partial_b — their partials
    // are stale and must not be merged
    const int first = (win > 0 && T > win) ? (T - win) / ACHUNK : 0;
    __shared__ float wts[64];        // max_seq <= 4096 => nsplit <= 64
    __shared__ float lg;
    if (threadIdx.x < 32) {
        const int lane = threadIdx.x;
        float mv = -1e30f;
        for (int s2 = first + lane; s2 < used; s2 += 32) mv = fmaxf(mv, pm[gh*nsplit+s2]);
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) mv = fmaxf(mv, __shfl_xor_sync(0xffffffffu, mv, o));
        float l = 0.f;
        for (int s2 = first + lane; s2 < used; s2 += 32) {
            const float w = __expf(pm[gh*nsplit+s2] - mv);
            wts[s2] = w;
            l += pl[gh*nsplit+s2] * w;
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) l += __shfl_xor_sync(0xffffffffu, l, o);
        if (lane == 0) lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    // a / lg, NOT a * (1/lg): the reciprocal-multiply rounds differently and this kernel's
    // outputs must stay bit-identical to the v1 reduce (same weights, same fma order).
    for (int d = threadIdx.x; d < hd; d += blockDim.x) {
        float a = 0.f;
        for (int s2 = first; s2 < used; ++s2)
            a = fmaf(wts[s2], pacc[((long)gh*nsplit+s2)*hd + d], a);
        out[(long)gh * hd + d] = f2bf(a / lg);
    }
}

extern "C" __global__ void attn_reduce_b(const float* __restrict__ pm, const float* __restrict__ pl,
                                         const float* __restrict__ pacc, float* __restrict__ out,
                                         const int hd, const int nsplit, const int n_heads,
                                         const int* __restrict__ pos_p, const int win) {
    const int gh = blockIdx.x;
    // Bound by THIS slot's context, not by nsplit: nsplit is sized from max_seq, so a server
    // configured for 4096 was merging 16 chunks per head on a 23-token context.
    const int T = pos_p[gh / n_heads] + 1;
    const int used = min(nsplit, (T - 1 + ACHUNK) / ACHUNK);
    const int first = (win > 0 && T > win) ? (T - win) / ACHUNK : 0;
    __shared__ float mg, lg;
    if (threadIdx.x == 0) {
        float mv = -1e30f;
        for (int s2 = first; s2 < used; ++s2) mv = fmaxf(mv, pm[gh*nsplit+s2]);
        float l = 0.f;
        for (int s2 = first; s2 < used; ++s2) l += pl[gh*nsplit+s2] * __expf(pm[gh*nsplit+s2] - mv);
        mg = mv; lg = (l > 0.f) ? l : 1.f;
    }
    __syncthreads();
    for (int d = threadIdx.x; d < hd; d += blockDim.x) {
        float a = 0.f;
        for (int s2 = first; s2 < used; ++s2)
            a = fmaf(__expf(pm[gh*nsplit+s2] - mg), pacc[((long)gh*nsplit+s2)*hd + d], a);
        out[(long)gh * hd + d] = a / lg;
    }
}

// Gemma4 v_norm: SCALELESS RMSNorm over each V head, in place, in the f32 qkv scratch — BEFORE
// the rope kernel snapshots V into the cache/f16 copy. On global layers the V region holds the
// raw k_proj output (the loader fused q|k|K), so this one kernel covers both layer types.
// One block per (row, v-head); rows = T (prefill) or M (decode).
extern "C" __global__ void v_derive(float* __restrict__ qkv, const int n_heads, const int n_kv,
                                    const int hd, const float eps, const int rows) {
    const int r = blockIdx.x / n_kv, vh = blockIdx.x - r * n_kv;
    if (r >= rows) return;
    const int per = n_heads + 2 * n_kv;
    float* v = qkv + (long)r * per * hd + (long)(n_heads + n_kv + vh) * hd;
    __shared__ float red[128];
    float s = 0.f;
    for (int i = threadIdx.x; i < hd; i += blockDim.x) s = fmaf(v[i], v[i], s);
    red[threadIdx.x] = s; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) red[threadIdx.x] += red[threadIdx.x + o];
        __syncthreads();
    }
    const float inv = rsqrtf(red[0] / (float)hd + eps);
    for (int i = threadIdx.x; i < hd; i += blockDim.x) v[i] *= inv;
}

// Gemma4 sandwich step: x = (x + rms(h)·postw) · ls, then out16 = bf16(rms(x)·prew) for the next
// GEMM. The norm sits on the BRANCH OUTPUT before the residual add (pre-norm models norm the
// STREAM after adding — rmsnorm_add_bf), and ls is the layer_scalar multiplying the whole
// stream at layer end (1.0 mid-layer). One block per row; rows = T (prefill) or M (decode).
extern "C" __global__ void sandwich_bf(float* __restrict__ x, const float* __restrict__ h,
                                       const float* __restrict__ postw,
                                       const float* __restrict__ prew,
                                       unsigned short* __restrict__ out,
                                       const int n, const float eps, const float ls) {
    const int m = blockIdx.x;
    float* xm = x + (long)m * n;
    const float* hm = h + (long)m * n;
    unsigned short* om = out + (long)m * n;
    // Warp-shuffle sums, as in rmsnorm_add_bf. This kernel is launched with the 1024-thread
    // config, so the old red[256] smem tree wrote out of bounds for every thread above 255 —
    // on Gemma 4's live path, where sandwich norms run twice a layer.
    __shared__ float red[32];
    const int lane = threadIdx.x & 31, wid = threadIdx.x >> 5, nwarp = blockDim.x >> 5;
    float s = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) s = fmaf(hm[i], hm[i], s);
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
    if (lane == 0) red[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float hinv = rsqrtf(red[0] / (float)n + eps);
    __syncthreads();
    float s2 = 0.f;
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        const float v = (xm[i] + hm[i] * hinv * postw[i]) * ls;
        xm[i] = v;
        s2 = fmaf(v, v, s2);
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s2 += __shfl_down_sync(0xffffffffu, s2, o);
    if (lane == 0) red[wid] = s2;
    __syncthreads();
    if (wid == 0) {
        float t = (lane < nwarp) ? red[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) t += __shfl_down_sync(0xffffffffu, t, o);
        if (lane == 0) red[0] = t;
    }
    __syncthreads();
    const float xinv = rsqrtf(red[0] / (float)n + eps);
    for (int i = threadIdx.x; i < n; i += blockDim.x) om[i] = f2bf(xm[i] * xinv * prew[i]);
}

// MEASURED: this was 338 us of a 6.6 ms decode step at M=64 — 39 MB of logits read at
// 115 GB/s, 26x off HBM. 64 blocks x 256 threads is 16k threads chewing 9.7M floats, each
// walking ~600 elements serially with nothing to hide the latency. 1024 threads + float4
// gives 16x the in-flight loads; the tie-break rule (lowest index wins) is preserved exactly
// so sampled tokens are unchanged.
extern "C" __global__ void argmax_b(const float* __restrict__ v, unsigned int* __restrict__ out,
                                    const int n) {
    const int m = blockIdx.x;
    const float* vm = v + (long)m * n;
    __shared__ float bv[1024];
    __shared__ unsigned int bi[1024];
    float best = -1e30f; unsigned int bidx = 0u;
    const int n4 = n >> 2;
    const float4* v4 = reinterpret_cast<const float4*>(vm);
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        const float4 q = v4[i];
        const unsigned int b0 = (unsigned int)(i << 2);
        if (q.x > best) { best = q.x; bidx = b0; }
        if (q.y > best) { best = q.y; bidx = b0 + 1u; }
        if (q.z > best) { best = q.z; bidx = b0 + 2u; }
        if (q.w > best) { best = q.w; bidx = b0 + 3u; }
    }
    for (int i = (n4 << 2) + (int)threadIdx.x; i < n; i += blockDim.x)
        if (vm[i] > best) { best = vm[i]; bidx = (unsigned int)i; }
    bv[threadIdx.x] = best; bi[threadIdx.x] = bidx; __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o) {
            // strict > keeps the LOWEST index on ties, matching the 256-thread version
            if (bv[threadIdx.x+o] > bv[threadIdx.x]) { bv[threadIdx.x] = bv[threadIdx.x+o]; bi[threadIdx.x] = bi[threadIdx.x+o]; }
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) out[m] = bi[0];
}

// Per-slot record + advance. `active[m] == 0` freezes a slot (finished or unoccupied) so the rest of
// the batch keeps stepping — the other half of continuous batching.
extern "C" __global__ void advance_b(int* __restrict__ pos, const unsigned int* __restrict__ tok,
                                     unsigned int* __restrict__ out, const int max_out, const int M,
                                     const unsigned int* __restrict__ active) {
    const int m = blockIdx.x * blockDim.x + threadIdx.x;
    if (m >= M) return;
    if (active && !active[m]) return;
    const int p = pos[m];
    if (p < max_out) out[(long)m * max_out + p] = tok[m];
    pos[m] = p + 1;
}

extern "C" __global__ void silu_mul(float* __restrict__ g, const float* __restrict__ u, const int n,
                                    const int act) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) { const float v = g[i]; g[i] = ffn_act(v, act) * u[i]; }
}

extern "C" __global__ void add_inplace(float* __restrict__ a, const float* __restrict__ b, const int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) a[i] += b[i];
}

// MEASURED: 326 us of a 2.56 ms M=1 decode step — 12.7% of it, in the single-row twin of the
// kernel already widened for the batched path. Same treatment: 1024 threads and float4 loads,
// and the same strict-> tie-break so the sampled token is bit-identical.
extern "C" __global__ void argmax_f32(const float* __restrict__ v, unsigned int* __restrict__ out,
                                      const int n) {
    __shared__ float bv[1024];
    __shared__ unsigned int bi[1024];
    float best = -1e30f; unsigned int bidx = 0;
    const int n4 = n >> 2;
    const float4* v4 = reinterpret_cast<const float4*>(v);
    for (int i = threadIdx.x; i < n4; i += blockDim.x) {
        const float4 q = v4[i];
        const unsigned int b0 = (unsigned int)(i << 2);
        if (q.x > best) { best = q.x; bidx = b0; }
        if (q.y > best) { best = q.y; bidx = b0 + 1u; }
        if (q.z > best) { best = q.z; bidx = b0 + 2u; }
        if (q.w > best) { best = q.w; bidx = b0 + 3u; }
    }
    for (int i = (n4 << 2) + (int)threadIdx.x; i < n; i += blockDim.x)
        if (v[i] > best) { best = v[i]; bidx = (unsigned int)i; }
    bv[threadIdx.x] = best; bi[threadIdx.x] = bidx;
    __syncthreads();
    for (int o = blockDim.x >> 1; o > 0; o >>= 1) {
        if ((int)threadIdx.x < o && bv[threadIdx.x + o] > bv[threadIdx.x]) {
            bv[threadIdx.x] = bv[threadIdx.x + o];
            bi[threadIdx.x] = bi[threadIdx.x + o];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) *out = bi[0];
}
"#;