memra_engine/

decode.rs

//! Incremental decode (T=1) with the dual cache + greedy generation loop. Serves end-to-end.
//! Reuses the validated kernels; threads KV (full-attn) and conv/SSM state (linear-attn) across steps.

use crate::cache::{Cache, RecurLayer};
use crate::forward::argmax;
use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer};
use crate::Engine;
use cudarc::driver::CudaSlice;
use std::collections::HashMap;

/// Persistent CUDA-graph decode state (CUDA-GRAPH-PLAN Phase 3). Holds the device-resident counters
/// the captured graph reads/writes (`token_d` = current/next token id, `pos_d` = rope position) — both
/// at FIXED addresses baked into every captured graph — plus the per-`t_kv`-bucket graph cache. The
/// bucket key is the eager `(fa_vec, n_splits)` pair (see `Engine::fa_bucket_key`): every t_kv that
/// maps to the same key reproduces eager's split geometry, so one captured graph replays bit-identically
/// for the whole bucket. A new key triggers a re-capture (n_splits changes ~every 64 tokens).
pub struct GraphDecodeState {
    pub token_d: CudaSlice<u32>, // [1] resident next-token id (argmax writes, embed reads)
    pub pos_d: CudaSlice<i32>,   // [1] resident rope position counter
    pub graphs: HashMap<(bool, usize), cudarc::driver::CudaGraph>,
    pub bucket_max: HashMap<(bool, usize), usize>, // bucket key -> bucket_max fed to the capture
    pub captures: usize,                           // count of (re)captures, for reporting
}

/// Long-lived step-wise CUDA-graph decode session (see HybridModel::graph_session_new).
/// One replay per step(); the only steady-state D2H is the 4-byte next-token read.
pub struct GraphSession {
    pub gs: GraphDecodeState,
    pub cache: Cache,
    /// LOAD-BEARING hold: the captured graph's embed-gather node references this
    /// allocation — dropping it would free memory the graph still reads.
    #[allow(dead_code)]
    embd_gpu: CudaSlice<u8>,
    graph: cudarc::driver::CudaGraph,
    plan: Vec<crate::graph_update::FaMain>,
    /// session budget: last valid t_kv (pos + max_new + 1 at creation).
    pub bucket_max: usize,
    /// current capture's kernel-class segment end — step() recaptures past it
    /// (round 45: exec-update retunes splits, it cannot swap kernels; see
    /// graph_decode_loop's SEGMENTS note).
    seg_end: usize,
    qt: i32,
    row_bytes: usize,
    n_vocab: usize,
    /// GRAMMAR MASK (constrained decoding, 2026-08-03): packed llguidance bitset the
    /// captured graph reads (mask_logits_f32 between lm_head and the in-graph argmax).
    /// STABLE POINTER — baked at capture, carried across recaptures; the caller uploads
    /// fresh contents (upload_mask) before every step. None = no mask node captured.
    mask_dev: Option<CudaSlice<u32>>,
    mask_words: usize,
}

impl GraphSession {
    /// One graph-replay decode step. Returns the next token (already fed back into the
    /// resident token_d — the following step consumes it). Errors past bucket_max
    /// (the caller sized max_new at capture). Transparently recaptures when the eager
    /// kernel class changes (fa_vec floor / v4 max / fa512 floor crossings).
    pub fn step(&mut self, e: &Engine, m: &crate::hybrid::HybridModel)
                -> Result<u32, Box<dyn std::error::Error>> {
        if self.cache.pos + 1 >= self.bucket_max {
            return Err("GraphSession: past bucket_max (generation budget exceeded)".into());
        }
        if self.cache.pos + 1 > self.seg_end {
            m.graph_session_recapture(e, self)?;
        }
        crate::graph_update::fa_apply(&self.graph, &mut self.plan, self.cache.pos + 1,
                                      crate::fa_split_keys)?;
        self.graph.launch()?;
        self.cache.pos += 1;
        for kvl in self.cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
            kvl.len += 1;
        }
        e.dtoh_u32_one(&self.gs.token_d)
    }

    /// GRAMMAR MASK upload (constrained graph sessions): fresh packed-bitset contents into
    /// the STABLE buffer the captured graph reads — call before every step(). The word
    /// count is a capture-time kernel arg (constant per model: the tokenizer vocab is
    /// fixed), so the length must match the capture exactly.
    pub fn upload_mask(&mut self, e: &Engine, words: &[u32])
                       -> Result<(), Box<dyn std::error::Error>> {
        let Some(d) = self.mask_dev.as_mut() else {
            return Err("upload_mask: session captured without a mask node".into());
        };
        if words.len() != self.mask_words {
            return Err(format!("upload_mask: {} words != captured {}",
                               words.len(), self.mask_words).into());
        }
        e.htod_u32_into(d, words)
    }

    /// Profiling decomposition of step() (graph-session-gate MEMRA_GS_PROF): the three
    /// phases exposed separately. prof_launch is ASYNC (no sync) — prof_read carries the
    /// sync+D2H. Advances the session exactly like step().
    pub fn prof_apply(&mut self, _e: &Engine) -> Result<(), Box<dyn std::error::Error>> {
        crate::graph_update::fa_apply(&self.graph, &mut self.plan, self.cache.pos + 1,
                                      crate::fa_split_keys)
    }
    pub fn prof_launch(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        self.graph.launch()?;
        self.cache.pos += 1;
        for kvl in self.cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
            kvl.len += 1;
        }
        Ok(())
    }
    pub fn prof_read(&mut self, e: &Engine) -> Result<u32, Box<dyn std::error::Error>> {
        e.dtoh_u32_one(&self.gs.token_d)
    }
}

impl GraphDecodeState {
    pub fn new(e: &Engine) -> Result<Self, Box<dyn std::error::Error>> {
        Ok(GraphDecodeState {
            token_d: e.stream().clone_htod(&[0u32])?,
            pos_d: e.htod_i32(&[0])?,
            graphs: HashMap::new(),
            bucket_max: HashMap::new(),
            captures: 0,
        })
    }
}

/// Generation parameters for the reusable serving API (`generate_with`).
#[derive(Clone, Debug)]
pub struct GenParams {
    pub max_new: usize,         // hard cap on generated tokens
    pub max_ctx: Option<usize>, // context-length guard; None => prompt+max_new+8
    pub eos: Vec<u32>,          // stop on any of these token ids (eos/eog + specials)
}
impl Default for GenParams {
    fn default() -> Self {
        GenParams {
            max_new: 128,
            max_ctx: None,
            eos: Vec::new(),
        }
    }
}

/// Why generation stopped.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StopReason {
    Eos,
    MaxNew,
    ContextFull,
    Callback,
}

/// Result of `generate_with`: the generated token ids + why it stopped.
pub struct GenOutput {
    pub tokens: Vec<u32>,
    pub stop_reason: StopReason,
}

/// Diagnostic-only snapshots of Hy3 layer 0 in the eager T=1 serving path.
/// Each buffer is one residual-width device row captured before the next stage can reuse it.
pub struct Hy3Layer0Stages {
    pub attention_output: CudaSlice<f32>,
    pub after_attention: CudaSlice<f32>,
    pub mlp_output: CudaSlice<f32>,
    pub residual: CudaSlice<f32>,
}

impl HybridModel {
    /// Device embed table for the dc fast loops (lazy ~0.5GB upload). On OOM — tight fits
    /// where resident experts + KV leave no headroom (35B ct-NVFP4 artifact at default
    /// budget, 2026-07-17) — returns None and the caller stays on the host-embd eager loop
    /// instead of panicking. Double-init race is benign (identical bytes, loser dropped).
    fn embd_gpu_try(&self, e: &Engine) -> Option<&cudarc::driver::CudaSlice<u8>> {
        if let Some(v) = self.embd_gpu.get() {
            return Some(v);
        }
        match e.upload_u8(&self.embd.raw) {
            Ok(buf) => Some(self.embd_gpu.get_or_init(|| buf)),
            Err(err) => {
                eprintln!("[embd-gpu] upload failed ({err}); dc loop disabled, host-embd eager loop serves");
                None
            }
        }
    }
}

impl HybridModel {
    /// One decode step for `token` at cache.pos; returns logits [n_vocab] (host f32). Advances cache.
    pub fn decode_step(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
        Ok(self.decode_step_h(e, token, cache)?.0)
    }

    /// Dense-FFN SwiGLU (T=1 decode): `down @ (silu(gate@z) * (up@z))`. Two fused levers stack here:
    ///  - RANK3 LEVER 2: gate+up NVFP4 macro-scales fold into ONE `silu_mul_scaled*` launch (via
    ///    `matmul_pre_noscale`), saving the two separate `scale_inplace` launches.
    ///  - RANK2 LEVER (q8_1 quant-fold): when ffn_down is ALSO on the q8_1 fast path, the SwiGLU
    ///    epilogue EMITS the q8_1 quantization of `act` directly (`silu_mul_scaled_q8_1`) and feeds
    ///    ffn_down via `matmul_pre`, removing ffn_down's standalone `quantize_q8_1` launch (the
    ///    down-proj activation has one consumer, so the quant folds into its producer for free).
    /// BIT-IDENTICAL to matmul_pre(gate)+matmul_pre(up)+silu_mul+quantize_q8_1+matmul(down): same
    /// float silu*mul, same amax/127 q8_1 rounding, same dp4a/mmvq dot. Falls back to the f32 `act`
    /// + plain matmul(down) path whenever any of the three is off the fast path.
    fn ffn_swiglu_decode(
        &self,
        e: &Engine,
        ffn_gate: &crate::model::GpuTensor,
        ffn_up: &crate::model::GpuTensor,
        ffn_down: &crate::model::GpuTensor,
        z: &CudaSlice<f32>,
        n_embd: usize,
        n_ff: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        // M3 dense layers use swigluoai (clamped) — the silu_mul fused fast paths below encode
        // plain SiLU; route through ffn_act (macro-scales folded via matmul_pre) until clamped
        // fused twins exist.
        if self.cfg.m3.is_some() {
            let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
            let gate = e.matmul_pre(ffn_gate, &zq, &zd, z, 1)?;
            let up = e.matmul_pre(ffn_up, &zq, &zd, z, 1)?;
            let mut act = e.uninit(n_ff)?;
            Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
            return Ok(e.matmul(ffn_down, &act, 1)?);
        }
        if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
            let (zq, zd) = e.quantize_q8_1(z, 1, n_embd)?;
            // DUAL mm-fusion first (NVFP4 gate+up in ONE launch), else two noscale launches.
            let pair = match e.matmul_pre_dual_noscale(ffn_gate, ffn_up, &zq, &zd, 1)? {
                Some((g, u)) => (Some(g), Some(u)),
                None => (
                    e.matmul_pre_noscale(ffn_gate, &zq, &zd, 1)?,
                    e.matmul_pre_noscale(ffn_up, &zq, &zd, 1)?,
                ),
            };
            match pair {
                (Some((gate, gs)), Some((up, us))) => {
                    // RANK2 fold: if ffn_down is q8_1-fast, emit act PRE-QUANTIZED and skip the
                    // standalone quantize_q8_1 before ffn_down.
                    if e.uses_q8_1_fast(ffn_down) {
                        let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff)?;
                        return Ok(e.matmul_pre(
                            ffn_down, &aq, &ad, /*x_fallback unused on fast path*/ &gate, 1,
                        )?);
                    }
                    let mut act = e.uninit(n_ff)?;
                    e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff)?;
                    return Ok(e.matmul(ffn_down, &act, 1)?);
                }
                _ => {
                    // one (or both) not on the separable-scale fast path: scaled matmul + plain silu_mul.
                    let gate = e.matmul_pre(ffn_gate, &zq, &zd, z, 1)?;
                    let up = e.matmul_pre(ffn_up, &zq, &zd, z, 1)?;
                    let mut act = e.uninit(n_ff)?;
                    Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
                    return Ok(e.matmul(ffn_down, &act, 1)?);
                }
            }
        }
        let gate = e.matmul(ffn_gate, z, 1)?;
        let up = e.matmul(ffn_up, z, 1)?;
        let mut act = e.uninit(n_ff)?;
        Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
        Ok(e.matmul(ffn_down, &act, 1)?)
    }

    /// Like `ffn_swiglu_decode` but the input is ALREADY q8_1-quantized `(zq, zd)` — used by the
    /// DECODE NORM-FUSION lever where `add_rms_norm_q8_1` emits the post-attn-normed activation
    /// pre-quantized (no f32 `z` materialized, no standalone quantize_q8_1 launch). Caller GUARANTEES
    /// ffn_gate and ffn_up are q8_1-fast (so `matmul_pre_noscale` returns Some at m=1). BIT-IDENTICAL
    /// to ffn_swiglu_decode(z) when (zq,zd) == quantize_q8_1(z): same matmul_pre_noscale, same
    /// silu_mul_scaled_q8_1 / silu_mul_scaled, same ffn_down dot.
    fn ffn_swiglu_decode_pre(
        &self,
        e: &Engine,
        ffn_gate: &crate::model::GpuTensor,
        ffn_up: &crate::model::GpuTensor,
        ffn_down: &crate::model::GpuTensor,
        zq: &CudaSlice<i8>,
        zd: &CudaSlice<f32>,
        n_ff: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let pair = match e.matmul_pre_dual_noscale(ffn_gate, ffn_up, zq, zd, 1)? {
            Some((g, u)) => (Some(g), Some(u)),
            None => (
                e.matmul_pre_noscale(ffn_gate, zq, zd, 1)?,
                e.matmul_pre_noscale(ffn_up, zq, zd, 1)?,
            ),
        };
        match pair {
            (Some((gate, gs)), Some((up, us))) => {
                if e.uses_q8_1_fast(ffn_down) {
                    let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, n_ff)?;
                    Ok(e.matmul_pre(ffn_down, &aq, &ad, &gate, 1)?)
                } else {
                    let mut act = e.uninit(n_ff)?;
                    e.silu_mul_scaled(&gate, &up, gs, us, &mut act, n_ff)?;
                    Ok(e.matmul(ffn_down, &act, 1)?)
                }
            }
            // Unreachable when the caller's q8_1-fast guarantee holds (m==1 + fast => Some). Guard
            // anyway: re-quant from the dequantized pair would need f32; surface a clear error.
            _ => Err("ffn_swiglu_decode_pre: gate/up not separable-scale at m=1 (caller must guarantee q8_1-fast)".into()),
        }
    }

    /// Shared post-attention residual + post-attn-norm + FFN for ONE decode layer, routed by ALL
    /// decode loops (eager + dc + dc_cap) so they stay bit-identical by construction. DECODE
    /// NORM-FUSION LEVER: when the layer is Dense AND ffn_gate/ffn_up are q8_1-fast (the daily NVFP4
    /// case), fuses residual-add + post_attn_norm + q8_1-quantize into ONE `add_rms_norm_q8_1` launch
    /// and feeds the FFN the pre-quantized activation (skipping its internal quantize_q8_1) — removing
    /// 1-2 launches + the f32 `z` HBM round-trip per layer. BIT-IDENTICAL to the unfused
    /// add_rms_norm(or add+rms_norm) + quantize_q8_1 + ffn (all proven bit-identical in kernel_check).
    /// MEMRA_NO_FUSE_NORMQ forces the unfused f32 path. Returns (x1 residual f32, ffn_out f32).
    /// True when ALL of a mixer's input projections are on the q8_1 fast path (so the attn-input
    /// rms_norm can emit q8_1 directly and the mixer skips its internal quantize_q8_1).
    pub(crate) fn mixer_in_q8_1_fast(&self, e: &Engine, mixer: &Mixer) -> bool {
        match mixer {
            Mixer::Full(fa) => {
                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv)
            }
            Mixer::Linear(la) => {
                e.uses_q8_1_fast(&la.wqkv)
                    && e.uses_q8_1_fast(&la.wqkv_gate)
                    && e.uses_q8_1_fast(&la.ssm_beta)
                    && e.uses_q8_1_fast(&la.ssm_alpha)
            }
            // MLA (increment 2, loader-only): predicate only — never claim the fused
            // norm+quantize chain for an arm that has no forward yet.
            Mixer::Mla(_) => false,
        }
    }

    /// attn_norm + mixer for the EAGER loop, with the attn-input NORM-FUSION. MEMRA_NO_FUSE_NORMQ
    /// forces the unfused (separate rms_norm + mixer-internal quantize) path.
    fn attn_in_norm_mixer(
        &self,
        e: &Engine,
        layer: &crate::hybrid::HybridLayer,
        x: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        pos: usize,
        cache: &mut Cache,
        il: usize,
        n_embd: usize,
        eps: f32,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let anorm = layer.attn_norm.float_data();
        let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
            && self.mixer_in_q8_1_fast(e, &layer.mixer);
        if fuse {
            let (hq, hd) = e.rms_norm_q8_1(x, anorm, n_embd, 1, eps)?;
            // h is unused on the fast path (matmul_pre x_fallback only used at m>=16); pass a zero-len.
            let h0 = e.zeros(0)?;
            match &layer.mixer {
                Mixer::Full(fa) => {
                    self.full_attn_decode_pre(e, fa, &h0, Some((&hq, &hd)), pos_d, pos, cache, il)
                }
                Mixer::Linear(la) => {
                    self.linear_attn_decode_pre(e, la, &h0, &hq, &hd, cache, il, false)
                }
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        } else {
            let mut h = e.uninit(n_embd)?;
            e.rms_norm(x, anorm, &mut h, n_embd, 1, eps)?;
            match &layer.mixer {
                Mixer::Full(fa) => self.full_attn_decode(e, fa, &h, pos_d, pos, cache, il),
                Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il),
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        }
    }

    /// attn_norm + mixer for the DEVICE-COUNTER loop (decode_step_dc). Full-attn uses the dc path;
    /// linear uses the eager-state path (persistent=false), same as decode_step_dc. NORM-FUSED.
    fn attn_in_norm_mixer_dc(
        &self,
        e: &Engine,
        layer: &crate::hybrid::HybridLayer,
        x: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
        n_embd: usize,
        eps: f32,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let anorm = layer.attn_norm.float_data();
        let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
            && self.mixer_in_q8_1_fast(e, &layer.mixer);
        if fuse {
            let (hq, hd) = e.rms_norm_q8_1(x, anorm, n_embd, 1, eps)?;
            let h0 = e.zeros(0)?;
            match &layer.mixer {
                Mixer::Full(fa) => {
                    self.full_attn_decode_dc_pre(e, fa, &h0, &hq, &hd, pos_d, cache, il)
                }
                Mixer::Linear(la) => {
                    self.linear_attn_decode_pre(e, la, &h0, &hq, &hd, cache, il, false)
                }
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        } else {
            let mut h = e.uninit(n_embd)?;
            e.rms_norm(x, anorm, &mut h, n_embd, 1, eps)?;
            match &layer.mixer {
                Mixer::Full(fa) => self.full_attn_decode_dc(e, fa, &h, pos_d, cache, il),
                Mixer::Linear(la) => self.linear_attn_decode(e, la, &h, cache, il),
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        }
    }

    /// attn_norm + mixer for the CAPTURE loop (decode_step_dc_cap). Full-attn uses the dc_cap path
    /// (fixed bucket_max); linear uses the persistent-state path. NORM-FUSED; capture-safe (rms_norm_q8_1
    /// + the *_pre mixers enqueue the same kernels every replay, stable buffers).
    fn attn_in_norm_mixer_dc_cap(
        &self,
        e: &Engine,
        layer: &crate::hybrid::HybridLayer,
        x: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
        bucket_max: usize,
        n_embd: usize,
        eps: f32,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let anorm = layer.attn_norm.float_data();
        let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
            && self.mixer_in_q8_1_fast(e, &layer.mixer);
        if fuse {
            let (hq, hd) = e.rms_norm_q8_1(x, anorm, n_embd, 1, eps)?;
            let h0 = e.zeros(0)?;
            match &layer.mixer {
                Mixer::Full(fa) => self.full_attn_decode_dc_cap_pre(
                    e, fa, &h0, &hq, &hd, pos_d, cache, il, bucket_max,
                ),
                Mixer::Linear(la) => {
                    self.linear_attn_decode_pre(e, la, &h0, &hq, &hd, cache, il, true)
                }
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        } else {
            let mut h = e.uninit(n_embd)?;
            e.rms_norm(x, anorm, &mut h, n_embd, 1, eps)?;
            match &layer.mixer {
                Mixer::Full(fa) => {
                    self.full_attn_decode_dc_cap(e, fa, &h, pos_d, cache, il, bucket_max)
                }
                Mixer::Linear(la) => self.linear_attn_decode_cap(e, la, &h, cache, il),
                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
            }
        }
    }

    fn residual_norm_ffn(
        &self,
        e: &Engine,
        layer: &crate::hybrid::HybridLayer,
        x: &CudaSlice<f32>,
        mixed: &CudaSlice<f32>,
        n_embd: usize,
        il: usize,
        eps: f32,
    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let pnorm = layer.post_attn_norm.float_data();
        match &layer.ffn {
            crate::hybrid::Ffn::Dense {
                ffn_gate,
                ffn_up,
                ffn_down,
            } => {
                let n_ff = ffn_gate.out_features();
                // cfg.m3: the fused-pre chain's silu_mul_scaled* epilogues are plain SiLU —
                // M3's swigluoai must route through ffn_swiglu_decode's m3 arm (FAST-gate
                // MISMATCH root cause #2, 2026-07-07: L0 dense FFN clamp skipped under FAST).
                let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
                    && self.cfg.m3.is_none()
                    && e.uses_q8_1_fast(ffn_gate)
                    && e.uses_q8_1_fast(ffn_up);
                if fuse {
                    let mut x1 = e.uninit(n_embd)?;
                    let (zq, zd) = e.add_rms_norm_q8_1(x, mixed, pnorm, &mut x1, n_embd, 1, eps)?;
                    let ffn_out =
                        self.ffn_swiglu_decode_pre(e, ffn_gate, ffn_up, ffn_down, &zq, &zd, n_ff)?;
                    Ok((x1, ffn_out))
                } else {
                    let mut x1 = e.uninit(n_embd)?;
                    let mut z = e.uninit(n_embd)?;
                    e.add_rms_norm(x, mixed, pnorm, &mut x1, &mut z, n_embd, 1, eps)?;
                    let ffn_out =
                        self.ffn_swiglu_decode(e, ffn_gate, ffn_up, ffn_down, &z, n_embd, n_ff)?;
                    Ok((x1, ffn_out))
                }
            }
            crate::hybrid::Ffn::Moe(m) => {
                let mut x1 = e.uninit(n_embd)?;
                let mut z = e.uninit(n_embd)?;
                // z-quantize fuse (add_rms_norm_zq8) measured NEGATIVE here (158.8 vs 160.6:
                // the fused warp-per-block quantize pass re-reads z slower than the dedicated
                // coalesced quantize_q8_1). Kernel + threading kept for graph-capture use where
                // launch count matters more; eager default = unfused (no gain = no change).
                e.add_rms_norm(x, mixed, pnorm, &mut x1, &mut z, n_embd, 1, eps)?;
                let ffn_out = self.moe_ffn_il_zq8(e, m, &z, None, 1, il as u16)?;
                Ok((x1, ffn_out))
            }
        }
    }

    /// EAGLE3 aux-hidden capture (EAGLE-PLAN N1): one decode step that ALSO returns the trunk
    /// residual-stream `x` taken AFTER each of the blocks in `aux_layers` (the EAGLE3 encoder feeds
    /// these 3 layer hiddens through `fc`). Returns (logits[n_vocab] host, aux: Vec<[n_embd] dev>),
    /// one device buffer per requested aux layer, in `aux_layers` order. The captured tensor is the
    /// residual `x` produced by that block (`x2` at the loop tail), cloned before the next block
    /// overwrites it — cheap (one clone_dtod of [n_embd] per aux layer). T=1 decode regime.
    pub fn decode_step_aux(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
        aux_layers: &[usize],
    ) -> Result<(Vec<f32>, Vec<CudaSlice<f32>>), Box<dyn std::error::Error>> {
        let (logits, aux, _) = self.decode_step_aux_inner(e, token, cache, aux_layers, false)?;
        Ok((logits, aux))
    }

    /// Diagnostic-only Hy3 layer-0 trace through the real eager T=1 serving path. Besides the
    /// final block residual, this captures the attention output before its residual add, the
    /// after-attention residual, and the dense-MLP output before the final residual add.
    pub fn decode_step_hy3_layer0_stages(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
    ) -> Result<(Vec<f32>, Hy3Layer0Stages), Box<dyn std::error::Error>> {
        if self.cfg.hy3.is_none() {
            return Err("decode_step_hy3_layer0_stages requires a Hy3 model".into());
        }
        if !matches!(
            self.layers.first().map(|layer| &layer.ffn),
            Some(crate::hybrid::Ffn::Dense { .. })
        ) {
            return Err("Hy3 diagnostic expected layer 0 to use a dense MLP".into());
        }
        let (logits, _, stages) = self.decode_step_aux_inner(e, token, cache, &[], true)?;
        Ok((
            logits,
            stages.ok_or("Hy3 layer-0 stages were not captured")?,
        ))
    }

    fn decode_step_aux_inner(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
        aux_layers: &[usize],
        capture_hy3_layer0: bool,
    ) -> Result<(Vec<f32>, Vec<CudaSlice<f32>>, Option<Hy3Layer0Stages>), Box<dyn std::error::Error>>
    {
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let pos = cache.pos;
        let pos_d = e.htod_i32(&[pos as i32])?;

        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;
        let mut aux: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
        let mut hy3_layer0 = None;

        for (il, layer) in self.layers.iter().enumerate() {
            // attn-input NORM-FUSION (eager); shared with decode_step_h.
            let mixed =
                self.attn_in_norm_mixer(e, layer, &x, &pos_d, pos, cache, il, n_embd, eps)?;
            // DECODE NORM-FUSION LEVER (residual_norm_ffn): residual add + post_attn RMSNorm +
            // q8_1-quantize fused into ONE add_rms_norm_q8_1 launch on the Dense q8_1-fast path, then
            // the FFN consumes the pre-quantized activation. Bit-identical to the unfused path.
            let (x1, ffn_out) = self.residual_norm_ffn(e, layer, &x, &mixed, n_embd, il, eps)?;
            let mut x2 = e.uninit(n_embd)?;
            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
            if capture_hy3_layer0 && il == 0 {
                hy3_layer0 = Some(Hy3Layer0Stages {
                    attention_output: e.clone_dtod(&mixed)?,
                    after_attention: e.clone_dtod(&x1)?,
                    mlp_output: e.clone_dtod(&ffn_out)?,
                    residual: e.clone_dtod(&x2)?,
                });
            }
            // EAGLE3 N1: capture this block's residual output if it is an aux layer.
            if aux_layers.contains(&il) {
                aux.push(e.clone_dtod(&x2)?);
            }
            x = x2;
        }
        // re-order aux to match aux_layers order (contains() pushes in il order; aux_layers is the
        // canonical order the encoder concats in — they coincide since aux_layers is ascending).
        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let logits = e.matmul(&self.output, &hn, 1)?;
        let host = e.dtoh(&logits)?;
        cache.pos += 1;
        Ok((host, aux, hy3_layer0))
    }

    /// Like `decode_step`, but ALSO returns the trunk's hidden state `x` taken BEFORE the final
    /// `output_norm` (MTP-PLAN §A: this is `h_seed` for the NextN head). Device buffer [n_embd].
    pub fn decode_step_h(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        if self.is_gemma4_e4b() {
            crate::pp::warn_unwired_once("gemma4-e4b eager decode");
            return self.gemma4_e4b_decode_step_h(e, token, cache);
        }
        if self.cfg.gemma4.is_some() {
            // pp2 door for the gemma4 arm lives inside gemma4_decode_step_h.
            return self.gemma4_decode_step_h(e, token, cache);
        }
        // M2 ppN door (crate::pp): N-stage split of this walk with an explicit activation
        // handoff at each boundary. Default OFF — unset env means this branch never taken.
        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
            return self.decode_step_h_ppn(e, token, cache, &fence);
        }
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let pos = cache.pos;
        let pos_d = e.htod_i32(&[pos as i32])?;

        // embed the single token -> [1, n_embd]
        let mut x = e.htod(&self.embd.gather(n_embd, &[token]))?;

        // CROSS-LAYER ADD+NORM FUSION (launch-arc 2026-07-07): layer il's post-FFN residual add
        // (x2 = x1 + ffn_out) and layer il+1's attn_norm+quantize are consecutive row-wise ops —
        // add_rms_norm_q8_1 does all three in ONE launch (bit-identity proven in kernel_check:
        // add_rms_norm == add then rms_norm; _q8_1 == then quantize_q8_1). Carry the un-added
        // (x1, ffn_out) pair into the next iteration; the fused launch materializes x2 (the
        // residual this layer needs) as its `res` output. Falls back to the separate add when
        // the next mixer is off the q8_1 fast path.
        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
        for (il, layer) in self.layers.iter().enumerate() {
            let anorm = layer.attn_norm.float_data();
            let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
                && self.mixer_in_q8_1_fast(e, &layer.mixer);
            // NOTE: take() FIRST, branch on fuse after — a tuple pattern like
            // `if let (Some(p), true) = (pending.take(), fuse)` DROPS the taken pair when
            // fuse is false (pattern fails post-take) and silently loses the residual add.
            let taken = pending.take();
            let mixed = match (taken, fuse) {
                (Some((x1, f1)), true) => {
                    // fused add + attn_norm + q8_1 (this layer's mixer input), res -> x2
                    let mut x2 = e.uninit(n_embd)?;
                    let (hq, hd) = e.add_rms_norm_q8_1(&x1, &f1, anorm, &mut x2, n_embd, 1, eps)?;
                    x = x2;
                    let h0 = e.zeros(0)?;
                    match &layer.mixer {
                        Mixer::Full(fa) => self.full_attn_decode_pre(
                            e,
                            fa,
                            &h0,
                            Some((&hq, &hd)),
                            &pos_d,
                            pos,
                            cache,
                            il,
                        )?,
                        Mixer::Linear(la) => {
                            self.linear_attn_decode_pre(e, la, &h0, &hq, &hd, cache, il, false)?
                        }
                        Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
                    }
                }
                (taken, _) => {
                    if let Some((x1, f1)) = taken {
                        let mut x2 = e.uninit(n_embd)?;
                        e.add(&x1, &f1, &mut x2, n_embd)?;
                        x = x2;
                    }
                    self.attn_in_norm_mixer(e, layer, &x, &pos_d, pos, cache, il, n_embd, eps)?
                }
            };

            // DECODE NORM-FUSION LEVER (residual_norm_ffn): add+post_attn_norm+q8_1 fused on the Dense
            // fast path. Bit-identical to add + rms_norm + ffn (add_rms_norm == add then rms_norm,
            // proven in kernel_check; add_rms_norm_q8_1 == add_rms_norm then quantize_q8_1).
            let (x1, ffn_out) = self.residual_norm_ffn(e, layer, &x, &mixed, n_embd, il, eps)?;
            pending = Some((x1, ffn_out));
        }
        // final layer's add (no next norm to fuse with — output_norm is f32-out)
        if let Some((x1, f1)) = pending.take() {
            let mut x2 = e.uninit(n_embd)?;
            e.add(&x1, &f1, &mut x2, n_embd)?;
            x = x2;
        }

        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        // h_seed = trunk hidden BEFORE output_norm (default, §A) or AFTER it (MEMRA_SPEC_HPOST,
        // the reference engines' convention — see spec::spec_hpost).
        let h_seed = if crate::spec::spec_hpost() {
            e.clone_dtod(&hn)?
        } else {
            e.clone_dtod(&x)?
        };
        // head-MIPS feasibility probe (MEMRA_DUMP_HN=<path>): append pre-head hiddens for
        // offline bound analysis. Diagnostic only.
        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
            let hh = e.dtoh(&hn)?;
            use std::io::Write;
            let mut fo = std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(path)?;
            for v in &hh {
                fo.write_all(&v.to_le_bytes())?;
            }
        }
        let logits = e.matmul(&self.output, &hn, 1)?;
        let host = e.dtoh(&logits)?;
        cache.pos += 1;
        Ok((host, h_seed))
    }

    /// M1-PP2 stage subgraph: run layers [lo, hi) of the generic eager walk. Enters with a
    /// MATERIALIZED residual `x` (no pending fusion pair from outside the range) and exits
    /// with the range's final residual materialized (the trailing add executed, exactly like
    /// the last layer of an unsplit walk). Body is the `decode_step_h` loop verbatim with the
    /// cross-layer add+norm fusion carry LOCAL to the range — so the only state a stage
    /// boundary has to move is the [n_embd] hidden state. Bit-identity of the cut relies on
    /// the kernel-check-pinned `add_rms_norm_q8_1 == add then rms_norm_q8_1` identity
    /// (`pp2-gate` verifies end-to-end on real weights).
    #[allow(clippy::too_many_arguments)]
    fn decode_layers_eager(
        &self,
        e: &Engine,
        mut x: CudaSlice<f32>,
        lo: usize,
        hi: usize,
        pos_d: &CudaSlice<i32>,
        pos: usize,
        cache: &mut Cache,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let n_embd = self.cfg.n_embd as usize;
        let eps = self.cfg.rms_eps;
        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
        for il in lo..hi {
            let layer = &self.layers[il];
            let anorm = layer.attn_norm.float_data();
            let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
                && self.mixer_in_q8_1_fast(e, &layer.mixer);
            // take() FIRST, branch on fuse after (see decode_step_h: a tuple pattern drops
            // the taken pair when fuse is false and silently loses the residual add).
            let taken = pending.take();
            let mixed = match (taken, fuse) {
                (Some((x1, f1)), true) => {
                    let mut x2 = e.uninit(n_embd)?;
                    let (hq, hd) = e.add_rms_norm_q8_1(&x1, &f1, anorm, &mut x2, n_embd, 1, eps)?;
                    x = x2;
                    let h0 = e.zeros(0)?;
                    match &layer.mixer {
                        Mixer::Full(fa) => self.full_attn_decode_pre(
                            e,
                            fa,
                            &h0,
                            Some((&hq, &hd)),
                            pos_d,
                            pos,
                            cache,
                            il,
                        )?,
                        Mixer::Linear(la) => {
                            self.linear_attn_decode_pre(e, la, &h0, &hq, &hd, cache, il, false)?
                        }
                        Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
                    }
                }
                (taken, _) => {
                    if let Some((x1, f1)) = taken {
                        let mut x2 = e.uninit(n_embd)?;
                        e.add(&x1, &f1, &mut x2, n_embd)?;
                        x = x2;
                    }
                    self.attn_in_norm_mixer(e, layer, &x, pos_d, pos, cache, il, n_embd, eps)?
                }
            };
            let (x1, ffn_out) = self.residual_norm_ffn(e, layer, &x, &mixed, n_embd, il, eps)?;
            pending = Some((x1, ffn_out));
        }
        // range's final add (no next norm inside the range to fuse with)
        if let Some((x1, f1)) = pending.take() {
            let mut x2 = e.uninit(n_embd)?;
            e.add(&x1, &f1, &mut x2, n_embd)?;
            x = x2;
        }
        Ok(x)
    }

    /// M2: `decode_step_h` as N stage subgraphs, each on ITS OWN CUDA stream (and, under
    /// MEMRA_PP_DEVICES, its own device/engine), with the transport-selected boundary
    /// handoff at each fence cut. Stage 0 = embed + its layer range; each middle stage
    /// RXes boundary s-1 (waits its ev_tx), runs its range, TXes boundary s; the last
    /// stage adds output_norm + lm head. Per-layer KV/linear state stays owned by the
    /// stage that runs the layer; `cache.pos` is snapshotted once and advanced once.
    /// MEMRA_PP_STREAMS=0 = the increment-1 same-stream seam.
    /// Gate: `ppn-gate` (bit-identical logits vs unsplit at every N/knob combination).
    fn decode_step_h_ppn(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
        fence: &[usize],
    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        if crate::pp::pp2_streams_off() {
            return self.decode_step_h_ppn_samestream(e, token, cache, fence);
        }
        let rt = crate::pp::PpNRt::get(e)?;
        let n_st = fence.len() - 1;
        assert_eq!(
            rt.n_stages(), n_st,
            "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
        );
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let pos = cache.pos;

        // PER-STAGE pos_d (M2 pipelining law): every stage uploads its OWN copy of the
        // step's pos scalar on ITS stream, so the buffer is allocated, consumed, and
        // freed on one stream (a shared stage-0 pos_d freed at fn return breaks under
        // deferred readback: the free enqueues on stream 0 while stages 1..N-1 still
        // dereference it — the 2026-08-02 pipelined-gate all-logits divergence).

        // ---- STAGE 0 (its own stream): embed + layers [0, fence[1]) + boundary-0 TX ----
        let mut slot = {
            let _st0 = rt.enter(0);
            let e0 = rt.engine(0, e);
            let pos_d = e0.htod_i32(&[pos as i32])?;
            let x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
            let x = self.decode_layers_eager(e0, x, fence[0], fence[1], &pos_d, pos, cache)?;
            rt.tx(0, &x, n_embd)?
            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
        };

        // ---- MIDDLE STAGES s in [1, n_st-1): RX boundary s-1 -> range -> TX boundary s ----
        for s in 1..n_st - 1 {
            let _st = rt.enter(s);
            let es = rt.engine(s, e);
            let pos_d = es.htod_i32(&[pos as i32])?;
            let x = rt.rx(s - 1, slot, n_embd)?;
            let x = self.decode_layers_eager(es, x, fence[s], fence[s + 1], &pos_d, pos, cache)?;
            slot = rt.tx(s, &x, n_embd)?;
        }

        // ---- LAST STAGE: RX + layers [fence[n_st-1], n) + output_norm + lm head ----
        let _stl = rt.enter(n_st - 1);
        let el = rt.engine(n_st - 1, e);
        let pos_d = el.htod_i32(&[pos as i32])?;
        let x = rt.rx(n_st - 2, slot, n_embd)?;
        let x =
            self.decode_layers_eager(el, x, fence[n_st - 1], fence[n_st], &pos_d, pos, cache)?;
        let e = el; // head runs through the last stage's engine on its stream

        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let h_seed = if crate::spec::spec_hpost() {
            e.clone_dtod(&hn)?
        } else {
            e.clone_dtod(&x)?
        };
        // same diagnostics door as decode_step_h (MEMRA_DUMP_HN) so the arms stay observably
        // interchangeable.
        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
            let hh = e.dtoh(&hn)?;
            use std::io::Write;
            let mut fo = std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(path)?;
            for v in &hh {
                fo.write_all(&v.to_le_bytes())?;
            }
        }
        let logits = e.matmul(&self.output, &hn, 1)?;
        let host = e.dtoh(&logits)?;
        cache.pos += 1;
        Ok((host, h_seed))
    }

    /// MEMRA_PP_STREAMS=0 rollback seam: the increment-1 body generalized to N — every
    /// stage subgraph on the ambient compute stream, each boundary = two plain dtod copies.
    fn decode_step_h_ppn_samestream(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
        fence: &[usize],
    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let pos = cache.pos;
        let pos_d = e.htod_i32(&[pos as i32])?;

        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) ----
        let x = e.htod(&self.embd.gather(n_embd, &[token]))?;
        let mut x = self.decode_layers_eager(e, x, fence[0], fence[1], &pos_d, pos, cache)?;

        // ---- each later stage: explicit [n_embd] handoff (TX copy, RX copy) + range ----
        for s in 1..fence.len() - 1 {
            let boundary_tx = e.clone_dtod(&x)?;
            let boundary_rx = e.clone_dtod(&boundary_tx)?;
            x = self.decode_layers_eager(e, boundary_rx, fence[s], fence[s + 1], &pos_d, pos, cache)?;
        }

        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let h_seed = if crate::spec::spec_hpost() {
            e.clone_dtod(&hn)?
        } else {
            e.clone_dtod(&x)?
        };
        if let Ok(path) = std::env::var("MEMRA_DUMP_HN") {
            let hh = e.dtoh(&hn)?;
            use std::io::Write;
            let mut fo = std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(path)?;
            for v in &hh {
                fo.write_all(&v.to_le_bytes())?;
            }
        }
        let logits = e.matmul(&self.output, &hn, 1)?;
        let host = e.dtoh(&logits)?;
        cache.pos += 1;
        Ok((host, h_seed))
    }

    /// M2 increment 3 (DEFERRED READBACK — the pipelining seed): the ppN step WITHOUT the
    /// terminal logits D2H. Returns `PendingLogits` (device logits + completion event +
    /// the runtime's dedicated readback stream); the caller keeps 2+ tokens in flight by
    /// enqueueing step t+1 BEFORE waiting step t (with MEMRA_PP_OVERLAP=1 the
    /// double-buffered boundary slots actually alternate, so stage 0 of t+1 runs under
    /// stage 1..N-1 of t; the slot ev_tx/ev_rx chain keeps each token's math fully
    /// event-ordered either way — enqueueing deeper than 2 is CORRECT, the slots simply
    /// serialize device-side).
    ///
    /// EXACTNESS CONTRACT: per-token logits are BIT-IDENTICAL to the serial arm — same
    /// kernels, same per-token event order; only the host-side wait moves (scheduling
    /// change, never math). The pipelined replay arm of `ppn-gate` proves it per step.
    ///
    /// NOT produced here (both are trunk COPIES — no math feeding the logits changes):
    /// h_seed and the MEMRA_DUMP_HN diagnostic tap. The serving loop decides their
    /// deferred form when it adopts this API.
    ///
    /// The caller advances the token stream, so `cache.pos` advances at ENQUEUE (host
    /// state; device work is event-ordered regardless).
    pub fn decode_step_h_ppn_deferred(
        &self,
        e: &Engine,
        token: u32,
        cache: &mut Cache,
    ) -> Result<crate::pp::PendingLogits, Box<dyn std::error::Error>> {
        let fence = crate::pp::pp_cuts(self.layers.len())
            .ok_or("ppn deferred: pp door closed (MEMRA_PP_STAGES unset)")?;
        if crate::pp::pp2_streams_off() {
            return Err("ppn deferred needs per-stage streams (MEMRA_PP_STREAMS=0 set)".into());
        }
        if self.cfg.gemma4.is_some() {
            return Err("ppn deferred: generic eager arm only (gemma4 is 2-stage serial)".into());
        }
        if crate::pp::pp_multi_stream_same_device()
            && std::env::var("MEMRA_PP_FORCE_SAME_DEV_PIPELINED").as_deref() != Ok("1")
        {
            return Err(
                "ppn deferred: refused with 2+ stage streams on one device — repro'd \
                 nondeterministic logits (35% flake, 2026-08-02 x20 soak, root cause open: \
                 shared-Engine kernels concurrent on co-located streams). Use one device \
                 per stage (MEMRA_PP_DEVICES) or the serial arm. \
                 MEMRA_PP_FORCE_SAME_DEV_PIPELINED=1 overrides for soak/bisect measurement."
                    .into(),
            );
        }
        let rt = crate::pp::PpNRt::get(e)?;
        let n_st = fence.len() - 1;
        assert_eq!(
            rt.n_stages(), n_st,
            "PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
        );
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let pos = cache.pos;

        // Per-stage pos_d — see decode_step_h_ppn: under deferred readback a shared
        // pos_d's fn-end free races stages 1..N-1 (the free enqueues on stream 0 at
        // ENQUEUE time here, no terminal D2H to drain first). Each stage owns its copy.
        let mut slot = {
            let _st0 = rt.enter(0);
            let e0 = rt.engine(0, e);
            let pos_d = e0.htod_i32(&[pos as i32])?;
            let x = e0.htod(&self.embd.gather(n_embd, &[token]))?;
            let x = self.decode_layers_eager(e0, x, fence[0], fence[1], &pos_d, pos, cache)?;
            rt.tx(0, &x, n_embd)?
        };
        for s in 1..n_st - 1 {
            let _st = rt.enter(s);
            let es = rt.engine(s, e);
            let pos_d = es.htod_i32(&[pos as i32])?;
            let x = rt.rx(s - 1, slot, n_embd)?;
            let x = self.decode_layers_eager(es, x, fence[s], fence[s + 1], &pos_d, pos, cache)?;
            slot = rt.tx(s, &x, n_embd)?;
        }
        let _stl = rt.enter(n_st - 1);
        let el = rt.engine(n_st - 1, e);
        let pos_d = el.htod_i32(&[pos as i32])?;
        let x = rt.rx(n_st - 2, slot, n_embd)?;
        let x =
            self.decode_layers_eager(el, x, fence[n_st - 1], fence[n_st], &pos_d, pos, cache)?;

        let mut hn = el.uninit(n_embd)?;
        el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let logits = el.matmul(&self.output, &hn, 1)?;
        let ev = rt.record_done()?;
        cache.pos += 1;
        Ok(crate::pp::PendingLogits::new(logits, ev, rt.readback_stream().clone()))
    }

    /// LOCKSTEP MULTI-STREAM decode (lane-3 M1): m independent streams advance one token each
    /// through a single per-layer walk. Per-stream math is identical to `decode_step_h` (same
    /// fusion chain, same mixer and FFN calls against that stream's own `Cache`), so each
    /// stream's token sequence is bit-identical to its single-stream run. The lockstep order
    /// puts the m streams' layer-il MoE calls adjacent in time, so one stream's expert-cache
    /// fill serves its siblings within the step — the measured cross-stream io amortization
    /// (1.12x/1.32x/1.66x at m=2/4/8) lands without batching attention or the CPU ABI.
    pub fn decode_step_lockstep(
        &self,
        e: &Engine,
        tokens: &[u32],
        caches: &mut [Cache],
    ) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
        if tokens.len() != caches.len() || tokens.is_empty() {
            return Err("lockstep needs one token per stream cache".into());
        }
        if self.cfg.gemma4.is_some() {
            return Err("lockstep decode does not support the gemma4 paths".into());
        }
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let m = tokens.len();

        let mut pos_d = Vec::with_capacity(m);
        let mut x: Vec<CudaSlice<f32>> = Vec::with_capacity(m);
        for (s, &token) in tokens.iter().enumerate() {
            pos_d.push(e.htod_i32(&[caches[s].pos as i32])?);
            x.push(e.htod(&self.embd.gather(n_embd, &[token]))?);
        }
        let mut pending: Vec<Option<(CudaSlice<f32>, CudaSlice<f32>)>> =
            (0..m).map(|_| None).collect();

        // M2 (MEMRA_LOCKSTEP_GROUPED=1): MoE layers batch all m rows through
        // moe_ffn_lockstep — resident experts amortize weight reads across streams via the
        // grouped GEMM machinery; CPU-assigned experts keep per-row companion calls.
        let grouped = match std::env::var("MEMRA_LOCKSTEP_GROUPED").as_deref() {
            Ok("1") => true,
            Ok("0") => false,
            // Auto: grouped wins from m>=3 under the default q8 lanes (M2 gate 2026-07-23:
            // m=2 6.17 base vs 5.85 grouped; m=3 6.31 grouped; m=4 5.66 vs 5.34).
            _ => m >= 3,
        };
        // M4a (MEMRA_LOCKSTEP_BATCH_ATTN=1): EXPERIMENTAL DOOR, measured flat — default off.
        // Full-attention layers run their WEIGHT-BOUND work (q/k/v and output projections) once
        // at m instead of m times, KV-bound work stays per stream. Bit-identity PASS, but e2e
        // flat at m=2 (4.72/4.72) and -2% at m=3 (5.24 vs 5.35), 2026-07-25: full-attn is the
        // minority layer type here (GDN dominates), so the m-band weight-read saving covers few
        // layers and is cancelled by the norm->q8_1 fusion this path gives up on exactly those
        // layers, plus its gather/scatter copies. The primitive itself
        // (`full_attn_decode_batched`) stays as the m-band building block for a serve loop,
        // where batching happens across requests at higher m and no fused alternative exists.
        let batch_attn = matches!(
            std::env::var("MEMRA_LOCKSTEP_BATCH_ATTN").as_deref(), Ok("1")
        ) && m >= 2;
        let pos_cat = e.htod_i32(
            &caches.iter().take(m).map(|c| c.pos as i32).collect::<Vec<_>>(),
        )?;
        let n_embd_total = n_embd * m;
        let mut xcat = e.uninit(n_embd_total)?;
        for (il, layer) in self.layers.iter().enumerate() {
            let anorm = layer.attn_norm.float_data();
            let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
                && self.mixer_in_q8_1_fast(e, &layer.mixer);
            let mut mixed_rows: Vec<Option<CudaSlice<f32>>> = (0..m).map(|_| None).collect();
            if batch_attn && matches!(layer.mixer, Mixer::Full(_)) {
                // Unfused residual+norm into the contiguous m-band buffer. Bit-identical to the
                // fused arm by construction (add_rms_norm_q8_1 == add, rms_norm, quantize_q8_1);
                // the batched mixer quantizes all m rows in one call.
                for s in 0..m {
                    if let Some((x1, f1)) = pending[s].take() {
                        let mut x2 = e.uninit(n_embd)?;
                        e.add(&x1, &f1, &mut x2, n_embd)?;
                        x[s] = x2;
                    }
                    let mut hn = e.uninit(n_embd)?;
                    e.rms_norm(&x[s], anorm, &mut hn, n_embd, 1, eps)?;
                    e.copy_into(&mut xcat, s * n_embd, &hn, n_embd)?;
                }
                let Mixer::Full(fa) = &layer.mixer else { unreachable!() };
                let out_cat =
                    self.full_attn_decode_batched(e, fa, &xcat, m, &pos_cat, caches, il)?;
                for s in 0..m {
                    let mut mixed = e.uninit(n_embd)?;
                    e.copy_view_into(
                        &mut mixed, 0,
                        &out_cat.slice(s * n_embd..(s + 1) * n_embd), n_embd)?;
                    if grouped && matches!(&layer.ffn, crate::hybrid::Ffn::Moe(_)) {
                        mixed_rows[s] = Some(mixed);
                    } else {
                        let (x1, ffn_out) =
                            self.residual_norm_ffn(e, layer, &x[s], &mixed, n_embd, il, eps)?;
                        pending[s] = Some((x1, ffn_out));
                    }
                }
            } else {
            for s in 0..m {
                let pos = caches[s].pos;
                let taken = pending[s].take();
                let mixed = match (taken, fuse) {
                    (Some((x1, f1)), true) => {
                        let mut x2 = e.uninit(n_embd)?;
                        let (hq, hd) =
                            e.add_rms_norm_q8_1(&x1, &f1, anorm, &mut x2, n_embd, 1, eps)?;
                        x[s] = x2;
                        let h0 = e.zeros(0)?;
                        match &layer.mixer {
                            Mixer::Full(fa) => self.full_attn_decode_pre(
                                e,
                                fa,
                                &h0,
                                Some((&hq, &hd)),
                                &pos_d[s],
                                pos,
                                &mut caches[s],
                                il,
                            )?,
                            Mixer::Linear(la) => self.linear_attn_decode_pre(
                                e,
                                la,
                                &h0,
                                &hq,
                                &hd,
                                &mut caches[s],
                                il,
                                false,
                            )?,
                            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
                        }
                    }
                    (taken, _) => {
                        if let Some((x1, f1)) = taken {
                            let mut x2 = e.uninit(n_embd)?;
                            e.add(&x1, &f1, &mut x2, n_embd)?;
                            x[s] = x2;
                        }
                        self.attn_in_norm_mixer(
                            e,
                            layer,
                            &x[s],
                            &pos_d[s],
                            pos,
                            &mut caches[s],
                            il,
                            n_embd,
                            eps,
                        )?
                    }
                };
                if grouped && matches!(&layer.ffn, crate::hybrid::Ffn::Moe(_)) {
                    mixed_rows[s] = Some(mixed);
                } else {
                    let (x1, ffn_out) =
                        self.residual_norm_ffn(e, layer, &x[s], &mixed, n_embd, il, eps)?;
                    pending[s] = Some((x1, ffn_out));
                }
            }
            }
            if grouped {
                if let crate::hybrid::Ffn::Moe(moe_weights) = &layer.ffn {
                    // Per-stream add+norm (identical math to residual_norm_ffn's MoE arm),
                    // rows batched for the cross-stream MoE stage, outputs split back.
                    let pnorm = layer.post_attn_norm.float_data();
                    let mut zbatch = e.uninit(n_embd_total)?;
                    let mut x1s: Vec<CudaSlice<f32>> = Vec::with_capacity(m);
                    for s in 0..m {
                        let mixed = mixed_rows[s].take().expect("grouped MoE row missing");
                        let mut x1 = e.uninit(n_embd)?;
                        let mut z = e.uninit(n_embd)?;
                        e.add_rms_norm(&x[s], &mixed, pnorm, &mut x1, &mut z, n_embd, 1, eps)?;
                        e.copy_view_into(&mut zbatch, s * n_embd, &z.slice(0..n_embd), n_embd)?;
                        x1s.push(x1);
                    }
                    let max_block = self.max_moe_block();
                    let ffn_all =
                        self.moe_ffn_lockstep(e, moe_weights, &zbatch, m, il as u16, max_block)?;
                    for (s, x1) in x1s.into_iter().enumerate() {
                        let mut out = e.uninit(n_embd)?;
                        e.copy_view_into(
                            &mut out, 0,
                            &ffn_all.slice(s * n_embd..(s + 1) * n_embd), n_embd)?;
                        pending[s] = Some((x1, out));
                    }
                }
            }
        }

        let mut logits_host = Vec::with_capacity(m);
        for s in 0..m {
            if let Some((x1, f1)) = pending[s].take() {
                let mut x2 = e.uninit(n_embd)?;
                e.add(&x1, &f1, &mut x2, n_embd)?;
                x[s] = x2;
            }
            let mut hn = e.uninit(n_embd)?;
            e.rms_norm(&x[s], self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
            let logits = e.matmul(&self.output, &hn, 1)?;
            logits_host.push(e.dtoh(&logits)?);
            caches[s].pos += 1;
        }
        Ok(logits_host)
    }

    /// DEVICE-COUNTER decode step (CUDA-GRAPH-PLAN Phase 2). A clone of `decode_step_h` that removes
    /// the two per-step VARYING host kernel-args by reading them from device counters:
    ///   1. the KV-append write slot  -> per-layer `kvl.len_d` (device i32[1])
    ///   2. the fa_decode t_kv bound   -> the same `kvl.len_d` after `inc_seqlen`
    /// plus it keeps the token id + rope pos DEVICE-RESIDENT (embed_gather_device, device rope pos,
    /// argmax_token_device). NO graph capture yet — runs the kernels eagerly through the counter
    /// path. Must be BIT-IDENTICAL to `decode_step_h`'s token stream (the gate).
    ///
    /// Args: `token_d` = resident device token id [1] (this step's input token); `pos_d` = resident
    /// device rope pos i32[1] (== cache.pos at entry; INCREMENTED in-path); `embd_gpu` = resident embed
    /// table; (qt,row_bytes) from EmbedHost::qt_and_row_bytes. Returns the NEXT token id device buffer.
    /// `cache.pos` and each `kvl.len`/`kvl.len_d` are advanced to match `decode_step_h`.
    pub fn decode_step_dc(
        &self,
        e: &Engine,
        token_d: &CudaSlice<u32>,
        pos_d: &mut CudaSlice<i32>,
        embd_gpu: &CudaSlice<u8>,
        embd_qt: i32,
        embd_row_bytes: usize,
        cache: &mut Cache,
        n_vocab: usize,
    ) -> Result<CudaSlice<u32>, Box<dyn std::error::Error>> {
        // Route gemma4 to ITS dc twin (mirrors decode_step_h): the generic walk below is the
        // qwen-class layer stack — running gemma weights through it produced the argmax-INIT
        // passthrough the round-45 g12 gate caught (first Hopper gating of this lane).
        if self.is_gemma4_e4b() {
            return Err("e4b has no device-counter decode step (dc/graph unwired)".into());
        }
        if self.cfg.gemma4.is_some() {
            return self.gemma4_decode_step_dc(e, token_d, pos_d, embd_gpu, embd_qt,
                                              embd_row_bytes, cache, n_vocab, None);
        }
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;

        // embed the single (DEVICE-resident) token -> [1, n_embd], no host round-trip of the id.
        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_row_bytes)?;

        for (il, layer) in self.layers.iter().enumerate() {
            // attn-input NORM-FUSION (dc path); bit-identical to decode_step_h (Phase-2 gate).
            let mixed = self.attn_in_norm_mixer_dc(e, layer, &x, pos_d, cache, il, n_embd, eps)?;

            // DECODE NORM-FUSION LEVER (residual_norm_ffn): see decode_step_h. Shared helper -> dc
            // path stays bit-identical to decode_step_h's token stream (the Phase-2 gate).
            let (x1, ffn_out) = self.residual_norm_ffn(e, layer, &x, &mixed, n_embd, il, eps)?;
            let mut x2 = e.uninit(n_embd)?;
            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
            x = x2;
        }

        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let logits = e.matmul(&self.output, &hn, 1)?;
        // device argmax -> next token id stays resident (no logits dtoh).
        let next_tok = e.argmax_token_device(&logits, n_vocab)?;
        // advance rope pos counter on-device (replaces the per-step htod_i32(&[pos])).
        e.inc_seqlen(pos_d)?;
        cache.pos += 1;
        Ok(next_tok)
    }

    /// CAPTURE body for CUDA-graph replay (CUDA-GRAPH-PLAN Phase 3). One full decode step enqueued
    /// entirely on `e.stream()` with ZERO host sync and ZERO per-step varying host kernel-args:
    ///   - embed reads the PERSISTENT device `token_d` (last step's argmax), writes scratch `x`.
    ///   - full-attn layers size n_splits from `bucket_max` (fixed for this capture); the kernel reads
    ///     the ACTUAL t_kv from the device counter `kvl.len_d`. KV append + device-counter inc happen
    ///     in-graph. The host `kvl.len`/`cache.pos` are NOT advanced here (the driver advances the host
    ///     mirrors once per replay; only the DEVICE counters advance inside the graph).
    ///   - linear-attn layers use the persistent-state variant (copy-back, stable pointers).
    ///   - lm_head -> parallel 2-pass argmax (`argmax_partial_f32`+`argmax_final_f32`) writes the
    ///     next id into the PERSISTENT `token_d`.
    ///   - `inc_seqlen(pos_d)` advances the rope-pos device counter in-graph.
    /// Captured ONCE per `bucket_max`; replayed for every t_kv in that bucket. Bit-identical to eager
    /// when `bucket_max` reproduces eager's n_splits for the replayed t_kv (the bucket-key contract).
    pub fn decode_step_dc_cap(
        &self,
        e: &Engine,
        token_d: &mut CudaSlice<u32>,
        pos_d: &mut CudaSlice<i32>,
        embd_gpu: &CudaSlice<u8>,
        embd_qt: i32,
        embd_row_bytes: usize,
        cache: &mut Cache,
        n_vocab: usize,
        bucket_max: usize,
    ) -> Result<(), Box<dyn std::error::Error>> {
        self.decode_step_dc_cap_masked(e, token_d, pos_d, embd_gpu, embd_qt, embd_row_bytes,
                                       cache, n_vocab, bucket_max, None)
    }

    /// `decode_step_dc_cap` + GRAMMAR MASK (constrained decoding): with `mask =
    /// Some((buf, words))`, mask_logits_f32 bans the packed bitset's unset ids IN the
    /// captured graph — a stable-pointer read between lm_head and the in-graph argmax
    /// (the KV-pointer pattern: contents change per step, address is baked). `None` is
    /// bit-for-bit the unmasked capture.
    #[allow(clippy::too_many_arguments)]
    pub fn decode_step_dc_cap_masked(
        &self,
        e: &Engine,
        token_d: &mut CudaSlice<u32>,
        pos_d: &mut CudaSlice<i32>,
        embd_gpu: &CudaSlice<u8>,
        embd_qt: i32,
        embd_row_bytes: usize,
        cache: &mut Cache,
        n_vocab: usize,
        bucket_max: usize,
        mask: Option<(&CudaSlice<u32>, usize)>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;

        let mut x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_row_bytes)?;

        for (il, layer) in self.layers.iter().enumerate() {
            // attn-input NORM-FUSION (capture path); capture-safe + bit-identical to eager.
            let mixed = self.attn_in_norm_mixer_dc_cap(
                e, layer, &x, pos_d, cache, il, bucket_max, n_embd, eps,
            )?;
            // DECODE NORM-FUSION LEVER (residual_norm_ffn): see decode_step_aux. Shared helper keeps
            // the capture path bit-identical to eager by construction.
            let (x1, ffn_out) = self.residual_norm_ffn(e, layer, &x, &mixed, n_embd, il, eps)?;
            let mut x2 = e.uninit(n_embd)?;
            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
            x = x2;
        }

        let mut hn = e.uninit(n_embd)?;
        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
        let mut logits = e.matmul(&self.output, &hn, 1)?;
        // GRAMMAR MASK: ban before the argmax reads the row (masked argmax == host
        // masked-argmax — -FLT_MAX is the argmax kernels' init sentinel).
        if let Some((m, words)) = mask {
            e.mask_logits_col(&mut logits, m, 0, n_vocab, words)?;
        }
        // argmax into the PERSISTENT token_d (next step's embed reads it) — same buffer pointer baked
        // at capture, written each replay, so the token id never round-trips to host in steady state.
        e.argmax_token_device_into(&logits, token_d, n_vocab)?;
        e.inc_seqlen(pos_d)?;
        Ok(())
    }

    /// CUDA-GRAPH decode driver (CUDA-GRAPH-PLAN Phase 3). Primes the prompt EAGERLY (device-counter
    /// `decode_step_dc`, advancing host + device counters together), then generates `max_new` tokens by
    /// CUDA-graph REPLAY: per step it picks the t_kv bucket key, captures a graph on first sight of that
    /// key (re-using the SAME persistent counters/cache so replays continue the sequence), and replays.
    /// The argmax-written next token stays device-resident in `gs.token_d`; we read back only the [1]
    /// u32 after each launch (the gate compares it; a real server can defer this). Returns the generated
    /// token ids. Greedy. Bit-identical to eager `decode_step` (the gate).
    ///
    /// CAPTURE STATE HYGIENE: `capture_graph` runs the step body 3x (2 warmup + 1 capture), each of
    /// which mutates the device KV/conv/ssm/counter state. We SNAPSHOT the cache + device counters +
    /// token id before capturing and RESTORE them after, so the 3 throwaway runs leave zero residue and
    /// replay resumes from the true pre-capture state.
    pub fn generate_graph(
        &self,
        e: &Engine,
        gs: &mut GraphDecodeState,
        prompt: &[u32],
        max_new: usize,
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        let n_embd = self.cfg.n_embd as usize;
        let head_dim = self.cfg.head_dim_k as usize;
        let (qt, row_bytes) = self.embd.qt_and_row_bytes(n_embd);

        // EVENT TRACKING OFF for the WHOLE graph-decode session. cudarc records a per-CudaSlice event
        // (the Engine is in multi-stream mode via copy_stream) and inserts `stream.wait(event)` on every
        // kernel arg whose buffer was touched — those waits are illegal inside a capture region. The
        // captured decode step is strictly single-stream, so this tracking is unnecessary. Disable it
        // BEFORE allocating ANY buffer the captured graph will reference (cache, embd, counters,
        // scratch) so none of them carry events. SAFETY: decode-dc touches only gpu.stream.
        let was_tracking = e.ctx().is_event_tracking();
        if was_tracking {
            unsafe {
                e.ctx().disable_event_tracking();
            }
        }
        let r = self.generate_graph_inner(e, gs, prompt, max_new, n_embd, head_dim, qt, row_bytes);
        if was_tracking {
            unsafe {
                e.ctx().enable_event_tracking();
            }
        }
        r
    }

    fn generate_graph_inner(
        &self,
        e: &Engine,
        gs: &mut GraphDecodeState,
        prompt: &[u32],
        max_new: usize,
        n_embd: usize,
        head_dim: usize,
        qt: i32,
        row_bytes: usize,
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        let _ = n_embd;
        let embd_gpu = e.upload_u8(&self.embd.raw)?;
        let max_ctx = prompt.len() + max_new + 8;
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;

        // (Re)create the persistent counters tracking-OFF so they carry no events (the caller's
        // GraphDecodeState::new may have allocated them with tracking on).
        gs.pos_d = e.htod_i32(&[0])?;
        gs.token_d = e.stream().clone_htod(&[0u32])?;
        // PRIME eagerly: feed each prompt token; advance host + device counters together.
        let mut next_in = 0u32;
        for &tok in prompt {
            e.set_u32_one(&mut gs.token_d, tok)?;
            let nt = self.decode_step_dc(
                e,
                &gs.token_d,
                &mut gs.pos_d,
                &embd_gpu,
                qt,
                row_bytes,
                &mut cache,
                /*n_vocab*/ self.output.out_features(),
            )?;
            next_in = e.dtoh_u32_one(&nt)?;
        }
        // gs.token_d now must hold the first generated INPUT token (= argmax of the last prime step).
        e.set_u32_one(&mut gs.token_d, next_in)?;

        // gemma4 rides ITS graph machinery (per-bucket captures + alloc-free slots; same token
        // stream convention: first generated token is out[0]) — graph_decode_loop below captures
        // the qwen-class dc step (the round-45 g12 illegal-address find).
        if self.cfg.gemma4.is_some() {
            let (toks, _reason) = self.gemma4_generate_graph(
                e, cache.pos, next_in, &mut cache, max_new, &[], |_| true)?;
            gs.captures += 1;
            return Ok(toks);
        }

        let mut out = Vec::with_capacity(max_new);
        self.graph_decode_loop(e, gs, &mut cache, &embd_gpu, qt, row_bytes, head_dim, max_new,
                               |tok| { out.push(tok); None })?;
        Ok(out)
    }

    /// The CUDA-graph EXEC-UPDATE replay loop over an already-primed cache (2026-07-15,
    /// the E4B graph-exec pattern generalized): capture the dc step per KERNEL-CLASS
    /// SEGMENT, classify its fa nodes (`graph_update::fa_plan` — symbol list is
    /// model-generic), then per token retune the fa split geometry to the LIVE eager
    /// ladder (`fa_apply` keeps graph and eager in FP lockstep — bit-exact) and replay.
    /// The previous per-bucket-key capture map recaptured on every ladder rung
    /// (32 recaptures/256 tokens = 97 vs 128 tok/s eager; decode-bench 2026-07-15).
    ///
    /// SEGMENTS (round 45, the q35 graph-gate dig): exec-update can retune split counts
    /// but can NOT swap kernels — a session spanning an eager KERNEL-CLASS boundary
    /// (fa_vec floor, the v4 max, the fa512 floor) replayed the capture-time kernel
    /// against a different eager kernel below the boundary: valid softmax, different
    /// fold order, and the first near-tie flips the stream (q35: deterministic 144/256
    /// from step 110, exactly the scalar->vec crossing; regime pinned either way =
    /// BIT-IDENTICAL 256/256). One capture per crossed class boundary (2-3/session,
    /// not per rung) keeps graph and eager on the SAME kernel at every t_kv.
    ///
    /// Callers must have synced gs.token_d (= the FIRST generated token), gs.pos_d
    /// (= cache.pos) and every kvl.len_d (= kvl.len). Event tracking must be OFF.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn graph_decode_loop(&self, e: &Engine, gs: &mut GraphDecodeState,
                                    cache: &mut Cache, embd_gpu: &CudaSlice<u8>,
                                    qt: i32, row_bytes: usize, head_dim: usize, max_new: usize,
                                    mut emit: impl FnMut(u32) -> Option<StopReason>)
                                    -> Result<StopReason, Box<dyn std::error::Error>> {
        let _ = head_dim;
        let n_vocab = self.output.out_features();
        let final_max = cache.pos + max_new + 1;

        // first generated token = argmax of the last prime step (emit before replay 1).
        let first = e.dtoh_u32_one(&gs.token_d)?;
        if let Some(r) = emit(first) { return Ok(r); }
        let mut done = 1usize;
        while done < max_new {
            let (graph, mut plan, seg_end) = self.graph_capture_segment(
                e, cache, gs, embd_gpu, qt, row_bytes, n_vocab, final_max)?;

            while done < max_new && cache.pos + 1 <= seg_end {
                // retune fa geometry to the live t_kv AFTER this replay's in-graph append.
                crate::graph_update::fa_apply(&graph, &mut plan, cache.pos + 1,
                                              crate::fa_split_keys)?;
                graph.launch()?;
                cache.pos += 1;
                for kvl in cache.kv.iter_mut().filter_map(|k| k.as_mut()) {
                    kvl.len += 1;
                }
                // read back the [1] u32 next token (the only D2H in steady state).
                let tok = e.dtoh_u32_one(&gs.token_d)?;
                done += 1;
                if let Some(r) = emit(tok) { return Ok(r); }
            }
        }
        Ok(StopReason::MaxNew)
    }

    /// Step-wise CUDA-graph decode session (ARCHITECTURE-H100.md graph-serving lane,
    /// 2026-07-26): generate_graph's prime+capture lifted into a long-lived session so a
    /// SERVING scheduler can replay ONE step per tick instead of blocking a whole
    /// generation. Serving policy (measured): graphs win only at B=1 (214 solo vs 425
    /// aggregate batched-eager at B=4) — this is the single-interactive-session path.
    /// Capture discipline is generate_graph's verbatim: event tracking must be OFF for
    /// every buffer the graph references (new() toggles it), capture at bucket_max =
    /// pos + max_new + 1, fa geometry retuned per step (fa_apply, FP lockstep with eager).
    pub fn graph_session_new(
        &self,
        e: &Engine,
        prompt: &[u32],
        max_new: usize,
    ) -> Result<(GraphSession, u32), Box<dyn std::error::Error>> {
        let n_embd = self.cfg.n_embd as usize;
        let (qt, row_bytes) = self.embd.qt_and_row_bytes(n_embd);
        let was_tracking = e.ctx().is_event_tracking();
        if was_tracking {
            unsafe { e.ctx().disable_event_tracking(); }
        }
        let r = self.graph_session_new_inner(e, prompt, max_new, qt, row_bytes);
        if was_tracking {
            unsafe { e.ctx().enable_event_tracking(); }
        }
        r
    }

    fn graph_session_new_inner(
        &self,
        e: &Engine,
        prompt: &[u32],
        max_new: usize,
        qt: i32,
        row_bytes: usize,
    ) -> Result<(GraphSession, u32), Box<dyn std::error::Error>> {
        let n_vocab = self.output.out_features();
        let embd_gpu = e.upload_u8(&self.embd.raw)?;
        let max_ctx = prompt.len() + max_new + 8;
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;
        let mut gs = GraphDecodeState::new(e)?;
        gs.pos_d = e.htod_i32(&[0])?;
        gs.token_d = e.stream().clone_htod(&[0u32])?;
        // prime (dc path — device counters advance with the host)
        let mut next_in = 0u32;
        for &tok in prompt {
            e.set_u32_one(&mut gs.token_d, tok)?;
            let nt = self.decode_step_dc(e, &gs.token_d, &mut gs.pos_d, &embd_gpu,
                                         qt, row_bytes, &mut cache, n_vocab)?;
            next_in = e.dtoh_u32_one(&nt)?;
        }
        e.set_u32_one(&mut gs.token_d, next_in)?;
        self.graph_session_capture(e, cache, gs, embd_gpu, max_new, qt, row_bytes, n_vocab,
                                   None, 0)
    }

    /// GraphSession over an ALREADY-PRIMED cache (round 35): keeps the chunked-prefill
    /// TTFT. graph_session_new's token-wise re-prime made solo long-prompt promotion a
    /// net ~3x END-TO-END LOSS (measured live: 871-tok prompt + 400 gen = 6.4s vs ~2.2s
    /// eager). Device counters sync from host state; capture recipe unchanged.
    /// Requires event tracking OFF (engine default; MEMRA_EVT=1 callers must not use this
    /// — the primed cache's buffers would carry events, illegal inside capture).
    pub fn graph_session_from_cache(
        &self,
        e: &Engine,
        cache: Cache,
        first_token: u32,
        max_new: usize,
    ) -> Result<(GraphSession, u32), Box<dyn std::error::Error>> {
        self.graph_session_from_cache_masked(e, cache, first_token, max_new, None)
    }

    /// `graph_session_from_cache` + GRAMMAR MASK (constrained decoding, 2026-08-03):
    /// `mask_init = Some(packed bitset)` allocates the session's stable mask buffer
    /// (tracking is OFF here — capture-legal), seeds it with the FIRST step's mask, and
    /// captures mask_logits_f32 into the graphed step. The caller re-uploads contents
    /// per step via `GraphSession::upload_mask` — same stable-pointer discipline as the
    /// KV len_d counters. `None` = the unmasked session, byte-identical.
    pub fn graph_session_from_cache_masked(
        &self,
        e: &Engine,
        mut cache: Cache,
        first_token: u32,
        max_new: usize,
        mask_init: Option<&[u32]>,
    ) -> Result<(GraphSession, u32), Box<dyn std::error::Error>> {
        if e.ctx().is_event_tracking() {
            return Err("graph_session_from_cache requires event tracking OFF (MEMRA_EVT unset)".into());
        }
        let n_embd = self.cfg.n_embd as usize;
        let (qt, row_bytes) = self.embd.qt_and_row_bytes(n_embd);
        let n_vocab = self.output.out_features();
        let embd_gpu = e.upload_u8(&self.embd.raw)?;
        let mut gs = GraphDecodeState::new(e)?;
        gs.pos_d = e.htod_i32(&[cache.pos as i32])?;
        gs.token_d = e.stream().clone_htod(&[first_token])?;
        for kvl in cache.kv.iter_mut().flatten() {
            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
        }
        let mask_dev = match mask_init {
            Some(w) => Some(e.htod_u32_v(w)?),
            None => None,
        };
        let mask_words = mask_init.map(|w| w.len()).unwrap_or(0);
        self.graph_session_capture(e, cache, gs, embd_gpu, max_new, qt, row_bytes, n_vocab,
                                   mask_dev, mask_words)
    }

    /// Eager fa kernel-class fingerprint at a given t_kv: the fa_vec pick plus the
    /// intra-vec variant switches (v4 max, fa512 floor) plus the split-ladder rung.
    /// fa_apply handles split-count changes WITHIN a rung; anything that changes this
    /// tuple needs a fresh capture (bucket_max drives the capture-time kernel pick).
    /// Round 45; LADDER RUNG ADDED 2026-08-02 (lane/ladder-3072): the dc kernels derive
    /// their in-kernel partition from the CAPTURED split_keys arg (ns_eff =
    /// ceil(T_kv/split_keys) — the ONE-PARTITION law), and fa_apply retunes only
    /// n_splits/grid. A capture whose segment straddled a ladder rung therefore replayed
    /// the far side's partition against eager's near side — same math, different FP fold
    /// order, and the first near-tie flips the stream (latent at the old 3072 rung: kat
    /// P=3000 passed on logit margins; exposed by the 512 rung: kat P=400 flipped 97/160).
    /// With the rung in the fingerprint a capture never straddles it, so the captured
    /// split_keys equals the live ladder on every replay — bit-exact at every t_kv.
    pub(crate) fn fa_class_of(&self, e: &Engine, t_kv: usize) -> (bool, bool, bool, usize) {
        let head_dim = self.cfg.head_dim_k as usize;
        let nkv = self.cfg.n_head_kv as usize;
        let g_fp8 = Engine::kv_fp8_on();
        (e.fa_geom_eager(t_kv, head_dim, nkv, g_fp8).0,
         crate::fa_v4_at_pub(t_kv),
         head_dim == 512 && t_kv >= crate::fa512_min_tkv(),
         crate::fa_split_keys_pub(t_kv, nkv))
    }

    /// Last t_kv (clamped to `final_max`) sharing `start`'s eager kernel class.
    pub(crate) fn fa_segment_end(&self, e: &Engine, start: usize, final_max: usize) -> usize {
        let cls = self.fa_class_of(e, start);
        let mut end = start;
        while end < final_max && self.fa_class_of(e, end + 1) == cls { end += 1; }
        end
    }

    /// Capture one kernel-class segment: snapshot/rollback the warmup runs, capture the
    /// dc step at bucket_max = the segment's last t_kv, fa_plan. Shared by the session
    /// creation, the session's recapture-on-cross, and graph_decode_loop.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn graph_capture_segment(
        &self,
        e: &Engine,
        cache: &mut Cache,
        gs: &mut GraphDecodeState,
        embd_gpu: &CudaSlice<u8>,
        qt: i32,
        row_bytes: usize,
        n_vocab: usize,
        final_max: usize,
    ) -> Result<(cudarc::driver::CudaGraph, Vec<crate::graph_update::FaMain>, usize),
                Box<dyn std::error::Error>> {
        self.graph_capture_segment_masked(e, cache, gs, embd_gpu, qt, row_bytes, n_vocab,
                                          final_max, None)
    }

    /// `graph_capture_segment` + optional in-graph grammar mask (see decode_step_dc_cap_masked).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn graph_capture_segment_masked(
        &self,
        e: &Engine,
        cache: &mut Cache,
        gs: &mut GraphDecodeState,
        embd_gpu: &CudaSlice<u8>,
        qt: i32,
        row_bytes: usize,
        n_vocab: usize,
        final_max: usize,
        mask: Option<(&CudaSlice<u32>, usize)>,
    ) -> Result<(cudarc::driver::CudaGraph, Vec<crate::graph_update::FaMain>, usize),
                Box<dyn std::error::Error>> {
        let t0 = cache.pos + 1;
        let seg_end = self.fa_segment_end(e, t0, final_max);
        let bucket_max = seg_end;
        let snap = cache.snapshot(e)?;
        let pos_save = e.dtoh_i32_one(&gs.pos_d)?;
        let len_save: Vec<Option<i32>> = cache.kv.iter()
            .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
        let tok_save = e.dtoh_u32_one(&gs.token_d)?;
        let graph = {
            let GraphDecodeState { token_d, pos_d, .. } = gs;
            let token_d: &mut CudaSlice<u32> = token_d;
            let pos_d: &mut CudaSlice<i32> = pos_d;
            let cache_ref = &mut *cache;
            e.capture_graph(|e| {
                self.decode_step_dc_cap_masked(e, token_d, pos_d, embd_gpu, qt, row_bytes,
                                               cache_ref, n_vocab, bucket_max, mask)
            })?
        };
        gs.captures += 1;
        cache.rollback(e, &snap, 0)?;
        e.set_i32_one(&mut gs.pos_d, pos_save)?;
        for (il, ls) in len_save.iter().enumerate() {
            if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
                e.set_i32_one(&mut kvl.len_d, *v)?;
            }
        }
        e.set_u32_one(&mut gs.token_d, tok_save)?;
        let plan = crate::graph_update::fa_plan(&graph)?;
        if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
            eprintln!("[graph-census] segment t_kv {t0}..={seg_end} fa_plan mains: {}",
                      plan.len());
            if let Ok(c) = crate::graph_update::node_census(&graph) {
                eprintln!("[graph-census] {c:?}");
            }
        }
        Ok((graph, plan, seg_end))
    }

    /// Session recapture at a kernel-class boundary (called by GraphSession::step).
    /// The mask node (when present) re-bakes the SAME stable buffer — contents carry over.
    pub(crate) fn graph_session_recapture(&self, e: &Engine, sess: &mut GraphSession)
                                          -> Result<(), Box<dyn std::error::Error>> {
        let mask = sess.mask_dev.take();
        let (graph, plan, seg_end) = self.graph_capture_segment_masked(
            e, &mut sess.cache, &mut sess.gs, &sess.embd_gpu,
            sess.qt, sess.row_bytes, sess.n_vocab, sess.bucket_max,
            mask.as_ref().map(|d| (d, sess.mask_words)))?;
        sess.mask_dev = mask;
        sess.graph = graph;
        sess.plan = plan;
        sess.seg_end = seg_end;
        Ok(())
    }

    /// Shared capture tail: capture the FIRST kernel-class segment, build the session.
    #[allow(clippy::too_many_arguments)]
    fn graph_session_capture(
        &self,
        e: &Engine,
        mut cache: Cache,
        mut gs: GraphDecodeState,
        embd_gpu_owned: CudaSlice<u8>,
        max_new: usize,
        qt: i32,
        row_bytes: usize,
        n_vocab: usize,
        mask_dev: Option<CudaSlice<u32>>,
        mask_words: usize,
    ) -> Result<(GraphSession, u32), Box<dyn std::error::Error>> {
        let embd_gpu = embd_gpu_owned;
        let bucket_max = cache.pos + max_new + 1;
        let (graph, plan, seg_end) = self.graph_capture_segment_masked(
            e, &mut cache, &mut gs, &embd_gpu, qt, row_bytes, n_vocab, bucket_max,
            mask_dev.as_ref().map(|d| (d, mask_words)))?;
        let first = e.dtoh_u32_one(&gs.token_d)?;
        Ok((GraphSession {
            gs, cache, embd_gpu, graph, plan, bucket_max, seg_end, qt, row_bytes, n_vocab,
            mask_dev, mask_words,
        }, first))
    }

    /// Device-counter full-attention decode (CUDA-GRAPH-PLAN Phase 2): clone of `full_attn_decode`
    /// using the `_dc` KV-append (write slot from `kvl.len_d`) + `_dc` fa_decode (t_kv from `kvl.len_d`
    /// after inc), and the resident device rope `pos_d`. Bit-identical to `full_attn_decode` (the
    /// `_dc` kernels reproduce the same math; fa_decode_dc with bucket_max==t_kv reproduces the same
    /// n_splits/per/combine). Advances `kvl.len`/`kvl.len_d`.
    pub(crate) fn full_attn_decode_dc(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        // eager-mirror path: advance host counters and size n_splits from the live t_kv (bit-identical
        // to fa_decode). The capture path uses full_attn_decode_dc_cap (fixed bucket_max, no host
        // advance, full-buffer K/V view).
        self.full_attn_decode_dc_inner(e, fa, h, None, pos_d, cache, il, None)
    }

    /// PRE-QUANTIZED-INPUT dc full-attn (device-counter path). See full_attn_decode_pre. BIT-IDENTICAL.
    pub(crate) fn full_attn_decode_dc_pre(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        hq: &CudaSlice<i8>,
        hd: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.full_attn_decode_dc_inner(e, fa, h, Some((hq, hd)), pos_d, cache, il, None)
    }

    /// PRE-QUANTIZED-INPUT CAPTURE dc full-attn (graph path, fixed bucket_max). BIT-IDENTICAL.
    pub(crate) fn full_attn_decode_dc_cap_pre(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        hq: &CudaSlice<i8>,
        hd: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
        bucket_max: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.full_attn_decode_dc_inner(e, fa, h, Some((hq, hd)), pos_d, cache, il, Some(bucket_max))
    }

    /// CAPTURE variant of `full_attn_decode_dc` (CUDA-GRAPH-PLAN Phase 3). `bucket_max` sizes the
    /// fa_decode_dc grid (n_splits) at capture time; the kernel reads the ACTUAL t_kv from the device
    /// counter `kvl.len_d`. Does NOT advance the host `kvl.len` (only the DEVICE counter via inc_seqlen,
    /// which is captured and replays each launch). Views the FULL K/V cache buffer so the kernel may
    /// safely read up to any t_kv within the bucket on replay. Bit-identical to eager when
    /// `bucket_max` yields the same n_splits as eager for the replayed t_kv (the bucket-key contract).
    pub(crate) fn full_attn_decode_dc_cap(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
        bucket_max: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.full_attn_decode_dc_inner(e, fa, h, None, pos_d, cache, il, Some(bucket_max))
    }

    fn full_attn_decode_dc_inner(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
        pos_d: &CudaSlice<i32>,
        cache: &mut Cache,
        il: usize,
        cap_bucket_max: Option<usize>,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let n_head = cfg.n_head as usize;
        let n_head_kv = cfg.n_head_kv as usize;
        let head_dim = cfg.head_dim_k as usize;
        let eps = cfg.rms_eps;
        let scale = 1.0 / (head_dim as f32).sqrt();

        let n_embd = cfg.n_embd as usize;
        // Q8 TRUNK-FUSION (2026-07-05): wq+wk+wv share input h — on the 35B every full-attn
        // projection is Q8_0, so ONE fused3 launch (block-offset split, out_f 8192/512/512)
        // replaces three launch-latency-class m=1 launches. BIT-IDENTICAL per (tensor,row) to
        // the three matmul_pre MMVQ dispatches (same kernel body). MEMRA_Q8_DUAL=0 rollback.
        let qkv_fused = |e: &Engine,
                         hq: &CudaSlice<i8>,
                         hd: &CudaSlice<f32>|
         -> Result<
            (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>),
            Box<dyn std::error::Error>,
        > {
            if let Some((qf, k, v)) = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)? {
                return Ok((qf, k, v));
            }
            Ok((
                e.matmul_pre(&fa.wq, hq, hd, h, 1)?,
                e.matmul_pre(&fa.wk, hq, hd, h, 1)?,
                e.matmul_pre(&fa.wv, hq, hd, h, 1)?,
            ))
        };
        let (qf, mut k, v) =
            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
                match pre_q {
                    Some((hq, hd)) => qkv_fused(e, hq, hd)?,
                    None => {
                        let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
                        qkv_fused(e, &hq, &hd)?
                    }
                }
            } else {
                (
                    e.matmul(&fa.wq, h, 1)?,
                    e.matmul(&fa.wk, h, 1)?,
                    e.matmul(&fa.wv, h, 1)?,
                )
            };
        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
        let gated = self.cfg.attn_out_gate();
        let (mut q, gate) = if gated {
            let mut q = e.uninit(n_head * head_dim)?;
            let mut gate = e.uninit(n_head * head_dim)?;
            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
            (q, Some(gate))
        } else {
            (qf, None)
        };

        let mut qn = e.uninit(n_head * head_dim)?;
        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
        q = qn;
        let mut kn = e.uninit(n_head_kv * head_dim)?;
        e.rms_norm(
            &k,
            fa.k_norm.float_data(),
            &mut kn,
            head_dim,
            n_head_kv,
            eps,
        )?;
        k = kn;
        let rope_dims = cfg.rope_dim_count as usize;
        // rope pos from the resident device counter (no per-step host upload).
        e.rope_neox(
            &mut q,
            pos_d,
            head_dim,
            rope_dims,
            n_head,
            1,
            cfg.rope_freq_base,
            1.0,
        )?;
        e.rope_neox(
            &mut k,
            pos_d,
            head_dim,
            rope_dims,
            n_head_kv,
            1,
            cfg.rope_freq_base,
            1.0,
        )?;

        let kvl = cache.kv[il].as_mut().unwrap();
        // (1) append at the device write slot kvl.len_d (== old len).
        e.append_kv_quantized_dc(
            &k,
            &v,
            &mut kvl.k,
            &mut kvl.v,
            &kvl.len_d,
            kvl.kv_dim_k,
            kvl.kv_dim_v,
            kvl.k_tok_bytes,
            kvl.v_tok_bytes,
            crate::Engine::kv_fp8_on(),
        )?;
        // (2) advance the device counter: kvl.len_d now holds new len == t_kv.
        e.inc_seqlen(&mut kvl.len_d)?;
        // n_splits sizing + K/V view extent:
        //  - eager path (cap_bucket_max==None): advance host len; size from live t_kv == bit-identical
        //    to fa_decode; view exactly t_kv*tok_bytes.
        //  - capture path (Some(bucket_max)): DO NOT touch host len (replay advances only the device
        //    counter); size n_splits from bucket_max; view the FULL cache buffer so any in-bucket t_kv
        //    is in range on replay.
        let (bucket_max, k_view, v_view) = match cap_bucket_max {
            None => {
                kvl.len += 1;
                let t_kv = kvl.len;
                (
                    t_kv,
                    e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes),
                    e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes),
                )
            }
            Some(bm) => (
                bm,
                e.view_u8(&kvl.k, kvl.k.len()),
                e.view_u8(&kvl.v, kvl.v.len()),
            ),
        };
        let (ktb, vtb) = (kvl.k_tok_bytes, kvl.v_tok_bytes);
        let mut attn = e.uninit(n_head * head_dim)?;
        if std::env::var("MEMRA_NOFA").is_ok() {
            return Err(
                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV cache; \
                        unset MEMRA_NOFA to use fa_decode_dc"
                    .into(),
            );
        }
        // (3) fa_decode reads t_kv from kvl.len_d; bucket_max yields the eager n_splits -> bit-identical.
        e.fa_decode_dc(
            &q,
            &k_view,
            &v_view,
            &mut attn,
            head_dim,
            n_head,
            n_head_kv,
            &kvl.len_d,
            bucket_max,
            scale,
            ktb,
            vtb,
            crate::Engine::kv_fp8_on(),
        )?;

        let attn_g = match &gate {
            Some(gate) => {
                let mut gsig = e.uninit(n_head * head_dim)?;
                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
                let mut ag = e.uninit(n_head * head_dim)?;
                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
                ag
            }
            None => attn,
        };
        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
    }

    /// Greedy generation: prime with prompt tokens (decode them in sequence to build state),
    /// then generate `max_new` tokens. Returns the generated token ids. (Back-compat: greedy,
    /// no EOS/stop — used by the decode==prefill validation gate. New code uses `generate_with`.)
    pub fn generate(
        &self,
        e: &Engine,
        prompt: &[u32],
        max_new: usize,
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        let max_ctx = prompt.len() + max_new + 8;
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;
        let mut last_logits = Vec::new();
        // prime: BATCHED cache prime (prime_cache — the prefill-throughput path, the measured #1
        // e2e gap: tokenwise primed at ~102/38 tok/s vs ~2000-5900 tok/s batched). Prompts below
        // PRIME_MIN_T, MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the
        // tokenwise loop. Frozen mixed residency would otherwise transiently stage the missing
        // expert bank through the GPU on every prompt replay.
        let t_prime = std::time::Instant::now();
        let batched_prime = prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
        if batched_prime {
            let (l, _h_seed, _hiddens) = self.prime_cache(e, prompt, &mut cache)?;
            last_logits = l;
        } else {
            for &tok in prompt {
                last_logits = self.decode_step(e, tok, &mut cache)?;
            }
        }
        e.stream().synchronize()?;
        // Harness timing contract: prime wall time published for gen-only throughput math
        // (bench binaries read this right after the call; subtraction-from-total breaks down
        // when prime >> gen — measured ±80% error at 6k-token prompts).
        crate::PRIME_NANOS.store(
            t_prime.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        let mut out = Vec::with_capacity(max_new);
        if self.cfg.gemma4.is_some()
            && let Some(embd_gpu) = self.embd_gpu_try(e) {
            // Graph serving probed FLAT vs this dc loop (2026-07-12, 1.7k N=2: 174.6/174.2 vs
            // 174.5/174.3) — the GRAPH-GATE's +2.5% is over the plain-eager loop, and the dc
            // arc already banked that; the gate (IDENTICAL at every ctx since the wkv
            // capture-arm fix) stays as the correctness harness.
            // DEVICE-COUNTER greedy loop (the dc arc): stream-identical to eager (DC-GATE).
            // E4B rides its own dc step (same trunk fns as its eager chain).
            let n_vocab = self.output.out_features();
            let (qt, rb) = self.embd.qt_and_row_bytes(self.cfg.n_embd as usize);
            for kvl in cache.kv.iter_mut().flatten() {
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            let e4b = self.is_gemma4_e4b();
            // 26B/31B WHOLE-TOKEN GRAPH SERVING door (MEMRA_GEMMA_GRAPH=1): measured FLAT on
            // the 26B (jsonl 2026-07-12) but the 31B carries ~4% launch-gap share (HANDOVER
            // graph-arc note) and was never measured — the plain-short 1.00x cell probe.
            if !e4b && std::env::var("MEMRA_GEMMA_GRAPH").as_deref() == Ok("1") {
                let first = argmax(&last_logits) as u32;
                let (toks, _reason) = self.gemma4_generate_graph(
                    e, cache.pos, first, &mut cache, max_new, &[], |_| true)?;
                out.extend(toks);
                return Ok(out);
            }
            let mut token_d = e.stream().clone_htod(&[argmax(&last_logits) as u32])?;
            let mut pos_d = e.htod_i32(&[cache.pos as i32])?;
            // E4B GRAPH-EXEC-UPDATE SERVING: one capture at bucket=win, per-token fa
            // geometry retune, replay. The 2026-07-12 park ("flat 173.5, stream 64/64") did
            // NOT reproduce — the capture warmups are real self-feeding steps and the old
            // door dropped their 2 tokens (E4B-GRAPH-GATE 3/64). Snapshot/rollback (the 26B
            // graph-loop pattern) fixes the stream; the exec-update kills the bucket-split
            // tax (42 fa launches at 64 splits vs eager's ~ceil(t_kv/8)).
            // DEFAULT: budget-gated ON (2026-07-13 valid-window A/B: steady-state replay
            // beats eager but the one-time capture ~30ms crosses over near 200 tokens —
            // 128tok −1.3%, 400tok +0.9%). MEMRA_E4B_GRAPH=1 forces, =0 kills.
            let win = self.cfg.gemma4.as_ref().map(|g| g.sliding_window as usize).unwrap_or(0);
            let e4b_graph = match std::env::var("MEMRA_E4B_GRAPH").as_deref() {
                Ok("1") => true, Ok("0") => false, _ => max_new >= 256,
            };
            if e4b && cache.pos + max_new + 2 < win && e4b_graph {
                self.gemma4_e4b_graph_exec_loop(
                    e, &mut cache, &mut token_d, &mut pos_d, embd_gpu, qt, rb, n_vocab, win,
                    max_new, usize::MAX, |tok| { out.push(tok); None })?;
                return Ok(out);
            }
            for _ in 0..max_new {
                out.push(e.dtoh_u32(&token_d)?[0]);
                token_d = if e4b {
                    self.gemma4_e4b_decode_step_dc(
                        e, &token_d, &mut pos_d, embd_gpu, qt, rb, &mut cache, n_vocab,
                    )?
                } else {
                    self.gemma4_decode_step_dc(
                        e, &token_d, &mut pos_d, embd_gpu, qt, rb, &mut cache, n_vocab, None,
                    )?
                };
            }
            return Ok(out);
        }
        // QWEN DC-EAGER route (2026-07-15, MEMRA_QWEN_DC=0 seam — mirror of generate_with's
        // serving loop; see the note there. The graph route probed −11% first.)
        static QWEN_DC2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        let qwen_dc = *QWEN_DC2.get_or_init(||
            std::env::var("MEMRA_QWEN_DC").as_deref() != Ok("0"));
        if qwen_dc && max_new > 0
            && let Some(embd_gpu) = self.embd_gpu_try(e) {
            let n_vocab = self.output.out_features();
            let (qt, rb) = self.embd.qt_and_row_bytes(self.cfg.n_embd as usize);
            for kvl in cache.kv.iter_mut().flatten() {
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            let mut pos_d = e.htod_i32(&[cache.pos as i32])?;
            let mut token_d = e.stream().clone_htod(&[argmax(&last_logits) as u32])?;
            for _ in 0..max_new {
                out.push(e.dtoh_u32(&token_d)?[0]);
                token_d = self.decode_step_dc(e, &token_d, &mut pos_d, embd_gpu, qt, rb,
                                              &mut cache, n_vocab)?;
            }
            return Ok(out);
        }
        for _ in 0..max_new {
            let next = argmax(&last_logits) as u32;
            out.push(next);
            last_logits = self.decode_step(e, next, &mut cache)?;
        }
        Ok(out)
    }

    /// E4B whole-token GRAPH-EXEC-UPDATE serving loop (shared by `generate` and
    /// `generate_with`): capture ONE self-feeding dcg step at bucket=`win`, then per token
    /// retune the fa nodes' split geometry to the live eager counts
    /// (`graph_update::fa_apply`) before replaying the instantiated exec.
    ///
    /// The capture's two warmup runs are REAL executions (self-feeding: they consume two
    /// tokens and advance KV/counters) — snapshot/rollback around the capture (the 26B
    /// graph-loop pattern) restores device+host state, or the stream drops those tokens
    /// (E4B-GRAPH-GATE 3/64 break, 2026-07-12). `emit` sees each token BEFORE its
    /// successor's replay; returning `Some(reason)` stops the loop. Caller owns the
    /// under-window gate (`cache.pos + budget + 2 < win`).
    #[allow(clippy::too_many_arguments)]
    fn gemma4_e4b_graph_exec_loop(
        &self, e: &Engine, cache: &mut Cache, token_d: &mut CudaSlice<u32>,
        pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, qt: i32, rb: usize,
        n_vocab: usize, win: usize, budget: usize, ctx_cap: usize,
        mut emit: impl FnMut(u32) -> Option<StopReason>,
    ) -> Result<StopReason, Box<dyn std::error::Error>> {
        // BISECT ARM (MEMRA_E4B_DCG_EAGER=1): run the dcg step EAGERLY per token at the
        // exact live bucket — no capture/replay/exec-update. Separates "the dc-bucket path
        // diverges from dc-eager numerically" from "the replay/update mechanism is wrong".
        if let Ok(m) = std::env::var("MEMRA_E4B_DCG_EAGER") {
            // =1: exact live bucket per token; =2: the capture's fixed win bucket.
            let mut reason = StopReason::MaxNew;
            for _ in 0..budget {
                let tok = e.dtoh_u32_one(token_d)?;
                if let Some(r) = emit(tok) { reason = r; break; }
                if cache.pos >= ctx_cap { reason = StopReason::ContextFull; break; }
                let b = if m == "2" { win } else { cache.pos + 1 };
                self.gemma4_e4b_decode_step_dcg(e, token_d, pos_d, embd_gpu, qt, rb,
                                                cache, n_vocab, b)?;
                cache.pos += 1;
                for kvl in cache.kv.iter_mut().flatten() { kvl.len += 1; }
            }
            return Ok(reason);
        }
        // snapshot device+host state (the 2 capture-warmup runs must leave no residue).
        let snap = cache.snapshot(e)?;
        let pos_save = e.dtoh_i32_one(pos_d)?;
        let len_save: Vec<Option<i32>> = cache.kv.iter()
            .map(|k| k.as_ref().map(|kvl| e.dtoh_i32_one(&kvl.len_d).unwrap())).collect();
        let tok_save = e.dtoh_u32_one(token_d)?;
        let (graph, keeper) = e.capture_graph_retained(|e| {
            self.gemma4_e4b_decode_step_dcg(e, token_d, pos_d, embd_gpu, qt, rb,
                                            cache, n_vocab, win)
        })?;
        cache.rollback(e, &snap, 0)?;
        e.set_i32_one(pos_d, pos_save)?;
        for (il, ls) in len_save.iter().enumerate() {
            if let (Some(kvl), Some(v)) = (cache.kv[il].as_mut(), ls) {
                e.set_i32_one(&mut kvl.len_d, *v)?;
            }
        }
        e.set_u32_one(token_d, tok_save)?;
        let mut plan = crate::graph_update::fa_plan(&graph)?;
        if std::env::var("MEMRA_GRAPH_NODES_DUMP").as_deref() == Ok("1") {
            let nodes = crate::graph_update::kernel_nodes(&graph)?;
            let mut counts: std::collections::BTreeMap<String, (usize, (u32, u32, u32))> =
                std::collections::BTreeMap::new();
            for n in &nodes {
                counts.entry(n.name.clone())
                    .or_insert((0, (n.params.gridDimX, n.params.gridDimY, n.params.gridDimZ)))
                    .0 += 1;
            }
            eprintln!("[graph-nodes] {} kernel nodes, {} fa update units (bucket={win})",
                      nodes.len(), plan.len());
            for (name, (c, grid)) in &counts {
                eprintln!("[graph-nodes]   {c:4}x {name} grid={grid:?}");
            }
        }
        let mut reason = StopReason::MaxNew;
        let timing = std::env::var("MEMRA_E4B_GRAPH_TIMING").as_deref() == Ok("1");
        let (mut t_dtoh, mut t_apply, mut t_launch) =
            (std::time::Duration::ZERO, std::time::Duration::ZERO, std::time::Duration::ZERO);
        for _ in 0..budget {
            let t0 = std::time::Instant::now();
            let tok = e.dtoh_u32_one(token_d)?;
            let t1 = std::time::Instant::now();
            if let Some(r) = emit(tok) { reason = r; break; }
            if cache.pos >= ctx_cap { reason = StopReason::ContextFull; break; }
            // live t_kv AFTER this replay's in-graph append = pos + 1.
            crate::graph_update::fa_apply(&graph, &mut plan, cache.pos + 1,
                                          crate::fa_split_keys)?;
            let t2 = std::time::Instant::now();
            graph.launch()?;
            if timing {
                let t3 = std::time::Instant::now();
                t_dtoh += t1 - t0; t_apply += t2 - t1; t_launch += t3 - t2;
            }
            cache.pos += 1;
            for kvl in cache.kv.iter_mut().flatten() { kvl.len += 1; }
        }
        if timing {
            eprintln!("[e4b-graph timing] dtoh(sync-wait) {:?} apply {:?} launch {:?}",
                      t_dtoh, t_apply, t_launch);
        }
        drop(keeper);   // capture-retained transients must outlive every replay
        Ok(reason)
    }

    /// The reusable serving generation API (BASE-3). Primes the prompt, then samples up to
    /// `params.max_new` tokens, stopping on EOS, any stop-token, or the context-length guard.
    /// Calls `on_token(id)` after each emitted token (for streaming; return `false` to stop early).
    /// Returns `GenOutput { tokens, stop_reason }`. Does NOT detokenize — the caller (which owns
    /// the tokenizer) handles text + stop-STRING matching on the detokenized tail.
    pub fn generate_with<F: FnMut(u32) -> bool>(
        &self,
        e: &Engine,
        prompt: &[u32],
        params: &GenParams,
        sampler: &mut crate::sampler::Sampler,
        mut on_token: F,
    ) -> Result<GenOutput, Box<dyn std::error::Error>> {
        // Context guard: prompt + generated must fit max_ctx (caller-supplied or model default).
        let ctx_cap = params.max_ctx.unwrap_or(prompt.len() + params.max_new + 8);
        if prompt.len() >= ctx_cap {
            return Ok(GenOutput {
                tokens: Vec::new(),
                stop_reason: StopReason::ContextFull,
            });
        }
        let room = ctx_cap - prompt.len();
        let budget = params.max_new.min(room);

        let mut cache = Cache::new(e, &self.cfg, ctx_cap)?;
        let mut last_logits = Vec::new();
        // BATCHED PRIME (2026-07-06 fix — generate_with was still tokenwise! run-gen's "decode"
        // numbers folded a ~40-100 tok/s tokenwise prime into the rate) + PRIME_NANOS contract.
        // Frozen Hy3 CPU/GPU expert serving is the deliberate exception: its batched MoE path
        // bypasses the CPU tier and rereads the spilled expert bank.
        let t_prime = std::time::Instant::now();
        let batched = prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
        if batched {
            let (l, _h, _x) = self.prime_cache(e, prompt, &mut cache)?;
            last_logits = l;
            for &tok in prompt {
                sampler.accept(tok);
            }
        } else {
            for &tok in prompt {
                last_logits = self.decode_step(e, tok, &mut cache)?;
                sampler.accept(tok);
            }
        }
        e.stream().synchronize()?;
        crate::PRIME_NANOS.store(
            t_prime.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        // MEMRA_PROFILE_GEN=2: profiler capture starts HERE — after the prime — so an
        // `nsys -c cudaProfilerApi` capture contains ONLY the decode loop (the run-spec
        // MEMRA_PROFILE_SPEC=2 pattern; =1 in run_gen brackets prime+decode).
        if std::env::var("MEMRA_PROFILE_GEN").as_deref() == Ok("2") {
            unsafe extern "C" {
                fn cudaProfilerStart() -> i32;
            }
            unsafe {
                cudaProfilerStart();
            }
        }
        let mut out = Vec::with_capacity(budget);
        let mut reason = StopReason::MaxNew;
        // gemma4 DEVICE-COUNTER greedy serving loop (the dc arc): token/pos/kv-lens live in
        // device counters, argmax on device — host sees 4B/token. Stream-identical to the
        // eager chain (DC-GATE). Penalties/temp fall through to the host-logits loop.
        if self.cfg.gemma4.is_some()
            && sampler.is_greedy() && sampler.penalty_last_n() == 0
            && let Some(embd_gpu) = self.embd_gpu_try(e) {
            let n_vocab = self.output.out_features();
            let (qt, rb) = self.embd.qt_and_row_bytes(self.cfg.n_embd as usize);
            for kvl in cache.kv.iter_mut().flatten() {
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            let first = crate::forward::argmax(&last_logits) as u32;
            let e4b = self.is_gemma4_e4b();
            let mut token_d = e.stream().clone_htod(&[first])?;
            let mut pos_d = e.htod_i32(&[cache.pos as i32])?;
            // E4B GRAPH-EXEC-UPDATE serving door (under-window regime) — mirror of the
            // `generate` door incl the budget-gated default; run-gen/serving measure here.
            let win = self.cfg.gemma4.as_ref().map(|g| g.sliding_window as usize).unwrap_or(0);
            let e4b_graph = match std::env::var("MEMRA_E4B_GRAPH").as_deref() {
                Ok("1") => true, Ok("0") => false, _ => budget >= 256,
            };
            if e4b && cache.pos + budget + 2 < win && e4b_graph {
                let (out_cell, sampler_cell) = (&mut out, &mut *sampler);
                let reason = self.gemma4_e4b_graph_exec_loop(
                    e, &mut cache, &mut token_d, &mut pos_d, embd_gpu, qt, rb, n_vocab, win,
                    budget, ctx_cap, |tok| {
                        sampler_cell.accept(tok);
                        out_cell.push(tok);
                        if params.eos.contains(&tok) { return Some(StopReason::Eos); }
                        if !on_token(tok) { return Some(StopReason::Callback); }
                        None
                    })?;
                return Ok(GenOutput { tokens: out, stop_reason: reason });
            }
            // 12B/31B WHOLE-TOKEN GRAPH door (MEMRA_GEMMA_GRAPH=1), mirrored from `generate`:
            // run-gen/serving measure THIS path, and the `generate` door never covered it —
            // the 2026-07-22 graph A/B read flat because the env engaged nothing here.
            if !e4b && std::env::var("MEMRA_GEMMA_GRAPH").as_deref() == Ok("1") {
                let (out_cell, sampler_cell) = (&mut out, &mut *sampler);
                let eos = params.eos.clone();
                let (toks, greason) = self.gemma4_generate_graph(
                    e, cache.pos, first, &mut cache, budget, &eos, |tok| {
                        sampler_cell.accept(tok);
                        out_cell.push(tok);
                        on_token(tok)
                    })?;
                let _ = toks;
                return Ok(GenOutput { tokens: out, stop_reason: greason });
            }
            let mut next = first;
            for _ in 0..budget {
                sampler.accept(next);
                out.push(next);
                if params.eos.contains(&next) {
                    reason = StopReason::Eos;
                    break;
                }
                if !on_token(next) {
                    reason = StopReason::Callback;
                    break;
                }
                if cache.pos >= ctx_cap {
                    reason = StopReason::ContextFull;
                    break;
                }
                token_d = if e4b {
                    self.gemma4_e4b_decode_step_dc(
                        e, &token_d, &mut pos_d, embd_gpu, qt, rb, &mut cache, n_vocab,
                    )?
                } else {
                    self.gemma4_decode_step_dc(
                        e, &token_d, &mut pos_d, embd_gpu, qt, rb, &mut cache, n_vocab, None,
                    )?
                };
                next = e.dtoh_u32(&token_d)?[0];
            }
            return Ok(GenOutput {
                tokens: out,
                stop_reason: reason,
            });
        }
        // QWEN DC-EAGER serving loop (2026-07-15, MEMRA_QWEN_DC=0 seam — the gemma dc-arc
        // pattern): the eager tail dtoh'd the FULL VOCAB logits + host-argmax'd every
        // token (the duty map's 10.3%-of-wall gap at 13% DRAM duty). decode_step_dc keeps
        // the token id + argmax device-resident — 4B/token host traffic, same tuned eager
        // kernels. Greedy + no-penalty only (sampling needs host logits).
        // (The CUDA-graph route was probed first and read −11%: the replay's dc-fa family
        // + capture rungs lag the tuned eager lanes; jsonl 2026-07-15.)
        static QWEN_DC: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        let qwen_dc = *QWEN_DC.get_or_init(||
            std::env::var("MEMRA_QWEN_DC").as_deref() != Ok("0"));
        if qwen_dc && sampler.is_greedy() && sampler.penalty_last_n() == 0 && budget > 0
            && let Some(embd_gpu) = self.embd_gpu_try(e) {
            let n_vocab = self.output.out_features();
            let (qt, rb) = self.embd.qt_and_row_bytes(self.cfg.n_embd as usize);
            for kvl in cache.kv.iter_mut().flatten() {
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            let mut pos_d = e.htod_i32(&[cache.pos as i32])?;
            let mut token_d = e.stream().clone_htod(
                &[crate::forward::argmax(&last_logits) as u32])?;
            // HYBRID GRAPH DOOR (round 35): graph_decode_loop over the batched-prime
            // cache — the E4B graph-exec door's hybrid mirror. Counters (pos_d/token_d/
            // len_d) synced above; event tracking is engine-default-OFF so capture over
            // these buffers is legal. PROMOTED default-ON at budget >= 256 (the E4B
            // door's amortization rule): official-shape A/B interleaved x5 = eager 190.3
            // -> graph 220.7 tok/s (+16.0%, 5/5, spread ±0.1); 128-tok stream IDENTICAL;
            // graph-decode-gate 256 steps x 16 buckets BIT-IDENTICAL. This REFUTES the
            // 2026-07-15 "-11%" qwen-graph verdict — it predated the exec-update rework
            // and the 07-26 FA family (stale-verdict law, round 35). =0 reverts.
            // Default ON at budget >= 256 on BOTH arches (unified-merge resolution,
            // 2026-07-30): main shipped this door budget-keyed on sm_120a (52222ddd,
            // E4B graph door) and every 5090 board row since measured with it; the H100
            // lane measured +16% x5. The branch-era arch-gate (79395a3e) cited the
            // stale 2026-07-15 "-11%" verdict, which predates main's promotion — the
            // rig-divergence law protects main's SHIPPED default, so the gate came off.
            // MEMRA_GEN_GRAPH=1 opts in anywhere; =0 reverts anywhere.
            let gen_graph = match std::env::var("MEMRA_GEN_GRAPH").as_deref() {
                Ok("1") => true,
                Ok("0") => false,
                _ => budget >= 256,
            };
            // SLRU expert cache is capture-ILLEGAL: a cache miss drains/H2Ds on the compute
            // stream mid-decode, which CUDA forbids while capturing (Ornith-35B Q4_K_M on the
            // 24GB rig died with CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED, 2026-08-01 — any MoE
            // model whose experts overflow the residency budget hit this at budget >= 256).
            // The door only opens with every MoE layer's experts device-resident; =1 cannot
            // legalize a capture, so this closes the forced door too.
            let moe_resident = self.layers.iter().all(|l| match &l.ffn {
                crate::hybrid::Ffn::Moe(m) => m.dev_exps.is_some(),
                _ => true,
            });
            if gen_graph && !moe_resident {
                static NOTICE: std::sync::Once = std::sync::Once::new();
                NOTICE.call_once(|| eprintln!(
                    "[gen-graph] door CLOSED: MoE experts on the SLRU cache path \
                     (capture-illegal) — eager decode"
                ));
            }
            if gen_graph && moe_resident && budget > 0 {
                let head_dim = self.cfg.head_dim_k as usize;
                let mut gs = GraphDecodeState::new(e)?;
                gs.pos_d = pos_d;
                gs.token_d = token_d;
                let (out_cell, sampler_cell) = (&mut out, &mut *sampler);
                let reason = self.graph_decode_loop(
                    e, &mut gs, &mut cache, embd_gpu, qt, rb, head_dim, budget, |tok| {
                        sampler_cell.accept(tok);
                        out_cell.push(tok);
                        if params.eos.contains(&tok) { return Some(StopReason::Eos); }
                        if !on_token(tok) { return Some(StopReason::Callback); }
                        None
                    })?;
                return Ok(GenOutput { tokens: out, stop_reason: reason });
            }
            let mut next = e.dtoh_u32(&token_d)?[0];
            for _ in 0..budget {
                sampler.accept(next);
                out.push(next);
                if params.eos.contains(&next) { reason = StopReason::Eos; break; }
                if !on_token(next) { reason = StopReason::Callback; break; }
                if cache.pos >= ctx_cap { reason = StopReason::ContextFull; break; }
                token_d = self.decode_step_dc(e, &token_d, &mut pos_d, embd_gpu, qt, rb,
                                              &mut cache, n_vocab)?;
                next = e.dtoh_u32(&token_d)?[0];
            }
            return Ok(GenOutput { tokens: out, stop_reason: reason });
        }
        for _ in 0..budget {
            let next = sampler.sample(&last_logits);
            sampler.accept(next);
            out.push(next);
            if params.eos.contains(&next) {
                reason = StopReason::Eos;
                break;
            }
            if !on_token(next) {
                reason = StopReason::Callback;
                break;
            }
            if cache.pos >= ctx_cap {
                reason = StopReason::ContextFull;
                break;
            }
            last_logits = self.decode_step(e, next, &mut cache)?;
        }
        Ok(GenOutput {
            tokens: out,
            stop_reason: reason,
        })
    }

    /// Full-attention decode: project q/gate/k/v for the new token, QK-norm, RoPE at pos,
    /// append k,v to the layer KV cache, attend over the full [0..=pos] context.
    pub(crate) fn full_attn_decode(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        pos: usize,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.full_attn_decode_pre(e, fa, h, None, pos_d, pos, cache, il)
    }

    /// PRE-QUANTIZED-INPUT eager full-attn (attn-input NORM-FUSION lever): caller passes the
    /// attn-normed activation already q8_1 `(hq,hd)` (rms_norm_q8_1) -> skips internal quantize_q8_1.
    /// `None` = quantize h here (the spec / non-fused path). BIT-IDENTICAL.
    pub(crate) fn full_attn_decode_pre(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        h: &CudaSlice<f32>,
        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
        pos_d: &CudaSlice<i32>,
        pos: usize,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let n_head = cfg.n_head as usize;
        let n_head_kv = cfg.n_head_kv as usize;
        let head_dim = cfg.head_dim_k as usize;
        let eps = cfg.rms_eps;
        let scale = 1.0 / (head_dim as f32).sqrt();

        // LATENCY-HIDING (MEMRA_KV_PREFETCH=1): warm this layer's KV stream into L2 while the
        // q/k/v projections run ahead of the fa (fa is latency-bound; its lines land warm).
        // Value-free scheduling — no numeric config change.
        static KV_PF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        if *KV_PF.get_or_init(|| std::env::var("MEMRA_KV_PREFETCH").as_deref() == Ok("1")) {
            let kvl = cache.kv[il].as_ref().unwrap();
            let t_kv = kvl.len + 1;
            e.prefetch_l2(&kvl.k, t_kv * kvl.k_tok_bytes)?;
            e.prefetch_l2(&kvl.v, t_kv * kvl.v_tok_bytes)?;
        }

        // wq|wk|wv all take the same input `h` (in_f = n_embd) — quantize q8_1 ONCE, feed all three.
        // Q8 TRUNK-FUSION: on Q8_0 trunks (35B) the three fold into ONE fused3 launch (same MMVQ
        // body per (tensor,row) — bit-identical; see full_attn_decode_dc_inner). MEMRA_Q8_DUAL=0 off.
        let n_embd = cfg.n_embd as usize;
        let qkv_fused = |e: &Engine,
                         hq: &CudaSlice<i8>,
                         hd: &CudaSlice<f32>|
         -> Result<
            (CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>),
            Box<dyn std::error::Error>,
        > {
            if let Some((qf, k, v)) = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)? {
                return Ok((qf, k, v));
            }
            Ok((
                e.matmul_pre(&fa.wq, hq, hd, h, 1)?,
                e.matmul_pre(&fa.wk, hq, hd, h, 1)?,
                e.matmul_pre(&fa.wv, hq, hd, h, 1)?,
            ))
        };
        let (qf, mut k, v) =
            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
                match pre_q {
                    Some((hq, hd)) => qkv_fused(e, hq, hd)?,
                    None => {
                        let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
                        qkv_fused(e, &hq, &hd)?
                    }
                }
            } else {
                (
                    e.matmul(&fa.wq, h, 1)?,
                    e.matmul(&fa.wk, h, 1)?,
                    e.matmul(&fa.wv, h, 1)?,
                )
            };
        // q|gate fused: [2*head_dim per head]. Split on-device (no dtoh/host-loop/htod).
        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
        let gated = self.cfg.attn_out_gate();
        let (mut q, gate) = if gated {
            let mut q = e.uninit(n_head * head_dim)?;
            let mut gate = e.uninit(n_head * head_dim)?;
            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
            (q, Some(gate))
        } else {
            (qf, None)
        };

        // QK-norm + RoPE at position `pos`
        let mut qn = e.uninit(n_head * head_dim)?;
        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
        q = qn;
        let mut kn = e.uninit(n_head_kv * head_dim)?;
        e.rms_norm(
            &k,
            fa.k_norm.float_data(),
            &mut kn,
            head_dim,
            n_head_kv,
            eps,
        )?;
        k = kn;
        let rope_dims = cfg.rope_dim_count as usize;
        e.rope_neox(
            &mut q,
            pos_d,
            head_dim,
            rope_dims,
            n_head,
            1,
            cfg.rope_freq_base,
            1.0,
        )?;
        e.rope_neox(
            &mut k,
            pos_d,
            head_dim,
            rope_dims,
            n_head_kv,
            1,
            cfg.rope_freq_base,
            1.0,
        )?;

        // append k,v into the RESIDENT GPU QUANTIZED KV cache at the current position (q8_0 K /
        // q5_1 V, on-device append-quantize kernel; no host round-trip). KVQUANT-PLAN §C/E2.
        let kvl = cache.kv[il].as_mut().unwrap();
        e.append_kv_quantized(
            &k,
            &v,
            &mut kvl.k,
            &mut kvl.v,
            kvl.len,
            kvl.kv_dim_k,
            kvl.kv_dim_v,
            kvl.k_tok_bytes,
            kvl.v_tok_bytes,
            crate::Engine::kv_fp8_on(),
        )?;
        kvl.len += 1;
        let t_kv = kvl.len;

        // attend: q[hd,nh,1] over the resident byte K/V (view first t_kv*tok_bytes BYTES).
        let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
        let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
        let (ktb, vtb) = (kvl.k_tok_bytes, kvl.v_tok_bytes);
        let mut attn = e.uninit(n_head * head_dim)?;
        if std::env::var("MEMRA_NOFA").is_ok() {
            return Err(
                "MEMRA_NOFA (naive f32 SDPA) is incompatible with the quantized KV cache; \
                        unset MEMRA_NOFA to use fa_decode"
                    .into(),
            );
        }
        e.fa_decode_kvmod(
            &q,
            &k_view,
            &v_view,
            &mut attn,
            head_dim,
            n_head,
            n_head_kv,
            t_kv,
            scale,
            ktb,
            vtb,
            crate::Engine::kv_fp8_on(),
        )?;
        let _ = pos;

        // output gate: attn * sigmoid(gate), then o-proj
        let attn_g = match &gate {
            Some(gate) => {
                let mut gsig = e.uninit(n_head * head_dim)?;
                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
                let mut ag = e.uninit(n_head * head_dim)?;
                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
                ag
            }
            None => attn,
        };
        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
    }

    /// BATCHED full-attention decode over `m` independent streams (one token each).
    ///
    /// Generic m-band primitive, not lockstep-specific: any caller holding `m` streams at the
    /// same layer (multi-stream decode, a continuous-batching serve loop) can use it. The split
    /// follows what the hardware cares about — WEIGHT-BOUND work runs once at `m` because all
    /// streams share the same projection weights (one weight read serves `m` tokens instead of
    /// `m` reads), while KV-BOUND work stays per stream because each stream owns its own cache.
    ///
    /// Bit-identity with the per-stream path holds by construction: `quantize_q8_1` and
    /// `rms_norm` are per-row, `rope_neox` takes a per-token position vector, the fused3/matmul
    /// m-band kernels are the same ones spec verify is gated on, and attention itself is
    /// untouched per stream.
    ///
    /// `xcat` is `[m, n_embd]` normed activations; `pos_cat` is the `m` rope positions;
    /// returns `[m, n_embd]` attention outputs.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn full_attn_decode_batched(
        &self,
        e: &Engine,
        fa: &FullAttnLayer,
        xcat: &CudaSlice<f32>,
        m: usize,
        pos_cat: &CudaSlice<i32>,
        caches: &mut [Cache],
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let n_head = cfg.n_head as usize;
        let n_head_kv = cfg.n_head_kv as usize;
        let head_dim = cfg.head_dim_k as usize;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let scale = 1.0 / (head_dim as f32).sqrt();
        let q_row = n_head * head_dim;
        let kv_row = n_head_kv * head_dim;

        // --- weight-bound: one quantize + one q/k/v projection for all m streams ---
        let (hq, hd) = e.quantize_q8_1(xcat, m, n_embd)?;
        let use_q8 =
            e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
        let (qf, mut k, v) = if use_q8 {
            match e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, &hq, &hd, m)? {
                Some(trio) => trio,
                None => (
                    e.matmul_pre(&fa.wq, &hq, &hd, xcat, m)?,
                    e.matmul_pre(&fa.wk, &hq, &hd, xcat, m)?,
                    e.matmul_pre(&fa.wv, &hq, &hd, xcat, m)?,
                ),
            }
        } else {
            (
                e.matmul(&fa.wq, xcat, m)?,
                e.matmul(&fa.wk, xcat, m)?,
                e.matmul(&fa.wv, xcat, m)?,
            )
        };

        // --- elementwise: batched by treating the m streams as extra rows/tokens ---
        let gated = cfg.attn_out_gate();
        let (mut q, gate) = if gated {
            let mut q = e.uninit(m * q_row)?;
            let mut gate = e.uninit(m * q_row)?;
            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, m)?;
            (q, Some(gate))
        } else {
            (qf, None)
        };
        let mut qn = e.uninit(m * q_row)?;
        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head * m, eps)?;
        q = qn;
        let mut kn = e.uninit(m * kv_row)?;
        e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, n_head_kv * m, eps)?;
        k = kn;
        let rope_dims = cfg.rope_dim_count as usize;
        e.rope_neox(&mut q, pos_cat, head_dim, rope_dims, n_head, m, cfg.rope_freq_base, 1.0)?;
        e.rope_neox(&mut k, pos_cat, head_dim, rope_dims, n_head_kv, m, cfg.rope_freq_base, 1.0)?;

        // --- KV-bound: each stream appends to and attends over its own cache ---
        let mut attn_cat = e.uninit(m * q_row)?;
        let mut q_s = e.uninit(q_row)?;
        let mut k_s = e.uninit(kv_row)?;
        let mut v_s = e.uninit(kv_row)?;
        for (s, cache) in caches.iter_mut().enumerate().take(m) {
            e.copy_view_into(&mut k_s, 0, &k.slice(s * kv_row..(s + 1) * kv_row), kv_row)?;
            e.copy_view_into(&mut v_s, 0, &v.slice(s * kv_row..(s + 1) * kv_row), kv_row)?;
            e.copy_view_into(&mut q_s, 0, &q.slice(s * q_row..(s + 1) * q_row), q_row)?;
            let kvl = cache.kv[il].as_mut().unwrap();
            e.append_kv_quantized(
                &k_s,
                &v_s,
                &mut kvl.k,
                &mut kvl.v,
                kvl.len,
                kvl.kv_dim_k,
                kvl.kv_dim_v,
                kvl.k_tok_bytes,
                kvl.v_tok_bytes,
                crate::Engine::kv_fp8_on(),
            )?;
            kvl.len += 1;
            let t_kv = kvl.len;
            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
            let mut attn = e.uninit(q_row)?;
            e.fa_decode_kvmod(
                &q_s,
                &k_view,
                &v_view,
                &mut attn,
                head_dim,
                n_head,
                n_head_kv,
                t_kv,
                scale,
                kvl.k_tok_bytes,
                kvl.v_tok_bytes,
                crate::Engine::kv_fp8_on(),
            )?;
            e.copy_into(&mut attn_cat, s * q_row, &attn, q_row)?;
        }

        // --- weight-bound again: gate epilogue + one output projection for all m streams ---
        let attn_g = match &gate {
            Some(gate) => {
                let mut gsig = e.uninit(m * q_row)?;
                e.sigmoid(gate, &mut gsig, m * q_row)?;
                let mut ag = e.uninit(m * q_row)?;
                e.mul(&attn_cat, &gsig, &mut ag, m * q_row)?;
                ag
            }
            None => attn_cat,
        };
        e.matmul(&fa.wo, &attn_g, m)
    }

    /// Linear-attention decode: conv with ring-buffer state, GDN scan carrying SSM state.
    pub fn linear_attn_decode(
        &self,
        e: &Engine,
        la: &LinearAttnLayer,
        h: &CudaSlice<f32>,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.linear_attn_decode_inner(e, la, h, None, cache, il, false)
    }

    /// PRE-QUANTIZED-INPUT variant (DECODE attn-input NORM-FUSION lever): the caller passes the
    /// post-attn-norm activation ALREADY q8_1-quantized `(hq,hd)` (produced by rms_norm_q8_1, fusing
    /// the attn_norm + the mixer's internal quantize_q8_1). Skips the internal quantize. Caller
    /// GUARANTEES the projections are q8_1-fast. `persistent` selects the capture-safe state plumbing.
    /// BIT-IDENTICAL to linear_attn_decode(h) when (hq,hd)==quantize_q8_1(rms_norm(x)*w).
    pub fn linear_attn_decode_pre(
        &self,
        e: &Engine,
        la: &LinearAttnLayer,
        h: &CudaSlice<f32>,
        hq: &CudaSlice<i8>,
        hd: &CudaSlice<f32>,
        cache: &mut Cache,
        il: usize,
        persistent: bool,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.linear_attn_decode_inner(e, la, h, Some((hq, hd)), cache, il, persistent)
    }

    /// CAPTURE variant of `linear_attn_decode` (CUDA-GRAPH-PLAN Phase 3). The GDN scan needs distinct
    /// in/out SSM-state buffers; the eager path SWAPS a fresh scratch into `rl.ssm_state` (new pointer
    /// each step), which is a CAPTURE HAZARD — the graph bakes capture-time pointers and never re-runs
    /// the host swap, so replay would read a stale state buffer. Here we instead COPY the scratch back
    /// into the STABLE `rl.ssm_state` buffer (memcpy_dtod, captured, same pointers every replay). Math
    /// is identical; only the buffer plumbing differs. `conv_state` is already mutated in place (no
    /// pointer change) so it is capture-safe as-is.
    pub(crate) fn linear_attn_decode_cap(
        &self,
        e: &Engine,
        la: &LinearAttnLayer,
        h: &CudaSlice<f32>,
        cache: &mut Cache,
        il: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        self.linear_attn_decode_inner(e, la, h, None, cache, il, true)
    }

    fn linear_attn_decode_inner(
        &self,
        e: &Engine,
        la: &LinearAttnLayer,
        h: &CudaSlice<f32>,
        pre_q: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
        cache: &mut Cache,
        il: usize,
        persistent_state: bool,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let cfg = &self.cfg;
        let ssm = cfg.ssm.as_ref().unwrap();
        let d_state = ssm.state_size as usize;
        let num_k = ssm.group_count as usize;
        let num_v = ssm.time_step_rank as usize;
        let d_conv = ssm.conv_kernel as usize;
        let head_k = d_state;
        let key_dim = head_k * num_k;
        let value_dim = d_state * num_v;
        let conv_dim = key_dim * 2 + value_dim;
        let eps = cfg.rms_eps;
        let scale = 1.0 / (d_state as f32).sqrt();

        // projections (T=1): wqkv, wqkv_gate, ssm_beta, ssm_alpha ALL take input `h` (in_f = n_embd)
        // -> quantize q8_1 ONCE, feed all four (was 4x redundant quantize_q8_1 of the same row).
        let n_embd = cfg.n_embd as usize;
        let all_fast = e.uses_q8_1_fast(&la.wqkv)
            && e.uses_q8_1_fast(&la.wqkv_gate)
            && e.uses_q8_1_fast(&la.ssm_beta)
            && e.uses_q8_1_fast(&la.ssm_alpha);
        // beta+alpha DUAL fuse (2026-07-05): ssm_beta and ssm_alpha are the same tiny shape
        // ([n_embd -> num_v=32]) — out_f=32 launches are pure launch latency (15-16us each,
        // HANDOVER b4-headroom note). The existing dual mr2 kernel (FFN gate+up) folds them into
        // ONE launch. Bit-identical per row: same MMVQ warp-per-row body, blockIdx.y picks the
        // weight; the separable macro-scale multiply is the same single f32 mul as matmul_pre's
        // in-kernel scale. Falls back to two matmul_pre when ineligible (Float layers 1/2/4 etc).
        let beta_alpha =
            |e: &Engine,
             hq: &CudaSlice<i8>,
             hd: &CudaSlice<f32>|
             -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
                if let Some(((mut b, bs), (mut a, as_))) =
                    e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, hq, hd, 1)?
                {
                    if bs != 1.0 {
                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
                    }
                    if as_ != 1.0 {
                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
                    }
                    return Ok((b, a));
                }
                // Q8_0 twin of the NVFP4 dual (9B GGUFs store ssm_beta/alpha as Q8_0 on most layers):
                // one fused2 launch, bit-identical per row, no macro-scale (q8_0 scale==1.0).
                if let Some((b, a)) = e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, hq, hd)? {
                    return Ok((b, a));
                }
                Ok((
                    e.matmul_pre(&la.ssm_beta, hq, hd, h, 1)?,
                    e.matmul_pre(&la.ssm_alpha, hq, hd, h, 1)?,
                ))
            };
        // Q8 TRUNK-FUSION (2026-07-05): wqkv+wqkv_gate share (hq,hd) and in_f — on the 35B both
        // are Q8_0 (out_f 8192/4096), so ONE fused2 launch replaces the two biggest
        // launch-latency-class m=1 launches of every linear layer. BIT-IDENTICAL per (tensor,row)
        // (same MMVQ body, block-offset split). Falls back per-tensor when ineligible.
        let qkv_pair =
            |e: &Engine,
             hq: &CudaSlice<i8>,
             hd: &CudaSlice<f32>|
             -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
                if let Some((qkv, z)) = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, hq, hd)? {
                    return Ok((qkv, z));
                }
                Ok((
                    e.matmul_pre(&la.wqkv, hq, hd, h, 1)?,
                    e.matmul_pre(&la.wqkv_gate, hq, hd, h, 1)?,
                ))
            };
        let (qkv_mixed, z, beta_raw, alpha) = if all_fast {
            // attn-input NORM-FUSION: use the caller's pre-quantized (hq,hd) when provided (the
            // attn_norm already emitted q8_1 via rms_norm_q8_1), else quantize h here. Bit-identical.
            match pre_q {
                Some((hq, hd)) => {
                    let (b, a) = beta_alpha(e, hq, hd)?;
                    let (qkv, z) = qkv_pair(e, hq, hd)?;
                    (qkv, z, b, a)
                }
                None => {
                    let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
                    let (b, a) = beta_alpha(e, &hq, &hd)?;
                    let (qkv, z) = qkv_pair(e, &hq, &hd)?;
                    (qkv, z, b, a)
                }
            }
        } else {
            // 35B trunk lands HERE: wqkv/wqkv_gate are Q8_0 but ssm_beta/alpha are F32, so
            // all_fast is false. Still fuse the two Q8_0 projections (one quantize + ONE launch
            // instead of two matmuls each re-quantizing h) — matmul_q8_fused2_x is bit-identical
            // to the two m=1 MMVQ dispatches. beta/alpha keep the Float cuBLAS path.
            let (qm, zg) = match e.matmul_q8_fused2_x(&la.wqkv, &la.wqkv_gate, h)? {
                Some(pair) => pair,
                None => (e.matmul(&la.wqkv, h, 1)?, e.matmul(&la.wqkv_gate, h, 1)?),
            };
            (
                qm,
                zg,
                e.matmul(&la.ssm_beta, h, 1)?,
                e.matmul(&la.ssm_alpha, h, 1)?,
            )
        };

        // RANK3 LEVER (conv fuse): assemble [conv_state | new col], depthwise causal conv + SiLU, and
        // roll the ring — ALL in ONE kernel (`ssm_conv1d_fused_decode`), never materializing conv_in
        // to HBM. Replaces conv_assemble_and_roll + ssm_conv1d. Bit-identical (same accumulation order).
        let rl = cache.recur[il].as_mut().unwrap();
        let mut conv_out = e.uninit(conv_dim)?; // [conv_dim, 1] channel-major, SiLU
        e.ssm_conv1d_fused_decode(
            &qkv_mixed,
            &mut rl.conv_state,
            la.ssm_conv1d.float_data(),
            &mut conv_out,
            conv_dim,
            d_conv,
        )?;

        // GDN scan: SSM state stays RESIDENT on GPU. gdn needs DISTINCT in/out state buffers.
        // DECODE DETERMINISM FIX: write the new state into the PERSISTENT spare buffer
        // (`ssm_state_alt`) and PING-PONG the two owned buffers in place — instead of allocating a
        // fresh `state_scratch` via `e.uninit` each step and swapping its pointer in. The old
        // per-step alloc/free churned the stream-ordered async pool; the freed prior state block was
        // recycled by a later step's scratch while a kernel still referenced the swapped-in state,
        // a use-after-reuse that made decode RUN-TO-RUN nondeterministic (two identical primes
        // diverged). With two stable resident buffers there is no per-step alloc/free and no pool
        // churn; the math is byte-identical. `o` is a true per-step output (consumed immediately by
        // gated_rmsnorm below) so it stays a normal scratch.
        let mut o = e.uninit(d_state * num_v)?;
        let n_state = d_state * d_state * num_v;
        let _ = head_k; // head_k == d_state; the kernels use head_k = d_state internally.
                        // GDN PREP, FUSED (2026-07-03): repack + q/k L2-norm + beta sigmoid + g_log in ONE
                        // gdn_prep_decode launch (was 5 tiny serialized kernels: qkv_to_gdn_repack, 2x l2_norm,
                        // sigmoid, gdn_glog). Same math; the L2 reduce runs a 32-lane warp tree instead of the
                        // 256-thread two-level tree (different FP sum order) — gates: argmax + run-spec exactness.
                        // (A prep+scan single-launch fusion — lane/gdnfuse, MEMRA_GDN_FUSE — measured NEUTRAL on
                        // eager decode 2026-07-08 and was removed in the flag audit; rig5090.jsonl holds the record.)
        {
            let mut q_l2 = e.uninit(d_state * num_v)?;
            let mut k_l2 = e.uninit(d_state * num_v)?;
            let mut v_gd = e.uninit(d_state * num_v)?;
            let mut beta = e.uninit(num_v)?;
            let mut g_log = e.uninit(num_v)?;
            e.gdn_prep_decode(
                &conv_out,
                &beta_raw,
                &alpha,
                la.ssm_dt.float_data(),
                la.ssm_a.float_data(),
                &mut q_l2,
                &mut k_l2,
                &mut v_gd,
                &mut beta,
                &mut g_log,
                d_state,
                num_v,
                num_k,
                key_dim,
                eps,
            )?;
            // gdn reads ssm_state, writes the spare ssm_state_alt (disjoint resident fields).
            let RecurLayer {
                ssm_state,
                ssm_state_alt,
                ..
            } = rl;
            e.gdn_scan_s128(
                &q_l2,
                &k_l2,
                &v_gd,
                &g_log,
                &beta,
                ssm_state,
                ssm_state_alt,
                &mut o,
                num_v,
                1,
                scale,
            )?;
        }
        if persistent_state {
            // CAPTURE-safe (graph replay): the canonical state every replay reads must stay at a
            // FIXED pointer (baked into the captured graph). Copy the freshly-written spare BACK
            // into ssm_state (captured, replays each launch). No host pointer swap.
            let alt = std::mem::replace(&mut rl.ssm_state_alt, e.zeros(0)?);
            e.copy_into(&mut rl.ssm_state, 0, &alt, n_state)?;
            rl.ssm_state_alt = alt;
        } else {
            // EAGER: swap the two OWNED resident buffers in place (stable pointers, no alloc/free).
            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
        }

        // gated RMSNorm + ssm_out. FUSED-QUANTIZE ARM (launch-arc): when ssm_out rides the
        // q8_1 fast path, emit q8_1 straight from the gated norm (bit-identical bytes to
        // gated_rmsnorm + quantize_q8_1) and feed matmul_pre — one launch instead of three
        // (norm, quantize, scale all fold away). Fallback = the original f32 chain.
        if e.uses_q8_1_fast(&la.ssm_out) {
            // norm is PER d_state-ROW (num_v rows), exactly like the f32 twin's grid; the q8_1
            // block stream is row-major so the flat bytes feed the matvec unchanged.
            let (gq, gd) =
                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v, eps)?;
            let g0 = e.zeros(0)?;
            return Ok(e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, 1)?);
        }
        let mut gn = e.uninit(d_state * num_v)?;
        e.gated_rmsnorm(
            &o,
            la.ssm_norm.float_data(),
            &z,
            &mut gn,
            d_state,
            num_v,
            eps,
        )?;
        Ok(e.matmul(&la.ssm_out, &gn, 1)?)
    }
}