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::Engine;
use crate::cache::{Cache, RecurLayer};
use crate::forward::argmax;
use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer};
use cudarc::driver::CudaSlice;
use memra_gguf::config::SwigluClamp;
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());
        }
        // GRAPH-LAUNCH HEADROOM GUARD (see spec::GRAPH_LAUNCH_MIN_FREE): a captured
        // session has NO per-tick eager twin — the session IS the graph — so below the
        // driver-free floor the step refuses RECOVERABLY. The worker ends THIS session
        // with an error event and every peer session (and the process) lives; unguarded,
        // cuGraphLaunch segfaults inside libcuda with zero log lines
        // (lane/graph-launch-guard-sweep-20260831, extending step37 defect 3).
        if !crate::spec::graph_launch_headroom_ok(e) {
            static NOTED: std::sync::Once = std::sync::Once::new();
            NOTED.call_once(|| crate::spec::graph_replay_suspended_note("graph-session"));
            return Err(format!(
                "graph-session replay refused: driver free below the {}MB launch floor \
                 (no eager twin for a captured session; ending the session recoverably \
                 instead of segfaulting cuGraphLaunch)",
                crate::spec::GRAPH_LAUNCH_MIN_FREE >> 20
            )
            .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).
    pub(crate) 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.
    #[allow(clippy::too_many_arguments)]
    pub(crate) 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,
        lim: Option<SwigluClamp>,
    ) -> 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. step35's per-layer `lim` is the same problem, same escape hatch:
        // silu_mul_scaled / silu_mul_scaled_q8_1 have no clamped twin.
        if self.cfg.m3.is_some() || lim.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_lim(e, &self.cfg, &gate, &up, 1.0, 1.0, lim, &mut act, n_ff)?;
            return 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 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 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 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)?;
        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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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) => {
                if fa.step_tp_qkv.is_some() {
                    return false;
                }
                // step35 also projects its head-wise GATE from the same attn-normed input, so
                // the fused (h-less) arm requires attn_gate on the q8_1 fast path too — without
                // this the gate matmul would get a zero-length `h`.
                let gate_ok = match &fa.attn_gate {
                    Some(g) => e.uses_q8_1_fast(g),
                    None => true,
                };
                gate_ok
                    && 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: the forward exists (increment 4) but deliberately does NOT claim the fused
            // norm+quantize chain — its first GEMM is wq_a/wkv_a off an f32 hidden, and no MLA
            // parity gate has covered a q8_1 activation path. Keeping this false routes every
            // MLA decode through the unfused arm, which is the gated one.
            Mixer::Mla(_) => false,
            // KDA projects through matmul_group, never the pre-quantized (hq,hd) pair, so the
            // fused norm+quantize chain has no KDA consumer to claim.
            Mixer::Kda(_) => 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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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_path_unimplemented("norm-fused decode"),
                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("norm-fused eager decode"),
            }
        } 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(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache),
                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il),
            }
        }
    }

    /// 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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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_path_unimplemented("device-counter decode"),
                Mixer::Kda(_) => {
                    crate::hybrid::kda_path_unimplemented("norm-fused device-counter decode")
                }
            }
        } 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_path_unimplemented("device-counter decode"),
                Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il),
            }
        }
    }

    /// 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).
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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_path_unimplemented("captured device-counter decode")
                }
                Mixer::Kda(_) => {
                    crate::hybrid::kda_path_unimplemented("norm-fused captured decode")
                }
            }
        } 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_path_unimplemented("captured device-counter decode")
                }
                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("captured decode"),
            }
        }
    }

    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    pub(crate) 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).
                // step35: SAME failure shape, per LAYER. A dense FFN's limit is the SHEXP array
                // (upstream's one build_ffn serves dense + shared expert, llama-graph.cpp:1751).
                let lim = self.cfg.clamp_shexp_at(il as u32);
                let fuse = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
                    && self.cfg.m3.is_none()
                    && lim.is_none()
                    && e.uses_q8_1_fast(ffn_gate)
                    && e.uses_q8_1_fast(ffn_up);
                if fuse {
                    // M2 safety: this q8 arm predates the deferred join and is never taken
                    // in the step37 config — refuse loudly rather than read unwritten mixed.
                    if crate::tp::take_oproj_tail().is_some() {
                        return Err(
                            "oproj tail handoff reached the q8 residual arm — unwired".into()
                        );
                    }
                    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)?;
                    if let Some((a0, a1)) = crate::tp::take_oproj_tail() {
                        e.join_add_rms_norm_raw(a0, a1, x, pnorm, &mut x1, &mut z, n_embd, eps)?;
                    } else {
                        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, lim)?;
                    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).
                // O-PROJ TAIL FUSION M2: when the direct join deferred its add, compose
                // mixed = a0+a1 in-register inside the norm (verbatim program).
                if let Some((a0, a1)) = crate::tp::take_oproj_tail() {
                    e.join_add_rms_norm_raw(a0, a1, x, pnorm, &mut x1, &mut z, n_embd, eps)?;
                } else {
                    e.add_rms_norm(x, mixed, pnorm, &mut x1, &mut z, n_embd, 1, eps)?;
                }
                // Feed the zq8 seam (orndecode B2): two consumers now share this quantize —
                // the dev expert arm (clones at t==1) and the shexp fused2 pair — so the
                // caller-side launch replaces two arm-side ones. Same kernel, same input,
                // byte-identical per the (1, Some) clone contract.
                let zq8 = e.quantize_q8_1(&z, 1, n_embd)?;
                let ffn_out = self.moe_ffn_il_zq8(e, m, &z, Some(&zq8), 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.
    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
    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>> {
        self.refuse_hyper("decode_step_aux")?;
        cache.ensure_usable("decode_step_aux")?;
        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>> {
        self.refuse_hyper("decode_step_hy3_layer0_stages")?;
        cache.ensure_usable("decode_step_hy3_layer0_stages")?;
        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")?,
        ))
    }

    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
    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)?;
            // MEMRA_TG_PROBE_LAYER diagnostics (token-graph bisection): dump layer K's
            // attention output and post-FFN residual through the real eager path.
            if std::env::var("MEMRA_TG_PROBE_LAYER")
                .ok()
                .and_then(|v| v.parse::<usize>().ok())
                == Some(il)
            {
                use std::io::Write;
                let mut xp = e.uninit(n_embd)?;
                e.add(&x1, &ffn_out, &mut xp, n_embd)?;
                let (pm, px) = (e.dtoh(&mixed)?, e.dtoh(&xp)?);
                for (path, data) in [
                    ("/root/eager-probe-mixed.bin", &pm),
                    ("/root/eager-probe-x.bin", &px),
                ] {
                    let mut fo = std::fs::OpenOptions::new()
                        .create(true)
                        .append(true)
                        .open(path)?;
                    for v in data {
                        fo.write_all(&v.to_le_bytes())?;
                    }
                }
            }
            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>> {
        cache.ensure_usable("decode_step_h")?;
        if self.hyper.is_some() {
            return self.decode_step_hyper(e, token, cache);
        }
        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.uses_gemma_program() {
            // 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()) {
            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
            }
            let rt = crate::pp::PpNRt::get(e)?;
            let _walk = rt.acquire_walk("decode_step_h_ppn")?;
            return self.decode_step_h_ppn(e, token, cache, &fence);
        }
        // Whole-token decode graph (step TP graph increment B, MEMRA_STEP_TP_GRAPH=1 +
        // the dcw/fused/router doors): one stitched multi-device launch per token.
        if self.uses_sliding_gated_moe_program()
            && let Some(result) = self.step35_token_graph_step(e, token, cache)?
        {
            return Ok(result);
        }
        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])?;
        // O-PROJ TAIL deferral eligibility: this walk flows into residual_norm_ffn.
        let _oproj_tail_scope = crate::tp::oproj_tail_scope();
        // RANK0 STREAM MERGE (MEMRA_RANK0_MERGE=1): rank0 shares dev0's PRIMARY context
        // with e (cudarc primary_ctx::retain), so its per-layer work can ride e's stream —
        // every e<->rank0 event hop becomes program order. Scheduling-only: BIT-IDENTICAL.
        let _r0merge = if crate::tp::rank0_merge_on() && self.uses_sliding_gated_moe_program() {
            Some(memra_runtime::rank0_redirect_scope(
                e.ctx().ordinal(),
                e.gpu.main_stream().clone(),
                e.gpu.blas(),
            ))
        } else {
            None
        };

        // MEMRA_DEV_EMBED=1 (RECEIPTED NEGATIVE, default OFF): device embed gather from
        // the resident table replaces the host row expand + 16KB pageable H2D with a 4B
        // id write + one gather launch. Bit-identical rows (2G-IDENTITY-MATCH), but
        // interleaved x3 measured FLAT (56.03 vs 56.06) — the host expand fully overlaps
        // GPU work — and the resident table costs ~2.1GB VRAM. Kept as an opt-in seam
        // (a future device-chained loop wants it; do not re-flip without a new receipt).
        static DEV_EMBED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        let dev_embed =
            *DEV_EMBED.get_or_init(|| std::env::var("MEMRA_DEV_EMBED").as_deref() == Ok("1"));
        // embed the single token -> [1, n_embd]
        let mut x = match (dev_embed, self.embd_gpu_try(e)) {
            (true, Some(embd_gpu)) => {
                static TOK_D: std::sync::Mutex<Option<(usize, CudaSlice<u32>)>> =
                    std::sync::Mutex::new(None);
                let mut guard = TOK_D.lock().map_err(|_| "dev-embed lock is poisoned")?;
                if guard.as_ref().is_none_or(|(d, _)| *d != e.ctx().ordinal()) {
                    *guard = Some((e.ctx().ordinal(), e.stream().clone_htod(&[0u32])?));
                }
                let (_, tok_d) = guard.as_mut().expect("armed above");
                e.set_u32_one(tok_d, token)?;
                let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
                e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?
            }
            _ => 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.
        // MEMRA_STEP_TP_TIMING=1: whole-token bucket split of the eager decode walk — mixer vs
        // FFN totals, the EP-tail layers (>= trunk-2) separated, plus the head. Each lap syncs
        // e's stream, so async work bills to the section that queued it. Diagnostic only.
        static B_MIX: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        static B_FFN: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        static B_MIX_TAIL: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        static B_FFN_TAIL: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        static B_HEAD: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        static B_TOKENS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        let timing = std::env::var("MEMRA_STEP_TP_TIMING").as_deref() == Ok("1");
        let lap = |timer: &std::sync::atomic::AtomicU64,
                   started: &mut Option<std::time::Instant>|
         -> Result<(), Box<dyn std::error::Error>> {
            let Some(start) = started.as_mut() else {
                return Ok(());
            };
            e.stream().synchronize()?;
            timer.fetch_add(
                start.elapsed().as_nanos() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
            *start = std::time::Instant::now();
            Ok(())
        };
        let mut lap_start = timing.then(std::time::Instant::now);
        let tail_from = self.layers.len().saturating_sub(2);
        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();
            // FUSION #2f (bf16-mixer decode, MEMRA_FUSE_ADD_NORM=0 reverts): off the q8_1
            // fast path the residual add and this layer's attn_norm ran as two launches;
            // add_rms_norm does both (kernel_check identity: add_rms_norm == add then
            // rms_norm; same rms_block()), then the mixer takes the pre-normed h directly.
            static FUSE_AN: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
            let fuse_add_norm =
                *FUSE_AN.get_or_init(|| std::env::var("MEMRA_FUSE_ADD_NORM").as_deref() != Ok("0"));
            let mixed = match (taken, fuse) {
                (Some((x1, f1)), false) if fuse_add_norm => {
                    let mut x2 = e.uninit(n_embd)?;
                    let mut h = e.uninit(n_embd)?;
                    e.add_rms_norm(&x1, &f1, anorm, &mut x2, &mut h, n_embd, 1, eps)?;
                    x = x2;
                    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(mla) => {
                            self.mla_attn_cached(e, mla, &h, &pos_d, 1, il, cache)?
                        }
                        Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?,
                    }
                }
                (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_path_unimplemented("decode_step_h"),
                        Mixer::Kda(_) => {
                            crate::hybrid::kda_path_unimplemented("norm-fused decode_step_h")
                        }
                    }
                }
                (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)?
                }
            };

            lap(
                if il >= tail_from { &B_MIX_TAIL } else { &B_MIX },
                &mut lap_start,
            )?;

            // 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)?;
            // MEMRA_TG_PROBE_LAYER diagnostics (token-graph bisection): dump layer K's
            // attention output and post-FFN residual through the real eager path.
            if std::env::var("MEMRA_TG_PROBE_LAYER")
                .ok()
                .and_then(|v| v.parse::<usize>().ok())
                == Some(il)
            {
                use std::io::Write;
                let mut xp = e.uninit(n_embd)?;
                e.add(&x1, &ffn_out, &mut xp, n_embd)?;
                let (pm, px) = (e.dtoh(&mixed)?, e.dtoh(&xp)?);
                for (path, data) in [
                    ("/root/eager-probe-mixed.bin", &pm),
                    ("/root/eager-probe-x.bin", &px),
                ] {
                    let mut fo = std::fs::OpenOptions::new()
                        .create(true)
                        .append(true)
                        .open(path)?;
                    for v in data {
                        fo.write_all(&v.to_le_bytes())?;
                    }
                }
            }
            lap(
                if il >= tail_from { &B_FFN_TAIL } else { &B_FFN },
                &mut lap_start,
            )?;
            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())?;
            }
        }
        // MEMRA_HEAD_SPLIT=1 (step TP only): split the lm-head rows across both devices —
        // dev1 idles at the token tail, rows are independent, and the per-row program is the
        // same matvec_bf16 kernel, so the concatenated logits are BIT-IDENTICAL to the
        // single-device head. Falls through to the plain matmul when ineligible.
        let host = 'head: {
            let split_on = {
                static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
                *ON.get_or_init(|| std::env::var("MEMRA_HEAD_SPLIT").as_deref() == Ok("1"))
            };
            if split_on
                && self.uses_sliding_gated_moe_program()
                && let Some(host) = self.head_split_matvec(e, &hn)?
            {
                break 'head host;
            }
            let logits = e.matmul(&self.output, &hn, 1)?;
            e.dtoh(&logits)?
        };
        lap(&B_HEAD, &mut lap_start)?;
        if timing {
            use std::sync::atomic::Ordering;
            let tokens = B_TOKENS.fetch_add(1, Ordering::Relaxed) + 1;
            if tokens.is_multiple_of(10) {
                let per = |t: &std::sync::atomic::AtomicU64| {
                    t.load(Ordering::Relaxed) as f64 / tokens as f64 / 1.0e6
                };
                eprintln!(
                    "[decode-bucket-timing] tokens={tokens} ms/token mix={:.2} ffn={:.2} \
                     mix_tail={:.2} ffn_tail={:.2} head={:.2}",
                    per(&B_MIX),
                    per(&B_FFN),
                    per(&B_MIX_TAIL),
                    per(&B_FFN_TAIL),
                    per(&B_HEAD),
                );
            }
        }
        cache.pos += 1;
        Ok((host, h_seed))
    }

    /// ASYNC-AHEAD DEVICE-CHAINED greedy or sampled decode (MEMRA_ASYNC_CHAIN=K): run up to `k`
    /// tokens with NO host sync inside the chain — the tail argmax writes the resident
    /// token_d on-device (host-identical tie-break, argmax_gate receipt), the next
    /// iteration embeds straight from it (embed_gather_device, bit-identical rows), and
    /// the host reads the id history ring ONCE per chunk. Unlike the graph chunk this
    /// keeps EAGER kernels and streams (full stream concurrency); the host submit runs
    /// ahead of the GPU, so the per-token host wall overlaps device work instead of
    /// serializing after it.
    /// Contract mirrors step35_token_graph_chunk: consumes `token` (already emitted by
    /// the caller) as launch 0's input and returns (hist[0..k], last token's logits).
    /// The caller emits hist[..k-1]. On a greedy chain, hist[k-1] == argmax(logits); on a
    /// sampled chain it is the device-drawn boundary id and MUST be fed directly instead
    /// of re-derived greedily from the returned row. When `MEMRA_HEAD_SPLIT=1`, both the
    /// greedy argmax and sampled draw consume the split path's materialized concatenated
    /// logits on device; the host reads that persistent row only once at chunk end.
    /// SAMPLED chain (owner rule: "we dont serve greedy, for real benchmarking we use
    /// sampling"). `samp` carries the serving sampler; the draw happens ON DEVICE inside the
    /// chain — `filter_stats` -> `gumbel_perturb_filtered_col` -> `argmax` into the resident
    /// `token_d` — so a sampled stream keeps the chain's whole point, which is that no host
    /// sync happens between tokens. The per-step counter advances so each token draws its own
    /// Gumbel noise. Without `MEMRA_HEAD_SPLIT`, the same draw runs on the plain head row.
    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
    pub fn decode_step_chain(
        &self,
        e: &Engine,
        token: u32,
        k_target: usize,
        cache: &mut Cache,
        samp: Option<&crate::decode_batch::DevSamp>,
    ) -> Result<Option<(Vec<u32>, Vec<f32>)>, Box<dyn std::error::Error>> {
        self.refuse_hyper("decode_step_chain")?;
        cache.ensure_usable("decode_step_chain")?;
        if !self.uses_sliding_gated_moe_program() {
            return Ok(None);
        }
        let k = k_target.min(16);
        if k < 2 {
            return Ok(None);
        }
        let Some(embd_gpu) = self.embd_gpu_try(e) else {
            return Ok(None);
        };
        let cfg = &self.cfg;
        let n_embd = cfg.n_embd as usize;
        let n_vocab = cfg.n_vocab as usize;
        let eps = cfg.rms_eps;
        let n_layers = self.layers.len();
        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);

        // Resident chain state (token id, id history ring, ring index), one set per device.
        #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
        static CHAIN: std::sync::Mutex<
            Option<(usize, CudaSlice<u32>, CudaSlice<u32>, CudaSlice<i32>)>,
        > = std::sync::Mutex::new(None);
        let mut guard = CHAIN.lock().map_err(|_| "chain state lock is poisoned")?;
        if guard.as_ref().is_none_or(|(d, ..)| *d != e.ctx().ordinal()) {
            *guard = Some((
                e.ctx().ordinal(),
                e.stream().clone_htod(&[0u32])?,
                e.stream().clone_htod(&[0u32; 16])?,
                e.htod_i32(&[0])?,
            ));
        }
        let (_, token_d, hist, hist_idx) = guard.as_mut().expect("armed above");

        // O-PROJ TAIL deferral eligibility (see decode_step_h).
        let _oproj_tail_scope = crate::tp::oproj_tail_scope();
        // RANK0 STREAM MERGE (see decode_step_h).
        let _r0merge = if crate::tp::rank0_merge_on() {
            Some(memra_runtime::rank0_redirect_scope(
                e.ctx().ordinal(),
                e.gpu.main_stream().clone(),
                e.gpu.blas(),
            ))
        } else {
            None
        };
        // Per-token pos buffers staged BEFORE the chain (the only H2D the chain needs).
        let mut pos_bufs = Vec::with_capacity(k);
        for step in 0..k {
            pos_bufs.push(e.htod_i32(&[(cache.pos + step) as i32])?);
        }
        e.set_u32_one(token_d, token)?;
        e.set_i32_one(hist_idx, 0)?;

        // MEMRA_CHAIN_PHASE=1 (P0 CEILING PROBE — WRONG OUTPUT BY DESIGN): alternate
        // tokens ride disjoint phase streams with NO cross-token event edges yet, so the
        // schedule shows the token-pipeline overlap ceiling while the ids race. Timing
        // receipts only; never gate a tape under this door.
        static PHASE_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        let phase_on =
            *PHASE_ON.get_or_init(|| std::env::var("MEMRA_CHAIN_PHASE").as_deref() == Ok("1"));

        let mut last_logits: Option<Option<CudaSlice<f32>>> = None;
        #[allow(clippy::needless_range_loop)]
        // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
        for step in 0..k {
            let _phase_ov = if phase_on {
                let (ps, pb) = e.gpu.phase_pair(step & 1)?;
                memra_runtime::set_decode_phase(Some(step & 1));
                Some(memra_runtime::push_stream_override(ps, pb))
            } else {
                None
            };
            let pos = cache.pos;
            let step_r = (|| -> Result<Option<CudaSlice<f32>>, Box<dyn std::error::Error>> {
                let x = e.embed_gather_device(embd_gpu, token_d, n_embd, embd_qt, embd_rb)?;
                let x = self.decode_layers_eager(e, x, 0, n_layers, &pos_bufs[step], pos, cache)?;
                let mut hn = e.uninit(n_embd)?;
                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
                // Split head when armed (MEMRA_HEAD_SPLIT env + eligibility): identical
                // concatenated logits, device argmax, no per-token readback.
                static HS_ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
                let hs =
                    *HS_ON.get_or_init(|| std::env::var("MEMRA_HEAD_SPLIT").as_deref() == Ok("1"));
                let sampling = samp.filter(|s| s.temp > 0.0);
                let split_done = if hs && self.uses_sliding_gated_moe_program() {
                    match sampling {
                        // Sampling keeps HEAD_SPLIT: the split path materializes the full
                        // concatenated row, so the device draw reads it instead of an argmax.
                        Some(s) => self.head_split_sample_device(
                            e,
                            &hn,
                            token_d,
                            s,
                            s.ctr.wrapping_add(step as u32),
                        )?,
                        None => self.head_split_argmax_device(e, &hn, token_d)?,
                    }
                } else {
                    false
                };
                let logits = if split_done {
                    None
                } else {
                    let logits = e.matmul(&self.output, &hn, 1)?;
                    match samp.filter(|s| s.temp > 0.0) {
                        None => e.argmax_token_device_into(&logits, token_d, n_vocab)?,
                        Some(s) => {
                            // Device draw, no host sync: thresholds for this row, Gumbel
                            // perturbation of the filtered row, argmax into token_d. Same
                            // kernels and the same (seed, ctr) draw the serve tick uses.
                            let ctr = s.ctr.wrapping_add(step as u32);
                            // Persistent per-chain scratch: allocating these per token cost
                            // more than the split head saved when this arm was first measured.
                            let filtered = s.top_k > 0 || s.top_p < 1.0 || s.min_p > 0.0;
                            if filtered {
                                let rows_d = e.htod_i32(&[0i32])?;
                                let mut th = e.zeros(1)?;
                                let mut z = e.zeros(1)?;
                                let mut mx = e.zeros(1)?;
                                e.filter_stats(
                                    &logits, n_vocab, &rows_d, &mut th, &mut z, &mut mx, n_vocab,
                                    1, s.temp, s.top_k, s.top_p, s.min_p,
                                )?;
                                let mut pb = e.zeros(n_vocab)?;
                                e.gumbel_perturb_filtered_col(
                                    &logits, 0, &mut pb, n_vocab, s.seed, ctr, s.temp, &mx, &th, 0,
                                )?;
                                e.argmax_token_device_col(&pb, 0, n_vocab, token_d, 0)?;
                            } else {
                                let mut pb = e.zeros(n_vocab)?;
                                e.gumbel_perturb_col(
                                    &logits, 0, &mut pb, n_vocab, s.seed, ctr, s.temp,
                                )?;
                                e.argmax_token_device_col(&pb, 0, n_vocab, token_d, 0)?;
                            }
                        }
                    }
                    Some(logits)
                };
                e.u32_hist_append(token_d, hist, hist_idx)?;
                Ok(logits)
            })();
            if phase_on {
                memra_runtime::set_decode_phase(None);
            }
            let logits = step_r?;
            cache.pos += 1;
            last_logits = Some(logits);
            // (None = split-head path; the persistent row holds this token's logits.)
        }
        if phase_on {
            // Drain both phases on every engine before the host readback.
            for p in 0..2 {
                e.gpu.phase_pair(p)?.0.synchronize()?;
            }
            if let Some(tp) = self.layers.first().and_then(|l| match &l.mixer {
                Mixer::Full(fa) => fa.step_tp_qkv.as_ref(),
                _ => None,
            }) {
                for rank in 0..tp.runtime.devices().len() {
                    if let Some(engine) = tp.runtime.rank_engine(rank) {
                        let _main = engine.gpu.enter_main()?;
                        for p in 0..2 {
                            engine.gpu.phase_pair(p)?.0.synchronize()?;
                        }
                    }
                }
            }
        }
        let hist_h = e.dtoh_u32(hist)?;
        let logits_h = match last_logits.expect("k >= 2") {
            Some(row) => e.dtoh(&row)?,
            None => self.head_split_logits_dtoh(e)?,
        };
        Ok(Some((hist_h[..k].to_vec(), logits_h)))
    }

    /// 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).
    /// `pub(crate)`: also the B=1 serve fast-path's trunk (decode_batch.rs
    /// `decode_step_b1_fast`, H3) — shared verbatim so the serve path inherits every m=1
    /// fusion instead of needing a batched twin per lever.
    #[allow(clippy::too_many_arguments)]
    pub(crate) 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();
            // FUSION #2f (same door as decode_step_h): off the q8_1 fast path, fuse the
            // residual add with this layer's attn_norm via add_rms_norm.
            static FUSE_AN_LE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
            let fuse_add_norm = *FUSE_AN_LE
                .get_or_init(|| std::env::var("MEMRA_FUSE_ADD_NORM").as_deref() != Ok("0"));
            let mixed = match (taken, fuse) {
                (Some((x1, f1)), false) if fuse_add_norm => {
                    let mut x2 = e.uninit(n_embd)?;
                    let mut h = e.uninit(n_embd)?;
                    e.add_rms_norm(&x1, &f1, anorm, &mut x2, &mut h, n_embd, 1, eps)?;
                    x = x2;
                    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(mla) => self.mla_attn_cached(e, mla, &h, pos_d, 1, il, cache)?,
                        Mixer::Kda(la) => crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?,
                    }
                }
                (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_path_unimplemented("decode_step_chain"),
                        Mixer::Kda(_) => {
                            crate::hybrid::kda_path_unimplemented("norm-fused decode_layers_eager")
                        }
                    }
                }
                (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)?;
            // MEMRA_TG_PROBE_LAYER diagnostics (token-graph bisection): dump layer K's
            // attention output and post-FFN residual through the real eager path.
            if std::env::var("MEMRA_TG_PROBE_LAYER")
                .ok()
                .and_then(|v| v.parse::<usize>().ok())
                == Some(il)
            {
                use std::io::Write;
                let mut xp = e.uninit(n_embd)?;
                e.add(&x1, &ffn_out, &mut xp, n_embd)?;
                let (pm, px) = (e.dtoh(&mixed)?, e.dtoh(&xp)?);
                for (path, data) in [
                    ("/root/eager-probe-mixed.bin", &pm),
                    ("/root/eager-probe-x.bin", &px),
                ] {
                    let mut fo = std::fs::OpenOptions::new()
                        .create(true)
                        .append(true)
                        .open(path)?;
                    for v in data {
                        fo.write_all(&v.to_le_bytes())?;
                    }
                }
            }
            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()
        );
        // #87 REVERSE PUBLICATION (lane/pp2spec-crash): this body's stage-stream
        // allocations may reuse pool blocks freed from a PREVIOUS ppn call's outputs
        // (h_seed, verify vx/ckpt) whose primary-stream consumers are still queued —
        // the reuse-write races the queued read. Order every stage stream behind the
        // caller's stream before the first stage allocation. Full anatomy:
        // `PpNRt::fence_stages_behind`.
        rt.fence_stages_behind(&e.stream())?;
        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>> {
        self.refuse_hyper("decode_step_h_ppn_deferred")?;
        cache.ensure_usable("decode_step_h_ppn_deferred")?;
        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.uses_gemma_program() {
            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 walk = rt.acquire_deferred_walk("decode_step_h_ppn_deferred")?;
        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(),
            walk,
        ))
    }

    /// 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>> {
        self.refuse_hyper("decode_step_lockstep")?;
        for cache in caches.iter() {
            cache.ensure_usable("decode_step_lockstep")?;
        }
        if tokens.len() != caches.len() || tokens.is_empty() {
            return Err("lockstep needs one token per stream cache".into());
        }
        if self.uses_gemma_program() {
            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_path_unimplemented("lockstep decode")
                                }
                                Mixer::Kda(_) => {
                                    crate::hybrid::kda_path_unimplemented("lockstep decode")
                                }
                            }
                        }
                        (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 && 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`.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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>> {
        self.refuse_hyper("decode_step_dc")?;
        cache.ensure_usable("decode_step_dc")?;
        // 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());
        }
        // PP DOOR: fail closed (pp2-hardening 2026-08-06). Same hole the batched path had —
        // the dc walk below is `for (il, layer) in self.layers.iter().enumerate()` on one
        // stream, with no stage split, so a sharded cross-device placement would peer-read
        // every remote layer's weights per step. Sits BEFORE the gemma4 delegate because
        // that twin has the same unsplit shape. The graph-capture path (`decode_step_dc_cap*`)
        // is covered transitively: it captures this same kernel chain, and its drivers reach
        // dc first — but a future capture path that does NOT is why the guard is a shared
        // helper (`pp::refuse_unsplit_if_remote`) rather than four copies.
        crate::pp::refuse_unsplit_if_remote(
            "decode_step_dc",
            "use the eager pp arm (decode_step_h), which IS stage-split",
        )?;
        if self.uses_gemma_program() {
            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)?;
            // MEMRA_TG_PROBE_LAYER diagnostics (token-graph bisection): dump layer K's
            // attention output and post-FFN residual through the real eager path.
            if std::env::var("MEMRA_TG_PROBE_LAYER")
                .ok()
                .and_then(|v| v.parse::<usize>().ok())
                == Some(il)
            {
                use std::io::Write;
                let mut xp = e.uninit(n_embd)?;
                e.add(&x1, &ffn_out, &mut xp, n_embd)?;
                let (pm, px) = (e.dtoh(&mixed)?, e.dtoh(&xp)?);
                for (path, data) in [
                    ("/root/eager-probe-mixed.bin", &pm),
                    ("/root/eager-probe-x.bin", &px),
                ] {
                    let mut fo = std::fs::OpenOptions::new()
                        .create(true)
                        .append(true)
                        .open(path)?;
                    for v in data {
                        fo.write_all(&v.to_le_bytes())?;
                    }
                }
            }
            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).
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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.refuse_hyper("decode_step_dc_cap")?;
        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>> {
        self.refuse_hyper("decode_step_dc_cap_masked")?;
        cache.ensure_usable("decode_step_dc_cap")?;
        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)?;
            // MEMRA_TG_PROBE_LAYER diagnostics (token-graph bisection): dump layer K's
            // attention output and post-FFN residual through the real eager path.
            if std::env::var("MEMRA_TG_PROBE_LAYER")
                .ok()
                .and_then(|v| v.parse::<usize>().ok())
                == Some(il)
            {
                use std::io::Write;
                let mut xp = e.uninit(n_embd)?;
                e.add(&x1, &ffn_out, &mut xp, n_embd)?;
                let (pm, px) = (e.dtoh(&mixed)?, e.dtoh(&xp)?);
                for (path, data) in [
                    ("/root/eager-probe-mixed.bin", &pm),
                    ("/root/eager-probe-x.bin", &px),
                ] {
                    let mut fo = std::fs::OpenOptions::new()
                        .create(true)
                        .append(true)
                        .open(path)?;
                    for v in data {
                        fo.write_all(&v.to_le_bytes())?;
                    }
                }
            }
            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>> {
        self.refuse_hyper("generate_graph")?;
        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeGraph) {
            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
                return Err("neither graph nor eager decode rewrite is qualified".into());
            }
            static ONCE: std::sync::Once = std::sync::Once::new();
            ONCE.call_once(|| {
                eprintln!(
                    "[rewrite] decode-graph.v1 unqualified; using receipt-backed native eager decode"
                );
            });
            return self.generate(e, prompt, max_new);
        }
        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
    }

    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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.uses_gemma_program() {
            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)?;

            #[allow(clippy::int_plus_one)]
            // allow: the +1 form states the documented boundary, not an off-by-one
            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>> {
        self.refuse_hyper("graph_session_new")?;
        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>> {
        cache.ensure_usable("graph_session_from_cache")?;
        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))
    }

    /// Measurement door for `graph_session_recapture` (graph-allocfree-probe): the capture
    /// path timed WITHOUT the prompt prime. Same call the live step() makes at a
    /// kernel-class crossing.
    pub fn graph_session_recapture_pub(
        &self,
        e: &Engine,
        sess: &mut GraphSession,
    ) -> Result<(), Box<dyn std::error::Error>> {
        self.graph_session_recapture(e, sess)
    }

    /// 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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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).
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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))
    }

    #[allow(clippy::too_many_arguments)]
    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
    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>> {
        // step35 has no device-counter twin yet: the `_dc` family needs a windowed dc fa_decode
        // (SWA layers read a token-OFFSET view, which the dc kernels' len_d-derived t_kv cannot
        // express) plus a per-layer-n_head capture. Refuse loudly instead of silently running
        // the generic geometry. The eager arm (`step35_decode_attn`) is the supported decode.
        if self.uses_sliding_gated_moe_program() {
            return Err(
                "step35 has no device-counter/graph decode arm (SWA needs an offset KV \
                        view the dc kernels cannot express) — use the eager decode"
                    .into(),
            );
        }
        let cfg = &self.cfg;
        let geometry = cfg.full_attention_geometry_at(il as u32);
        let n_head = geometry.n_head as usize;
        let n_head_kv = geometry.n_head_kv as usize;
        let head_dim = geometry.head_dim_k as usize;
        let eps = cfg.rms_eps;
        let scale = geometry.attention_scale();

        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 let Some(mut qkv) = self.full_attn_tp_qkv(e, fa, h, 1)? {
            let v = qkv.pop().ok_or("full-attention TP QKV omitted V")?;
            let k = qkv.pop().ok_or("full-attention TP QKV omitted K")?;
            let q = qkv.pop().ok_or("full-attention TP QKV omitted Q")?;
            if !qkv.is_empty() {
                return Err("full-attention TP QKV returned extra projections".into());
            }
            (q, k, v)
        } else 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 = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
        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 = geometry.n_rot 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,
            geometry.rope_base,
            1.0,
        )?;
        e.rope_neox(
            &mut k,
            pos_d,
            head_dim,
            rope_dims,
            n_head_kv,
            1,
            geometry.rope_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,
        };
        match self.full_attn_tp_o(e, fa, &attn_g, 1)? {
            Some(output) => Ok(output),
            None => 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, 0)?;
            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.uses_gemma_program()
            && 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.)
        // step35 is EXCLUDED: this route calls `decode_step_dc`, whose full-attn arm refuses
        // step35 by design (SWA layers need a token-OFFSET KV view the dc kernels' len_d-derived
        // t_kv cannot express). Without this gate the door opens for any greedy model and the
        // refusal surfaces as a user-visible generate() error — the first PP-2 boot of
        // Step-3.7-Flash died exactly there, AFTER a clean load and an argmax MATCH.
        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
            && !self.uses_sliding_gated_moe_program()
            && 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, 0)?;
            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,
        );
        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.uses_gemma_program()
            && 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.)
        // step35 is EXCLUDED here for the same reason as the `generate` mirror above: every route
        // inside this door (`decode_step_dc` and the `graph_decode_loop` capture) reaches
        // `full_attn_decode_dc_inner`, which refuses step35 because its SWA layers read a
        // token-OFFSET KV view the dc kernels cannot express. step35 takes the host-logits eager
        // loop at the bottom of this function (`decode_step` -> `step35_decode_attn`), which is
        // the supported decode for this arch. Removing this gate requires a windowed dc fa_decode
        // plus a per-layer-n_head capture, not a flag.
        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
            && !self.uses_sliding_gated_moe_program()
            && 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.
            //
            // KEY LOWERED 256 -> 48 (q27 deep dive, 2026-08-05, pro6000wk-runpod-community).
            // The 256 key was set by the E4B amortization rule, never by a measured crossover,
            // so every <=128-token generation — including the whole published board, which runs
            // --max-tokens 128 — was silently EAGER. Swept the actual crossover on TWO models
            // (the key is a cross-model default, so one artifact is not enough), interleaved
            // arms with the order alternated per rep, N=3, all runs argmax MATCH:
            //   Qwen3.6-27B-Q8_0     : n=16 -7.47% | n=32 -1.35% | n=48 +0.90% | n=64 +1.93%
            //                          n=128 +3.80% | n=512 +5.50%
            //   Qwen3.6-27B-NVFP4-MTP: n=16 -15.27% | n=32 +0.22% | n=48 +3.45%
            //                          n=64 +5.09% | n=128 +7.72%
            // Both models: clearly negative at 16, no reliable gain at 32, positive from 48 up,
            // monotone in budget from 48 on. 48 is the first budget where BOTH are positive, so
            // it is the key — the capture cost needs ~32 steps to amortize, not ~256. The n=32
            // nvfp4 cell is NOISY, not flat (graph arm 79.02/78.91/77.09, spread 1.93 vs an
            // eager spread of 0.04): it is not evidence of a win, and it is why the key sits at
            // 48 rather than 32. Exactness at the new key:
            // graph-decode-gate 256 steps BIT-IDENTICAL (buckets=16, captures=2),
            // graph-session-gate 96 tokens PASS, kernel-check ALL GREEN, run-spec K=1..8
            // self-consistency PASS. Board caveat: community board, RELATIVE deltas only.
            //
            // SM-GATED (5090-arbiter gate, 2026-08-05, research/q27-deepdive-20260805/local5090/):
            // the 48 key does NOT transfer to the 82-SM local rig. Same A/B protocol there
            // (tg128 d512, N=3 interleaved, order alternated, warmup discarded): q27-NVFP4-MTP
            // graph arm at n=128 = -1.61% (eager 45.86 / graph 45.12 median, 3/3 pairs lose),
            // and the crossover sweep stays negative through n=256 (-1.07%) and n=512 (-0.59%)
            // — on few-SM silicon the replay's fixed kernel forms lag the tuned eager lanes and
            // the launch-gap tax the graph amortizes is proportionally smaller. Key on SM count
            // (the fa_split_keys big_rig pattern, lib.rs fa_sm_count), threshold 180: the 48
            // crossover is MEASURED only at 188 SM (PRO 6000) and refuted at 82 SM; the 132-SM
            // H100 board and the 170-SM desktop 5090 are UNMEASURED at sub-256 budgets, so they
            // keep the shipped 256 key their board rows were measured with (rig-divergence +
            // stale-verdict laws). Widening the gate below 180 requires an on-box crossover
            // sweep on that silicon, not an inference from this comment.
            let big_rig = e.sm_count() >= 180;
            let gen_graph = match std::env::var("MEMRA_GEN_GRAPH").as_deref() {
                Ok("1") => true,
                Ok("0") => false,
                _ => budget >= if big_rig { 48 } else { 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.
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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.
    #[allow(clippy::too_many_arguments)]
    // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
    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>> {
        if self.uses_sliding_gated_moe_program() {
            return self.step35_decode_attn(e, fa, il, h, pre_q, pos_d, cache);
        }
        let cfg = &self.cfg;
        let geometry = cfg.full_attention_geometry_at(il as u32);
        if fa
            .step_tp_qkv
            .as_ref()
            .is_some_and(|tp| tp.attention.is_some())
        {
            if pre_q.is_some() {
                return Err(
                    "rank-local generic TP attention preserves BF16 activations and refuses the \
                     q8_1 pre-quantized decode path"
                        .into(),
                );
            }
            if geometry.attention_gate != memra_gguf::config::AttentionGateKind::None {
                return Err(
                    "rank-local generic TP attention currently requires an ungated attention \
                     plan; fused-Q and separate-head gates retain their existing qualified paths"
                        .into(),
                );
            }
            if !crate::tp::step_tp_decode_v2_enabled()? {
                return Err(
                    "MEMRA_PARALLEL_TP_ATTENTION=1 requires MEMRA_STEP_TP_DECODE_V2=1 for \
                     generic decode; the v1 driver is Step-gate-specific"
                        .into(),
                );
            }
            return self.step35_tp_decode_attn_resident_v2(e, fa, il, h, pos_d, cache);
        }
        let n_head = geometry.n_head as usize;
        let n_head_kv = geometry.n_head_kv as usize;
        let head_dim = geometry.head_dim_k as usize;
        let eps = cfg.rms_eps;
        let scale = geometry.attention_scale();

        // 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 = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
        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 = geometry.n_rot as usize;
        e.rope_neox(
            &mut q,
            pos_d,
            head_dim,
            rope_dims,
            n_head,
            1,
            geometry.rope_base,
            1.0,
        )?;
        e.rope_neox(
            &mut k,
            pos_d,
            head_dim,
            rope_dims,
            n_head_kv,
            1,
            geometry.rope_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,
        };
        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>> {
        if self.uses_sliding_gated_moe_program() {
            return Err(
                "step35 has no batched (m-stream) decode mixer — per-layer n_head, \
                        partial rope and the SWA offset view need a step35 twin"
                    .into(),
            );
        }
        let cfg = &self.cfg;
        let geometry = cfg.full_attention_geometry_at(il as u32);
        let n_head = geometry.n_head as usize;
        let n_head_kv = geometry.n_head_kv as usize;
        let head_dim = geometry.head_dim_k as usize;
        let n_embd = cfg.n_embd as usize;
        let eps = cfg.rms_eps;
        let scale = geometry.attention_scale();
        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 = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
        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 = geometry.n_rot as usize;
        e.rope_neox(
            &mut q,
            pos_cat,
            head_dim,
            rope_dims,
            n_head,
            m,
            geometry.rope_base,
            1.0,
        )?;
        e.rope_neox(
            &mut k,
            pos_cat,
            head_dim,
            rope_dims,
            n_head_kv,
            m,
            geometry.rope_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).
    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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)
    }

    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
    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 geometry = la.geometry;
        let d_state = geometry.key_head_dim as usize;
        let num_k = geometry.key_heads as usize;
        let num_v = geometry.value_heads as usize;
        let d_conv = geometry.conv_kernel as usize;
        let head_k = d_state;
        let key_dim = head_k * num_k;
        let value_dim = geometry.value_head_dim as usize * 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 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,
        )?;
        e.matmul(&la.ssm_out, &gn, 1)
    }
}