memra_engine/

gemma_spec.rs

//! gemma4 MTP spec-decode: the "gemma4-assistant" drafter (4-layer, Q-only attention over the
//! MAIN model's KV cache — no draft KV, no trims) + the greedy draft/verify loop.
//!
//! Wiring verified from llama gemma4-assistant.cpp + llama-model.cpp:2162 (HANDOVER "GEMMA4 MTP
//! DRAFTER — VERIFIED WIRING"): per draft token, x = MAIN tok_embd(token) * sqrt(2816);
//! xh = concat(x, h[2816]) -> pre_proj [5632->1024]; 4 gemma-style blocks whose attention
//! projects Q ONLY and attends the main cache (SWA layers 0..2 -> main layer n-2 = 28 windowed;
//! global layer 3 -> main layer n-1 = 29 full); dense GELU_PAR ffn; final output_norm ->
//! TIED 1024-dim head (no softcap); h_next = post_proj [1024->2816].

use crate::Engine;
use crate::cache::Cache;
use crate::hybrid::HybridModel;
use crate::model::GpuTensor;
use cudarc::driver::CudaSlice;
use memra_gguf::GgufFile;
use memra_gguf::source::{GgufSource, TensorSource};

pub struct GemmaDraftLayer {
    pub attn_norm: GpuTensor,
    pub wq: GpuTensor,
    pub wo: GpuTensor,
    pub q_norm: GpuTensor,
    pub post_attn_norm: GpuTensor,
    pub ffn_norm: GpuTensor,
    pub ffn_gate: GpuTensor,
    pub ffn_up: GpuTensor,
    pub ffn_down: GpuTensor,
    pub ffn_post_norm: GpuTensor,
    pub out_scale: f32,
    pub swa: bool,
    pub hd: usize,
    pub nh: usize,
}

pub struct GemmaDraft {
    pub layers: Vec<GemmaDraftLayer>,
    pub pre_proj: GpuTensor,  // [5632 -> 1024]
    pub post_proj: GpuTensor, // [1024 -> 2816]
    pub output_norm: GpuTensor,
    pub head: GpuTensor, // tied drafter token_embd [1024, n_vocab] (or FR-trimmed rows)
    /// FR-Spec trim map: draft-row index -> target token id (None = full head, identity).
    pub d2t: Option<Vec<u32>>,
    /// Device copy of `d2t` — the async round translates each drafted trim-idx in place
    /// (u32_map_k) before it seeds the next draft step or meets the verify argmax.
    pub d2t_dev: Option<CudaSlice<u32>>,
    /// Adaptive trim (coverage escapes are the entire trim cost — oracle-proven +2% on the
    /// cell the static trim lost by 17%, jsonl 2026-07-19): spare head slots learned at
    /// serve time from the prompt's own ids and verify-correction tokens.
    pub trim_adapt: Option<TrimAdapt>,
    pub rope_freqs: CudaSlice<f32>,
    pub ones: CudaSlice<f32>, // weightless-norm weight (max hd 512)
    pub n_embd: usize,        // 1024
    pub n_backbone: usize,    // 2816
    pub rope_base_global: f32,
    pub rope_base_swa: f32,
    pub sliding_window: usize,
}

/// Serve-time adaptive trim (MEMRA_GEMMA_TRIM_ADAPT=<spare slots>): the static FR trim's whole
/// loss is coverage escapes — tokens the base emits that the trim can't propose (guaranteed
/// rejections; the oracle control that injected the exact escapees flipped a -17% cell to +2%
/// at identical acceptance, jsonl 2026-07-19). Every escape self-identifies at serve time: it
/// arrives as a verify CORRECTION token (and its cousins ride in with the prompt), so the head
/// keeps `n_spare` extra rows and learns them — prompt ids up front, corrections as they land.
/// First miss pays one rejected round; every recurrence after is proposable. Rows are written
/// into the existing device buffers (no realloc — captured graphs keep their baked addresses).
pub struct TrimAdapt {
    /// full-vocab head rows (host copy) — the gather source for learned rows.
    src_rows: Vec<u8>,
    row_bytes: usize,
    n_vocab: usize,
    /// trim-set membership by token id (ranked + learned).
    present: Vec<bool>,
    /// spare slots live at [spare_base, spare_base + n_spare) in the gathered head.
    spare_base: usize,
    n_spare: usize,
    used: usize,
    logged_full: bool,
}

impl TrimAdapt {
    /// Add `tok`'s head row to the trim set if absent and a spare slot is free.
    fn maybe_add(
        &mut self,
        e: &Engine,
        tok: u32,
        head: &mut GpuTensor,
        d2t: &mut [u32],
        d2t_dev: &mut CudaSlice<u32>,
    ) -> Result<bool, Box<dyn std::error::Error>> {
        let t = tok as usize;
        if t >= self.n_vocab || self.present[t] {
            return Ok(false);
        }
        if self.used == self.n_spare {
            if !self.logged_full {
                self.logged_full = true;
                eprintln!(
                    "[trim-adapt] spare slots exhausted ({}) — later escapes stay unproposable",
                    self.n_spare
                );
            }
            return Ok(false);
        }
        let slot = self.spare_base + self.used;
        self.used += 1;
        self.present[t] = true;
        if let GpuTensor::Quant { bytes, .. } = head {
            e.htod_u8_into(
                bytes,
                slot * self.row_bytes,
                &self.src_rows[t * self.row_bytes..(t + 1) * self.row_bytes],
            )?;
        }
        d2t[slot] = tok;
        e.u32_set_k(d2t_dev, tok, slot)?;
        Ok(true)
    }
}

/// Union `toks` into the adaptive trim set (no-op when the draft has no adaptive state).
/// Split-borrow helper: the fields move together or not at all.
fn trim_adapt_learn(
    e: &Engine,
    d: &mut GemmaDraft,
    toks: &[u32],
) -> Result<(), Box<dyn std::error::Error>> {
    let GemmaDraft {
        trim_adapt,
        head,
        d2t,
        d2t_dev,
        ..
    } = d;
    let (Some(ta), Some(d2t), Some(d2t_dev)) =
        (trim_adapt.as_mut(), d2t.as_mut(), d2t_dev.as_mut())
    else {
        return Ok(());
    };
    for &tok in toks {
        ta.maybe_add(e, tok, head, d2t, d2t_dev)?;
    }
    Ok(())
}

impl GemmaDraft {
    /// Adaptive-trim stats: (slots used, slot budget). None when adaptation is off.
    pub fn trim_adapt_stats(&self) -> Option<(usize, usize)> {
        self.trim_adapt.as_ref().map(|ta| (ta.used, ta.n_spare))
    }

    /// Persist the learned trim rows: append ids not yet in the sidecar to
    /// `<ranks>.learned` (the load path pre-fills spare slots from it, so a distribution's
    /// escapes pay their first-miss round ONCE across the serve lifetime, not per request).
    pub fn trim_adapt_save(&self) -> std::io::Result<usize> {
        let (Some(ta), Some(d2t), Some(path)) = (
            self.trim_adapt.as_ref(),
            self.d2t.as_ref(),
            self.trim_learned_path(),
        ) else {
            return Ok(0);
        };
        let prior: std::collections::HashSet<u32> = std::fs::read_to_string(&path)
            .map(|t| t.lines().filter_map(|l| l.trim().parse().ok()).collect())
            .unwrap_or_default();
        let fresh: Vec<u32> = d2t[ta.spare_base..ta.spare_base + ta.used]
            .iter()
            .copied()
            .filter(|id| !prior.contains(id))
            .collect();
        if !fresh.is_empty() {
            use std::io::Write;
            let mut f = std::fs::OpenOptions::new()
                .create(true)
                .append(true)
                .open(&path)?;
            for id in &fresh {
                writeln!(f, "{id}")?;
            }
        }
        Ok(fresh.len())
    }

    fn trim_learned_path(&self) -> Option<String> {
        std::env::var("MEMRA_GEMMA_DRAFT_RANKS")
            .ok()
            .map(|p| format!("{p}.learned"))
    }
}

fn load_t(
    e: &Engine,
    src: &dyn TensorSource,
    name: &str,
) -> Result<GpuTensor, Box<dyn std::error::Error>> {
    GpuTensor::load_from_source(e, src, name)
}

impl GemmaDraft {
    pub fn load(e: &Engine, g: &GgufFile) -> Result<Self, Box<dyn std::error::Error>> {
        // two published spellings of the same arch: the 26B/31B drafters ship
        // "gemma4-assistant", the E4B assistant ships "gemma4_assistant" — the metadata
        // key prefix follows the arch string verbatim.
        let arch = match g.arch() {
            Some(a @ ("gemma4-assistant" | "gemma4_assistant")) => a.to_string(),
            other => panic!("not a gemma4-assistant drafter (arch {other:?})"),
        };
        let src = GgufSource(g);
        let meta_u = |k: &str| -> u32 {
            g.metadata
                .get(&format!("{arch}.{k}"))
                .and_then(|v| v.as_u64())
                .unwrap_or(0) as u32
        };
        let meta_f = |k: &str, d: f32| -> f32 {
            match g.metadata.get(&format!("{arch}.{k}")) {
                Some(memra_gguf::MetaValue::F32(v)) => *v,
                Some(memra_gguf::MetaValue::F64(v)) => *v as f32,
                _ => d,
            }
        };
        let n_layer = meta_u("block_count") as usize;
        let n_embd = meta_u("embedding_length") as usize;
        // 26B/31B carry the target width as embedding_length_out; the E4B assistant as
        // n_embd_backbone.
        let n_backbone = match meta_u("embedding_length_out") as usize {
            0 => meta_u("n_embd_backbone") as usize,
            v => v,
        };
        let hd_g = meta_u("attention.key_length") as usize;
        let hd_s = meta_u("attention.key_length_swa") as usize;
        let swa_pat: Vec<bool> = match g
            .metadata
            .get(&format!("{arch}.attention.sliding_window_pattern"))
        {
            Some(memra_gguf::MetaValue::Array(a)) => a
                .iter()
                .filter_map(|v| v.as_u64().map(|x| x != 0))
                .collect(),
            _ => return Err("drafter missing sliding_window_pattern".into()),
        };

        let mut layers = Vec::with_capacity(n_layer);
        #[allow(clippy::needless_range_loop)]
        // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
        for il in 0..n_layer {
            let p = |n: &str| format!("blk.{il}.{n}");
            let swa = swa_pat[il];
            let out_scale = {
                let t = src
                    .find(&p("layer_output_scale.weight"))
                    .ok_or("missing layer_output_scale")?;
                memra_gguf::dequant::dequantize(t.ggml_type, &t.bytes, 1)[0]
            };
            let hd = if swa { hd_s } else { hd_g };
            let wq = load_t(e, &src, &p("attn_q.weight"))?;
            // heads per layer from the projection shape (the E4B assistant keeps 4 heads on
            // BOTH classes — hd differs — while 26B/31B are uniform; the shape is the truth).
            let nh = wq.out_features() / hd;
            layers.push(GemmaDraftLayer {
                attn_norm: load_t(e, &src, &p("attn_norm.weight"))?,
                wq,
                wo: load_t(e, &src, &p("attn_output.weight"))?,
                q_norm: load_t(e, &src, &p("attn_q_norm.weight"))?,
                post_attn_norm: load_t(e, &src, &p("post_attention_norm.weight"))?,
                ffn_norm: load_t(e, &src, &p("ffn_norm.weight"))?,
                ffn_gate: load_t(e, &src, &p("ffn_gate.weight"))?,
                ffn_up: load_t(e, &src, &p("ffn_up.weight"))?,
                ffn_down: load_t(e, &src, &p("ffn_down.weight"))?,
                ffn_post_norm: load_t(e, &src, &p("post_ffw_norm.weight"))?,
                out_scale,
                swa,
                hd,
                nh,
            });
        }
        let rope_freqs = {
            let t = src
                .find("rope_freqs.weight")
                .ok_or("drafter missing rope_freqs")?;
            e.htod(&memra_gguf::dequant::dequantize(
                t.ggml_type,
                &t.bytes,
                t.ne.iter().product::<u64>() as usize,
            ))?
        };
        // FR-Spec head trim (MEMRA_GEMMA_DRAFT_RANKS=<ids file, rank order>): gather the ranked
        // rows of the drafter head + d2t map. (Top-N-IDS truncation measured NEGATIVE — id
        // order is not frequency; the CORPUS-ranked gather is the real FR-Spec.)
        // MEMRA_GEMMA_TRIM_ADAPT=<n> (default 512 when ranks are set, 0 = off) appends n spare
        // rows the serve loop fills from prompt ids + verify corrections (see TrimAdapt).
        let (head, d2t, trim_adapt) = {
            let t = src
                .find("token_embd.weight")
                .ok_or("drafter missing token_embd")?;
            let in_f = t.ne[0] as usize;
            let n_vocab = t.ne[1] as usize;
            match std::env::var("MEMRA_GEMMA_DRAFT_RANKS").ok() {
                Some(path) => {
                    // row gather is layout-agnostic given the per-row byte stride: Q4_0 (26B
                    // drafter) and Q8_0 (31B drafter) both ship 32-elem blocks row-major.
                    // (qtype, elems/block, bytes/block) — the gather is stride-agnostic.
                    let (qtype, blk_e, blk_b) = match t.ggml_type {
                        memra_gguf::GgmlType::Q4_0 => (crate::QT_Q4_0, 32, 18),
                        memra_gguf::GgmlType::Q8_0 => (crate::QT_Q8_0, 32, 34),
                        memra_gguf::GgmlType::Q6_K => (crate::QT_Q6_K, 256, 210),
                        other => panic!("drafter head trim: unsupported head type {other:?}"),
                    };
                    let ids: Vec<u32> = std::fs::read_to_string(&path)?
                        .lines()
                        .filter_map(|l| l.trim().parse().ok())
                        .filter(|&id| (id as usize) < n_vocab)
                        .collect();
                    let n_spare: usize = std::env::var("MEMRA_GEMMA_TRIM_ADAPT")
                        .ok()
                        .and_then(|v| v.parse().ok())
                        .unwrap_or(512);
                    let row_bytes = in_f / blk_e * blk_b;
                    let mut gathered = Vec::with_capacity((ids.len() + n_spare) * row_bytes);
                    for &id in &ids {
                        let off = id as usize * row_bytes;
                        gathered.extend_from_slice(&t.bytes[off..off + row_bytes]);
                    }
                    // spare slots start as copies of row ids[0] mapping to ids[0] — a real,
                    // already-present token, so however the argmax resolves the duplicate-
                    // logit tie, the d2t translation lands on the same token id.
                    for _ in 0..n_spare {
                        let off = ids[0] as usize * row_bytes;
                        gathered.extend_from_slice(&t.bytes[off..off + row_bytes]);
                    }
                    eprintln!(
                        "[gemma-draft] FR head trim: {} rows + {} adaptive ({} MB vs {} MB full)",
                        ids.len(),
                        n_spare,
                        (ids.len() + n_spare) * row_bytes / 1_000_000,
                        n_vocab * row_bytes / 1_000_000
                    );
                    let mut trim_adapt = (n_spare > 0).then(|| {
                        let mut present = vec![false; n_vocab];
                        for &id in &ids {
                            present[id as usize] = true;
                        }
                        TrimAdapt {
                            src_rows: t.bytes.to_vec(),
                            row_bytes,
                            n_vocab,
                            present,
                            spare_base: ids.len(),
                            n_spare,
                            used: 0,
                            logged_full: false,
                        }
                    });
                    let mut d2t = ids;
                    let spare_fill = d2t[0];
                    d2t.extend(std::iter::repeat_n(spare_fill, n_spare));
                    // pre-fill spare slots from the learned sidecar (trim_adapt_save):
                    // prior serves' escapes are proposable from round 1 of THIS serve.
                    if let Some(ta) = trim_adapt.as_mut() {
                        let learned: Vec<u32> = std::fs::read_to_string(format!("{path}.learned"))
                            .map(|t| t.lines().filter_map(|l| l.trim().parse().ok()).collect())
                            .unwrap_or_default();
                        let mut n_pre = 0usize;
                        for id in learned {
                            let i = id as usize;
                            if i < n_vocab && !ta.present[i] && ta.used < ta.n_spare {
                                let slot = ta.spare_base + ta.used;
                                ta.used += 1;
                                ta.present[i] = true;
                                let off = i * row_bytes;
                                gathered[slot * row_bytes..(slot + 1) * row_bytes]
                                    .copy_from_slice(&t.bytes[off..off + row_bytes]);
                                d2t[slot] = id;
                                n_pre += 1;
                            }
                        }
                        if n_pre > 0 {
                            eprintln!(
                                "[trim-adapt] {n_pre} learned rows pre-filled from {path}.learned"
                            );
                        }
                    }
                    // upload AFTER the sidecar pre-fill wrote its rows into `gathered`.
                    let bytes = e.htod_bytes(&gathered)?;
                    (
                        GpuTensor::Quant {
                            bytes,
                            qtype,
                            row_bytes,
                            ne: vec![in_f as u64, d2t.len() as u64],
                            scale: 1.0,
                            rp: false,
                            #[cfg(memra_cutlass)]
                            cutlass: None,
                            fp8: None,
                            blk: None,
                            rp4: None,
                            f16: None,
                        },
                        Some(d2t),
                        trim_adapt,
                    )
                }
                None => (load_t(e, &src, "token_embd.weight")?, None, None),
            }
        };
        // Q4_0 split-plane decode mirrors (MEMRA_Q4RP, same as the main trunk — see hybrid.rs):
        // the draft chain is 3 serial mmvq trips/round; the head alone is ~137MB/draft.
        // projection tensor prefix: 26B/31B "nextn.", the E4B assistant "mtp.".
        let proj_prefix = if src.find("nextn.pre_projection.weight").is_some() {
            "nextn"
        } else {
            "mtp"
        };
        let (mut pre_proj, mut post_proj) = (
            load_t(e, &src, &format!("{proj_prefix}.pre_projection.weight"))?,
            load_t(e, &src, &format!("{proj_prefix}.post_projection.weight"))?,
        );
        let mut head = head;
        let mut layers = layers;
        if crate::Engine::q4rp_enabled() {
            // adaptive-trim heads skip the split-plane mirror: the mmvq _rp twins read the
            // MIRROR, so an in-place row learn on `bytes` would be invisible to the matmul.
            let head_ws: &mut [&mut GpuTensor] = if trim_adapt.is_some() {
                &mut [&mut pre_proj, &mut post_proj]
            } else {
                &mut [&mut pre_proj, &mut post_proj, &mut head]
            };
            for w in head_ws.iter_mut() {
                e.build_q4_rp4(w)?;
            }
            for l in layers.iter_mut() {
                for w in [
                    &mut l.wq,
                    &mut l.wo,
                    &mut l.ffn_gate,
                    &mut l.ffn_up,
                    &mut l.ffn_down,
                ] {
                    e.build_q4_rp4(w)?;
                }
            }
        }
        let d2t_dev = match &d2t {
            Some(m) => Some(e.stream().clone_htod(&m[..])?),
            None => None,
        };
        Ok(GemmaDraft {
            layers,
            pre_proj,
            post_proj,
            output_norm: load_t(e, &src, "output_norm.weight")?,
            head,
            d2t,
            d2t_dev,
            trim_adapt,
            rope_freqs,
            ones: e.htod(&[1.0f32; 512])?,
            n_embd,
            n_backbone,
            rope_base_global: meta_f("rope.freq_base", 1e6),
            rope_base_swa: meta_f("rope.freq_base_swa", 1e4),
            sliding_window: meta_u("attention.sliding_window") as usize,
        })
    }
}

impl HybridModel {
    /// The MAIN layer whose KV cache a drafter layer attends (llama-model.cpp:2139):
    /// the last OWN-KV layer of the class — `boundary - 2` windowed / `boundary - 1`
    /// global, where boundary = n_layer - shared_kv_layers. Shared across every
    /// gemma4-assistant drafter (26B/31B: boundary = n_layer; E4B: 24).
    pub(crate) fn gemma4_draft_kv_target(&self, swa: bool) -> usize {
        let shared = self
            .cfg
            .gemma4
            .as_ref()
            .map(|g| g.shared_kv_layers as usize)
            .unwrap_or(0);
        let boundary = self.layers.len() - shared;
        boundary - if swa { 2 } else { 1 }
    }

    /// One drafter step: (token, h[2816 device]) at absolute position `pos` over the FROZEN main
    /// cache. Returns (draft logits host [n_vocab], h_next [2816 device]).
    pub fn gemma4_draft_step(
        &self,
        e: &Engine,
        d: &GemmaDraft,
        token: u32,
        h: &CudaSlice<f32>,
        pos: usize,
        cache: &Cache,
    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let (hn, h_next) = self.gemma4_draft_trunk(e, d, token, h, pos, cache)?;
        let logits = e.dtoh(&e.matmul(&d.head, &hn, 1)?)?;
        Ok((logits, h_next))
    }

    /// Drafter trunk with the token in DEVICE memory (a 1-elem view of the round's batch
    /// buffer) — zero host traffic.
    #[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 gemma4_draft_trunk_dev(
        &self,
        e: &Engine,
        d: &GemmaDraft,
        tok_v: &cudarc::driver::CudaView<u32>,
        h: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &Cache,
        dc_bucket: Option<usize>,
    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let nb = d.n_backbone;
        let embd_gpu = self
            .embd_gpu
            .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload"));
        let (qt, rb) = self.embd.qt_and_row_bytes(nb);
        let mut xs = e.embed_gather_device_tv(embd_gpu, tok_v, 1, nb, qt, rb)?;
        e.scale_inplace(&mut xs, (nb as f32).sqrt(), nb)?;
        self.gemma4_draft_trunk_from_x(e, d, &xs, h, pos_d, cache, dc_bucket)
    }

    /// Drafter trunk: returns (post-output_norm hidden [1024], h_next [2816]).
    fn gemma4_draft_trunk(
        &self,
        e: &Engine,
        d: &GemmaDraft,
        token: u32,
        h: &CudaSlice<f32>,
        pos: usize,
        cache: &Cache,
    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let nb = d.n_backbone;
        let mut xs = e.htod(&self.embd.try_gather(nb, &[token])?)?;
        e.scale_inplace(&mut xs, (nb as f32).sqrt(), nb)?;
        let pos_d = e.htod_i32(&[pos as i32])?;
        self.gemma4_draft_trunk_from_x(e, d, &xs, h, &pos_d, cache, None)
    }

    /// Trunk body from the pre-scaled main-embed row.
    #[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 gemma4_draft_trunk_from_x(
        &self,
        e: &Engine,
        d: &GemmaDraft,
        xs: &CudaSlice<f32>,
        h: &CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        cache: &Cache,
        dc_bucket: Option<usize>,
    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
        // pos rides a DEVICE slot (burst-arc step a, 2026-07-12): the round fills persistent
        // slots via set_i32_one (kernel-arg stores — no per-step htod/alloc) and the chain
        // becomes graph-capturable (an in-graph i32_copy_add can feed the slots later).
        let eps = self.cfg.rms_eps;
        let ne = d.n_embd;

        // xh = concat(x, h) [2*n_backbone]
        let nb = d.n_backbone;
        let mut xh = e.uninit(2 * nb)?;
        e.copy_into(&mut xh, 0, xs, nb)?;
        e.copy_into(&mut xh, nb, h, nb)?;

        let mut cur = e.matmul(&d.pre_proj, &xh, 1)?; // [1024]

        for dl in d.layers.iter() {
            // attention over the shared MAIN KV: swa -> the last OWN-KV windowed layer,
            // global -> the last OWN-KV global layer (llama-model.cpp:2139 rule). Plain
            // 26B/31B trunks have no shared tail, so this is n-2 / n-1 there; E4B's 18
            // KV-shared tail layers move the boundary to 24 -> targets 22 (swa) / 23.
            let main_il = self.gemma4_draft_kv_target(dl.swa);
            let kvl = cache.kv[main_il].as_ref().unwrap();
            let (hd, nhh) = (dl.hd, dl.nh);
            let nkv = kvl.kv_dim_k / hd;
            let base = if dl.swa {
                d.rope_base_swa
            } else {
                d.rope_base_global
            };

            let mut hn = e.uninit(ne)?;
            e.rms_norm(&cur, dl.attn_norm.float_data(), &mut hn, ne, 1, eps)?;
            let q0 = e.matmul(&dl.wq, &hn, 1)?;
            let mut q = e.uninit(nhh * hd)?;
            e.rms_norm(&q0, dl.q_norm.float_data(), &mut q, hd, nhh, eps)?;
            if dl.swa {
                e.rope_neox(&mut q, pos_d, hd, hd, nhh, 1, base, 1.0)?;
            } else {
                e.rope_neox_ff(&mut q, pos_d, hd, hd, nhh, 1, base, 1.0, &d.rope_freqs)?;
            }
            let avail = kvl.len;
            let win = d.sliding_window;
            let mut attn = e.uninit(nhh * hd)?;
            // drafter attends the MAIN cache — its format follows the main layer's class
            // (windowed L28 = wkv arm, global L29 = gkv arm; gkv routing is hd-keyed inside).
            // DEVICE-LEN arms (burst arc): the length rides the main layer's len_d counter
            // so the chain is replay-correct across rounds. dc_bucket = the RUNG the round
            // derived (power-of-2, shared by eager and captured replays — same n_splits,
            // same combine order; the main graph arc's bucket lesson). None = host-len arm.
            if let Some(bucket) = dc_bucket {
                let k_view = e.view_u8(&kvl.k, kvl.k.len());
                let v_view = e.view_u8(&kvl.v, kvl.v.len());
                if dl.swa && avail > win {
                    e.fa_decode_rows_w(
                        &q,
                        &k_view,
                        &v_view,
                        &mut attn,
                        hd,
                        nhh,
                        nkv,
                        &kvl.len_d,
                        -1,
                        1,
                        1.0,
                        win,
                        kvl.k_tok_bytes,
                        kvl.v_tok_bytes,
                        None,
                    )?;
                } else {
                    e.fa_decode_dc(
                        &q,
                        &k_view,
                        &v_view,
                        &mut attn,
                        hd,
                        nhh,
                        nkv,
                        &kvl.len_d,
                        bucket,
                        1.0,
                        kvl.k_tok_bytes,
                        kvl.v_tok_bytes,
                        dl.swa && crate::Engine::wkv_on(),
                    )?;
                }
            } else {
                let (off_tok, t_kv) = if dl.swa && avail > win {
                    (avail - win, win)
                } else {
                    (0, avail)
                };
                let k_view = e.view_u8_range(
                    &kvl.k,
                    off_tok * kvl.k_tok_bytes,
                    (off_tok + t_kv) * kvl.k_tok_bytes,
                );
                let v_view = e.view_u8_range(
                    &kvl.v,
                    off_tok * kvl.v_tok_bytes,
                    (off_tok + t_kv) * kvl.v_tok_bytes,
                );
                e.fa_decode_kvmod(
                    &q,
                    &k_view,
                    &v_view,
                    &mut attn,
                    hd,
                    nhh,
                    nkv,
                    t_kv,
                    1.0,
                    kvl.k_tok_bytes,
                    kvl.v_tok_bytes,
                    dl.swa && crate::Engine::wkv_on(),
                )?;
            }
            let o = e.matmul(&dl.wo, &attn, 1)?;

            let mut post = e.uninit(ne)?;
            e.rms_norm(&o, dl.post_attn_norm.float_data(), &mut post, ne, 1, eps)?;
            let mut attn_out = e.uninit(ne)?;
            e.add(&post, &cur, &mut attn_out, ne)?;

            let mut z = e.uninit(ne)?;
            e.rms_norm(&attn_out, dl.ffn_norm.float_data(), &mut z, ne, 1, eps)?;
            let n_ff = dl.ffn_gate.out_features();
            let gate = e.matmul(&dl.ffn_gate, &z, 1)?;
            let up = e.matmul(&dl.ffn_up, &z, 1)?;
            let mut act = e.uninit(n_ff)?;
            e.gelu_tanh_mul(&gate, &up, &mut act, n_ff)?;
            let f0 = e.matmul(&dl.ffn_down, &act, 1)?;
            let mut fpost = e.uninit(ne)?;
            e.rms_norm(&f0, dl.ffn_post_norm.float_data(), &mut fpost, ne, 1, eps)?;
            let mut xn = e.uninit(ne)?;
            e.add_scale(&fpost, &attn_out, dl.out_scale, &mut xn, ne)?;
            cur = xn;
        }

        let mut hn = e.uninit(ne)?;
        e.rms_norm(&cur, d.output_norm.float_data(), &mut hn, ne, 1, eps)?;
        let h_next = e.matmul(&d.post_proj, &hn, 1)?; // [2816]; head applied by callers (NO softcap)
        Ok((hn, h_next))
    }

    /// Greedy draft step: like gemma4_draft_step but the token argmax stays on device —
    /// host sees 4 bytes (no 1MB logits dtoh per draft). Returns (token, h_next).
    pub fn gemma4_draft_step_greedy(
        &self,
        e: &Engine,
        d: &GemmaDraft,
        token: u32,
        h: &CudaSlice<f32>,
        pos: usize,
        cache: &Cache,
    ) -> Result<(u32, CudaSlice<f32>), Box<dyn std::error::Error>> {
        let (hn, h_next) = self.gemma4_draft_trunk(e, d, token, h, pos, cache)?;
        let ld = e.matmul(&d.head, &hn, 1)?;
        let tok_d = e.argmax_token_device(&ld, d.head.out_features())?;
        let idx = e.dtoh_u32(&tok_d)?[0];
        let tok = match &d.d2t {
            Some(map) => map[idx as usize],
            None => idx,
        };
        Ok((tok, h_next))
    }
}

impl HybridModel {
    /// gemma4 MTP greedy spec loop: prime the prompt, then rounds of (chained K-token draft
    /// over the frozen main cache) + (ONE batched verify) + longest-prefix accept + KV rollback.
    /// Returns generated tokens; prints acceptance stats.
    #[allow(clippy::too_many_arguments)]
    #[allow(clippy::unnecessary_unwrap)] // allow: the Some-guards sit in multi-clause regime gates; if-let would reshape the arm structure
    #[allow(clippy::map_entry)] // allow: the init bodies are fallible (`?`); Entry::or_insert_with cannot propagate errors
    pub fn generate_spec_gemma(
        &self,
        e: &Engine,
        d: &mut GemmaDraft,
        prompt: &[u32],
        max_new: usize,
        k: usize,
        eos: &[u32],
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        let n_embd = self.cfg.n_embd as usize;
        let eps = self.cfg.rms_eps;
        let mut cache = Cache::new(e, &self.cfg, prompt.len() + max_new + k + 8)?;

        // Adaptive trim, learn point 1: the PROMPT's own ids — the measured escapees are the
        // prompt's domain content words echoed back (▁oceans, clouds, Explain...), so the
        // prompt is the cheapest predictor of what the trim is about to miss.
        trim_adapt_learn(e, d, prompt)?;

        let t_prime = std::time::Instant::now();
        // short prompts fall below prime_cache's T floor — the batched verify IS a prime.
        let (pl, h_seed) = if prompt.len() >= crate::hybrid_forward::PRIME_MIN_T {
            let (l, hs, _hh) = self.prime_cache(e, prompt, &mut cache, 0)?;
            (l, hs)
        } else if self.is_gemma4_e4b() {
            // E4B short-prompt prime: TOKENWISE — the batched e4b trunk at base_len==0
            // rides the PRIME-FA f32 arm (a different numerics class from the plain arm's
            // tokenwise prime), and the class skew flipped near-tie streams (3/64,
            // 2026-07-13). decode_step_h is the same chain the plain arm primes with.
            let n_embd_ = self.cfg.n_embd as usize;
            let mut ll = Vec::new();
            let mut hx = e.zeros(n_embd_)?;
            for &tok in prompt {
                let (l, hh) = self.gemma4_e4b_decode_step_h(e, tok, &mut cache)?;
                ll = l;
                hx = hh;
            }
            // decode_step_h returns the PRE-output_norm hidden; the short-prompt arm's
            // h convention below is POST-norm — norm here.
            let mut hp = e.uninit(n_embd_)?;
            e.rms_norm(&hx, self.output_norm.float_data(), &mut hp, n_embd_, 1, eps)?;
            (ll, hp)
        } else {
            let n_vocab = self.output.out_features();
            let (lv, hv) = self.gemma4_decode_step_t_h(e, prompt, 0, &mut cache)?;
            let t = prompt.len();
            let last = lv[(t - 1) * n_vocab..t * n_vocab].to_vec();
            // NOTE hv rows are POST-output_norm; h_seed convention below expects PRE-norm and
            // re-norms — so recover a pre-norm-free path: use the post-norm row DIRECTLY.
            let hvv = e.view(&hv, t * n_embd);
            let row = hvv.slice((t - 1) * n_embd..t * n_embd);
            let mut hrow = e.uninit(n_embd)?;
            e.copy_view_into(&mut hrow, 0, &row, n_embd)?;
            // mark: already post-norm — skip the re-norm below via the flag
            (last, hrow)
        };
        e.stream().synchronize()?;
        crate::PRIME_NANOS.store(
            t_prime.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        // drafter h = POST-output_norm hidden (llama h_nextn); prime returns PRE-norm h_seed,
        // the short-prompt verify path already returns post-norm rows.
        let mut h = if prompt.len() >= crate::hybrid_forward::PRIME_MIN_T {
            let mut hh = e.uninit(n_embd)?;
            e.rms_norm(
                &h_seed,
                self.output_norm.float_data(),
                &mut hh,
                n_embd,
                1,
                eps,
            )?;
            hh
        } else {
            h_seed
        };

        let mut last = crate::forward::argmax(&pl) as u32;
        // MEMRA_PROFILE_SPEC=2: capture starts at the ROUND LOOP (prime excluded) — pair
        // with `nsys -c cudaProfilerApi` (the qwen loop's pattern, spec.rs).
        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
            unsafe extern "C" {
                fn cudaProfilerStart() -> i32;
            }
            unsafe {
                cudaProfilerStart();
            }
        }
        let mut out: Vec<u32> = Vec::with_capacity(max_new);
        let (mut drafted, mut accepted, mut rounds) = (0usize, 0usize, 0usize);
        // per-position accept histogram (MEMRA_SPEC_STATS): [attempted, accepted] per slot —
        // the depth-K policy statistic (deep slots' marginal accept decides fixed-cap vs deep).
        let mut pos_att = [0usize; 16];
        let mut pos_acc = [0usize; 16];

        // ASYNC ROUND v2 (dc class): the whole draft chain + verify enqueue with ZERO host
        // syncs — token seeds via kernel-arg store (u32_set_k, no host-memory transfer), draft
        // argmaxes land in the batch buffer, verify argmaxes in vam_d; ONE pack + ONE dtoh of
        // (k drafts + k+1 vam) closes the round. (v1 with memcpy_htod seeding measured
        // NEGATIVE — the pageable-copy sync; this is the retry with the sync removed.)
        let mut batch_d = e.stream().alloc_zeros::<u32>(k + 1)?;
        let mut packed = e.stream().alloc_zeros::<u32>(2 * k + 1)?;
        // confidence-adaptive depth (MEMRA_SPEC_PMIN, default 0 = off): per-draft probs.
        let pmin: f32 = std::env::var("MEMRA_SPEC_PMIN")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.0);
        // IN-ROUND confidence cut (2026-07-28): llama's draft-mtp stops drafting the
        // moment a draft's top-1 prob falls below p-min; our MEMRA_SPEC_PMIN is one round
        // LATE by design (zero-sync round). This arm pays one small dtoh sync per draft
        // step (steps ~150µs; sync ~15µs) to cut the chain mid-round and verify at the
        // shrunk width. Eager arm only — burst/graph arms draft fixed depth.
        // DEFAULT is SELF-KEYED: active at depth (pos >= floor_ctx) and only in rounds
        // following a MISS — measured: depth cells with sub-0.9 acceptance win (26B
        // +1.4-3.2% @ 0.868-0.882 accept, 31B +2% @ 0.845-0.883), chat cells and the
        // 0.95-accept 12B depth lose under an ALWAYS-on cut (-0.9 to -6%) but their
        // rounds are mostly full-accept so the self-key idles there. Explicit
        // MEMRA_SPEC_PMIN_INROUND pins the cut at every position/round; =0 disables.
        let pmin_ir_env: Option<f32> = std::env::var("MEMRA_SPEC_PMIN_INROUND")
            .ok()
            .and_then(|v| v.parse().ok());
        const PMIN_IR_DEFAULT: f32 = 0.7;
        let mut prev_full = true; // round 1: no miss evidence yet — draft at full depth
        let mut p_d = e.stream().alloc_zeros::<f32>(k.max(1))?;

        // ADAPTIVE DRAFT LENGTH (default ON 2026-07-10; MEMRA_SPEC_ADAPT=0 reverts): llama's
        // draft-mtp reaches 0.64-0.70 acceptance on the SAME drafter (ours fixed-K: 0.52) by
        // drafting fewer tokens when unconfident (p-min gate). Zero-sync host proxy: next
        // round's depth = last round's accepted run + 1, clamped to [floor=1, k] — rounds
        // after a miss shrink, streaks re-deepen. The round's ONE dtoh already carries the
        // acceptance; no new syncs. Policy sweep (short chat, N=1 each): floor1/cap3 239.2
        // vs fixed-K3 231.1 (+3.5%, accept .52->.58); floor2 and cap4/5 all worse.
        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() != Ok("0");
        // ADAPTIVE FLOOR default is per-model (MEMRA_SPEC_ADAPT_FLOOR overrides): the floor-1
        // policy collapses to shallow drafts after any miss and pays a slow re-deepen; on
        // models with an expensive verify step the deep-draft upside dwarfs the wasted-draft
        // cost. Measured 2026-07-25 (chat cell, own-gen trim; peak grids both models):
        // 31B K=5 floor=4 120.2 vs floor=1 103.8 (+15.7%, N=3; floor 5-6 falls off);
        // 12B K=4-5 floor=4 240.5-240.8 vs floor=1 200.6 (+20%, floor 5+ falls off).
        // The floor clamps to k_cap, so shallow-K callers are unaffected.
        // 26B tier (2026-07-26 re-sweep under the f16pv spec flip): floor=2 wins BOTH its
        // cells — short 329.5 vs 307.0 floor1 (+7%, best at every K), depth 329.7 vs ~318
        // (the 2026-07-10 "floor2 worse" verdict predates the flip and is superseded).
        // E4B (n_embd < 2500) keeps floor=1 — unmeasured, cheap verify.
        let adapt_floor_default: usize = if self.cfg.n_embd >= 3500 {
            4
        } else if self.cfg.n_embd >= 2500 {
            2
        } else {
            1
        };
        // (stream-k spec key lives in HybridModel::load_from_source_impl — it must be set
        // before the PRIME's GEMMs autotune, not here.)
        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
            .ok()
            .and_then(|v| v.parse().ok());
        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
        // POSITION KEY (2026-07-26): the HIGH floor is a SHORT-CTX win. At depth the
        // per-position acceptance is lower and FORCED-DEEP drafts turn net-negative:
        // 31B d1736 floor4 99-101 and floor2 97.4-99.8 @ 0.758-0.778 vs floor1
        // 103.8-104.2 @ 0.817 (two perf-ci batteries + flip-tree N=2 — floor2 is a REAL
        // small loss there, not noise), while its chat cell holds +15-20% under floor4.
        // The 26B is the opposite at depth: its mild floor2 WINS (304-305 vs ~297).
        // Default: full floor while pos < floor_ctx; past it HIGH-floor models (>=4)
        // relax to 1, MILD-floor models keep their floor. MEMRA_SPEC_FLOOR_CTX overrides
        // the boundary; an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(1024);
        let floor_at = |pos: usize| -> usize {
            if adapt_floor_env.is_some() || pos < floor_ctx {
                adapt_floor
            } else if adapt_floor >= 4 {
                1
            } else {
                adapt_floor
            }
        };
        // cap ceiling 7 by default; MEMRA_SPEC_CAPMAX opens the b16 verify tier (t=9..16).
        // The historical cap>=8 "crash" was two host bugs, both fixed 2026-07-12: round 1
        // ran UNCLAMPED (`kc = k` — verify t=K+1 entered the b16 tier while it was gated)
        // and the b16 dispatch requested _r2 twins that were never compiled (mcols==16 now
        // forces the base variant). Stream gates arbitrate any raised cap.
        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(7);
        let k_cap = k.min(cap_max).max(1);
        // DRAFT-CHAIN GRAPHS (burst-arc step c, MEMRA_GEMMA_DRAFT_GRAPH=1): the whole k-step
        // draft chain replays as ONE captured graph — position slots fill in-graph,
        // the seed hidden rides the persistent g_seed buffer, KV lengths ride len_d (step b).
        // Keyed on (kr, rung, over_win): a new depth/rung/window regime captures lazily.
        let graph_on = std::env::var("MEMRA_GEMMA_DRAFT_GRAPH").as_deref() == Ok("1");
        #[allow(clippy::type_complexity)]
        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
        let mut draft_graphs: std::collections::HashMap<
            (usize, usize, bool),
            (
                cudarc::driver::CudaGraph,
                Vec<Box<dyn std::any::Any + Send>>,
            ),
        > = Default::default();
        let mut g_seed = e.zeros(n_embd)?;
        // seed len_d before round 1 (prime went through the host-len path).
        for kvl in cache.kv.iter_mut().flatten() {
            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
        }
        // persistent per-step rope-pos slots (device; filled by set_i32_one kernel-arg stores).
        let mut pos_slots: Vec<CudaSlice<i32>> = (0..k_cap.max(1))
            .map(|_| e.htod_i32(&[0]))
            .collect::<Result<_, _>>()?;
        // clamp round 1 too (the leak above).
        let mut kc = k_cap;
        // round-graph arm state (MEMRA_GEMMA_ROUND_GRAPH): stream buffers, the fill dummy, the
        // kv-len pointer table and the verify scratch, allocated on first use.
        let mut burst_state: Option<(
            crate::round_stream::StreamBufs,
            CudaSlice<f32>,
            CudaSlice<u64>,
            crate::hybrid_forward::VerifyStreamScratch,
        )> = None;
        let win_main = self
            .cfg
            .gemma4
            .as_ref()
            .map(|g| g.sliding_window as usize)
            .unwrap_or(0);
        let g4_shared = self
            .cfg
            .gemma4
            .as_ref()
            .map(|g| g.shared_kv_layers)
            .unwrap_or(0);
        'outer: while out.len() < max_new {
            let mut kr = if adapt { kc } else { k_cap };
            // power-of-2 rung bucket for the dc arms (shared by eager and captured replays);
            // MEMRA_GEMMA_DRAFT_DC=0 reverts to the host-len kvmod arm.
            let dc_bucket: Option<usize> = {
                static DC: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
                if *DC.get_or_init(|| std::env::var("MEMRA_GEMMA_DRAFT_DC").as_deref() != Ok("0")) {
                    let ml = cache
                        .kv
                        .iter()
                        .flatten()
                        .map(|kv| kv.len)
                        .max()
                        .unwrap_or(1);
                    Some((ml + k_cap + 2).next_power_of_two().max(512))
                } else {
                    None
                }
            };
            e.u32_set_k(&mut batch_d, last, 0)?;
            e.copy_into(&mut g_seed, 0, &h, n_embd)?;
            // the draft chain, step j: reads g_seed via the hc chain, pos from pos_slots[j]
            // (eager: host-filled; graph: filled in-graph).
            let run_chain = |e: &Engine,
                             d: &GemmaDraft,
                             batch_d: &mut CudaSlice<u32>,
                             p_d: &mut CudaSlice<f32>,
                             g_seed: &CudaSlice<f32>,
                             pos_slots: &Vec<CudaSlice<i32>>,
                             inround: f32|
             -> Result<usize, Box<dyn std::error::Error>> {
                // uninit+copy (NOT clone_dtod): clone_dtod's internal alloc bypasses the
                // capture-retain hooks — its address got pool-reused between replays and the
                // replayed chain read a corrupted seed (accept 0.52 vs 0.76).
                let mut hc = e.uninit(n_embd)?;
                e.copy_into(&mut hc, 0, g_seed, n_embd)?;
                for j in 0..kr {
                    let tv = batch_d.slice(j..j + 1);
                    let (hn, h_next) = self.gemma4_draft_trunk_dev(
                        e,
                        d,
                        &tv,
                        &hc,
                        &pos_slots[j],
                        &cache,
                        dc_bucket,
                    )?;
                    let ld = e.matmul(&d.head, &hn, 1)?;
                    e.argmax_token_device_col(&ld, 0, d.head.out_features(), batch_d, j + 1)?;
                    // confidence-adaptive depth (MEMRA_SPEC_PMIN): TRIM-space prob before d2t.
                    if pmin > 0.0 || inround > 0.0 {
                        e.prob_of_token_device_col(
                            &ld,
                            batch_d,
                            j + 1,
                            p_d,
                            j,
                            d.head.out_features(),
                        )?;
                    }
                    // FR-trimmed head: translate the trim-space argmax to the vocab id.
                    if let Some(map) = &d.d2t_dev {
                        e.u32_map_k(batch_d, map, j + 1)?;
                    }
                    hc = h_next;
                    // IN-ROUND cut: one small dtoh sync per step; stop drafting the moment
                    // confidence falls below the gate and verify at the shrunk width.
                    // (A DSpark-class marginal-rate window — S_{j+1}*T(j) > E[tok](j)*t_d
                    // with profiled t_draft/t_verify EMAs — measured FLAT here 2026-07-30:
                    // never cuts at accept >= 0.8, par-to-noise on 26B/31B depth x3
                    // interleaved; arm removed per flags doctrine, jsonl row is the record.)
                    if inround > 0.0 && j + 1 < kr {
                        let ph = e.dtoh(p_d)?;
                        if ph[j] < inround {
                            return Ok(j + 1);
                        }
                    }
                }
                Ok(kr)
            };
            let over_win = {
                let win = d.sliding_window;
                d.layers.iter().any(|dl| {
                    dl.swa
                        && cache.kv[self.gemma4_draft_kv_target(true)]
                            .as_ref()
                            .is_some_and(|kv| kv.len > win)
                })
            };
            // ---- ROUND-GRAPH ARM ---- (MEMRA_GEMMA_ROUND_GRAPH=1): the WHOLE round —
            // draft chain + stream verify + device accept/seed/rollback/commit + the
            // device adaptive-depth update — captured ONCE per (k_cap, rung, over_win)
            // regime and replayed as ONE graph launch per round (the llama round-cost
            // mechanism: ~600 per-round enqueues collapse to 1). The round is SELF-FEEDING
            // (pos_ctr/pend/brk/g_seed all advance in-graph), so the capture warmups are
            // simply two SERVED rounds — their tokens land in the ring and drain normally
            // (no snapshot/rollback needed, unlike the E4B token door).
            // Adaptive K rides brk[0] via spec_adapt_k: drafts always run k_cap deep (the
            // drafter is cheap) but the accept walk depth follows the host policy exactly.
            let round_graph_on = std::env::var("MEMRA_GEMMA_ROUND_GRAPH").as_deref() == Ok("1");
            if round_graph_on
                && dc_bucket.is_some()
                && pmin == 0.0
                && g4_shared == 0
                && !self.is_gemma4_e4b()
                && (cache.pos + 2 * (k_cap + 1) + k_cap + 4 < win_main || cache.pos > win_main)
                && (cache.pos + 2 * (k_cap + 1) + k_cap + 4 < crate::fa512_min_tkv()
                    || cache.pos + 1 >= crate::fa512_min_tkv())
                && e.fa_rows_eligible(cache.pos, 256)
                && cache.pos + 2 * (k_cap + 1) + k_cap + 2 <= cache.max_ctx
            {
                if burst_state.is_none() {
                    // ring sized for the capture warmups (2 rounds) + the live round.
                    let bufs = crate::round_stream::StreamBufs::new(e, k_cap, 3)?;
                    let fill_dummy = e.zeros(n_embd)?;
                    let ptrs =
                        crate::round_stream::kv_len_ptr_table(e, &cache, Some(&bufs.pos_ctr))?;
                    let scr = self.verify_stream_scratch(e, k_cap + 1)?;
                    burst_state = Some((bufs, fill_dummy, ptrs, scr));
                }
                // entry: `last` is the pending token (emitted at drain), h is the seed.
                let (bufs, fill_dummy, ptrs, scr) = burst_state.as_mut().unwrap();
                let n_rows = cache.kv.len() + 1;
                e.set_i32_one(&mut bufs.pos_ctr, cache.pos as i32)?;
                e.u32_set_k(&mut bufs.ring_d, 0, 0)?;
                e.u32_set_k(&mut bufs.pend_d, last, 0)?;
                e.u32_set_k(&mut bufs.brk_d, (if adapt { kc } else { k_cap }) as u32, 0)?;
                e.u32_set_k(&mut bufs.brk_d, 1, 1)?;
                e.copy_into(&mut g_seed, 0, &h, n_embd)?;
                // entry pend is emitted host-side (the ring only carries accepted drafts
                // + bonuses — the burst-arm contract).
                out.push(last);
                if eos.contains(&last) {
                    break 'outer;
                }
                if out.len() >= max_new {
                    break 'outer;
                }
                #[allow(clippy::unnecessary_unwrap)]
                // allow: the Some-guard sits in a multi-clause regime gate; if-let would reshape the arm structure
                let key = (usize::MAX - k_cap, dc_bucket.unwrap(), over_win);
                let mut fresh_rounds = 1usize; // rounds executed by this iteration
                // `hint` is the verify stream's ARM-GATING upper bound — it must sit on
                // the SAME side of every crossover as the live lengths this capture
                // serves, INCLUDING the arms' own margins (`hint + t < f512` gates the
                // global scalar arm; `hint + 1 >= win` gates rows_w), or the captured
                // verify bakes a different kernel class than the eager reference
                // (107-vs-106 / 4-64 drifts; the regime gate above guarantees the live
                // side with the same margins).
                let hint = if cache.pos > win_main {
                    dc_bucket.unwrap() + k_cap + 2 // over-window: rows_w regime
                } else if cache.pos + 1 >= crate::fa512_min_tkv() {
                    win_main - 2 // above f512, under window
                } else {
                    crate::fa512_min_tkv().saturating_sub(k_cap + 5) // under both
                };
                let bufs_ptr: *mut crate::round_stream::StreamBufs = &mut *bufs;
                let scr_ptr: *mut crate::hybrid_forward::VerifyStreamScratch = &mut *scr;
                let cache_ptr: *mut Cache = &mut cache;
                let batch_ptr: *mut CudaSlice<u32> = &mut batch_d;
                let seed_ptr: *mut CudaSlice<f32> = &mut g_seed;
                let slots_ptr: *mut Vec<CudaSlice<i32>> = &mut pos_slots;
                let mut round_body = |e: &Engine| -> Result<(), Box<dyn std::error::Error>> {
                    // SAFETY: single-threaded round body; the raw pointers alias the outer
                    // &mut only within this closure (no overlapping borrows).
                    let (bufs, scr, cache, batch_d, g_seed, pos_slots) = unsafe {
                        (
                            &mut *bufs_ptr,
                            &mut *scr_ptr,
                            &mut *cache_ptr,
                            &mut *batch_ptr,
                            &mut *seed_ptr,
                            &mut *slots_ptr,
                        )
                    };
                    e.i32_copy_add(&bufs.pos_ctr, &mut bufs.pos_start_d, 0)?;
                    e.u32_copy(&bufs.pend_d, batch_d)?;
                    for (j, slot) in pos_slots.iter_mut().take(k_cap).enumerate() {
                        e.i32_copy_add(&bufs.pos_ctr, slot, j as i32)?;
                    }
                    let mut hc = e.uninit(n_embd)?;
                    e.copy_into(&mut hc, 0, g_seed, n_embd)?;
                    #[allow(clippy::needless_range_loop)]
                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
                    for j in 0..k_cap {
                        let tv = batch_d.slice(j..j + 1);
                        let (hn, h_next) = self.gemma4_draft_trunk_dev(
                            e,
                            d,
                            &tv,
                            &hc,
                            &pos_slots[j],
                            cache,
                            dc_bucket,
                        )?;
                        let ld = e.matmul(&d.head, &hn, 1)?;
                        e.argmax_token_device_col(&ld, 0, d.head.out_features(), batch_d, j + 1)?;
                        if let Some(map) = &d.d2t_dev {
                            e.u32_map_k(batch_d, map, j + 1)?;
                        }
                        hc = h_next;
                    }
                    let (vam_d, vh) = self.gemma4_verify_t_am_stream(
                        e,
                        batch_d,
                        k_cap + 1,
                        &bufs.pos_ctr,
                        hint,
                        cache,
                        scr,
                    )?;
                    e.spec_accept_greedy_dc(
                        &vam_d,
                        batch_d,
                        &bufs.last_pred_d,
                        &bufs.brk_d,
                        &mut bufs.acc_d,
                    )?;
                    if std::env::var("MEMRA_DEBUG_SPEC").as_deref() == Ok("1")
                        && std::env::var("MEMRA_ROUND_GRAPH_CHECK").as_deref() == Ok("1")
                    {
                        let vhh = e.dtoh(&vh)?;
                        let nrm = |r: usize| {
                            vhh[r * n_embd..(r + 1) * n_embd]
                                .iter()
                                .map(|x| x * x)
                                .sum::<f32>()
                                .sqrt()
                        };
                        let vamh = e.dtoh_u32(&vam_d)?;
                        eprintln!(
                            "[rg-vh] |row0|={:.3} |row1|={:.3} |row2|={:.3} vam={:?}",
                            nrm(0),
                            nrm(1),
                            nrm(2),
                            &vamh[..(k_cap + 1).min(7)]
                        );
                    }
                    e.spec_seed_gather(&vh, fill_dummy, &bufs.acc_d, g_seed, 1, n_embd)?;
                    e.spec_rollback_stream(ptrs, &bufs.pos_start_d, &bufs.acc_d, 1, n_rows)?;
                    e.spec_ring_commit(
                        batch_d,
                        &bufs.acc_d,
                        &bufs.brk_d,
                        &mut bufs.ring_d,
                        &mut bufs.pend_d,
                    )?;
                    e.spec_adapt_k(&bufs.acc_d, &mut bufs.brk_d, floor_at(cache.pos), k_cap)?;
                    Ok(())
                };
                // MEMRA_ROUND_GRAPH_CHECK=1: run the body EAGERLY (no capture/replay) —
                // splits "body semantics wrong" from "replay mechanics wrong".
                let body_check = std::env::var("MEMRA_ROUND_GRAPH_CHECK").as_deref() == Ok("1");
                if body_check {
                    round_body(e)?;
                    if std::env::var("MEMRA_DEBUG_SPEC").as_deref() == Ok("1") {
                        let acc = e.dtoh_u32(&bufs.acc_d)?;
                        let brk = e.dtoh_u32(&bufs.brk_d)?;
                        let bt = e.dtoh_u32(&batch_d)?;
                        let tgt = self.gemma4_draft_kv_target(true);
                        let ld = e.dtoh_i32(&cache.kv[tgt].as_ref().unwrap().len_d)?[0];
                        let gs = e.dtoh(&g_seed)?;
                        let gn: f32 = gs.iter().map(|x| x * x).sum::<f32>().sqrt();
                        eprintln!(
                            "[rg-check] pos0={} batch={bt:?} n_acc={} bonus={} brk_next={:?} len_d[L{tgt}]={ld} |g_seed|={gn:.3}",
                            cache.pos, acc[0], acc[1], brk
                        );
                    }
                } else {
                    if !draft_graphs.contains_key(&key) {
                        let g = e.capture_graph_retained(&mut round_body)?;
                        draft_graphs.insert(key, g);
                        fresh_rounds += 2; // the capture warmups were served rounds
                    }
                    draft_graphs.get(&key).unwrap().0.launch()?;
                }
                // drain: ONE host sync per iteration (warmup rounds included on capture).
                let toks = bufs.drain_ring(e)?;
                let posh = e.dtoh_i32(&bufs.pos_ctr)?[0] as usize;
                if std::env::var("MEMRA_DEBUG_SPEC").as_deref() == Ok("1") {
                    eprintln!(
                        "[round-graph] fresh={fresh_rounds} drained={} posh={posh} toks={:?}",
                        toks.len(),
                        &toks[..toks.len().min(12)]
                    );
                }
                drafted += fresh_rounds * k_cap;
                rounds += fresh_rounds;
                accepted += toks.len().saturating_sub(fresh_rounds);
                let mut ended = false;
                for &tk in &toks[..toks.len() - 1] {
                    out.push(tk);
                    if eos.contains(&tk) || out.len() >= max_new {
                        ended = true;
                        break;
                    }
                }
                last = *toks.last().unwrap();
                cache.pos = posh;
                for kvl in cache.kv.iter_mut().flatten() {
                    kvl.len = posh;
                }
                // NO allocation between replays: a pool alloc here can land on a baked
                // transient address and corrupt the next replay (the draft-graph lesson).
                // g_seed already holds the next seed (in-graph gather); copy INTO the
                // existing h buffer for the (possible) eager-arm handoff.
                e.copy_into(&mut h, 0, &g_seed, n_embd)?;
                kc = k_cap; // device brk owns the walk depth; host kc only seeds entry
                // learn point 2 (round-graph drain): ring = accepted drafts + bonuses; only
                // bonuses can be escapes, and the present-bitmap check skips the rest cheap.
                trim_adapt_learn(e, d, &toks)?;
                if ended {
                    break 'outer;
                }
                continue 'outer;
            }
            if graph_on && dc_bucket.is_some() {
                #[allow(clippy::unnecessary_unwrap)]
                // allow: the Some-guard sits in a multi-clause regime gate; if-let would reshape the arm structure
                let key = (kr, dc_bucket.unwrap(), over_win);
                if !draft_graphs.contains_key(&key) {
                    // chain-only capture; pos slots are graph INPUTS (filled eagerly before
                    // each launch, like g_seed — the in-graph copy_add fills replayed one
                    // round stale, see jsonl).
                    let g = e.capture_graph_retained(|e| {
                        run_chain(e, d, &mut batch_d, &mut p_d, &g_seed, &pos_slots, 0.0)
                            .map(|_| ())
                    })?;
                    draft_graphs.insert(key, g);
                }
                for (j, slot) in pos_slots.iter_mut().take(kr).enumerate() {
                    e.set_i32_one(slot, (cache.pos + j) as i32)?;
                }
                draft_graphs.get(&key).unwrap().0.launch()?;
                // MEMRA_DRAFT_GRAPH_CHECK=1: re-run the chain eagerly from the same state and
                // diff the drafted slots (replay-vs-eager divergence bisect).
                if std::env::var("MEMRA_DRAFT_GRAPH_CHECK").as_deref() == Ok("1") {
                    // NON-DESTRUCTIVE: compare, then restore the graph's tokens so the round
                    // proceeds exactly as it would without the check.
                    let gtoks = e.dtoh_u32(&batch_d)?;
                    for (j, slot) in pos_slots.iter_mut().take(kr).enumerate() {
                        e.set_i32_one(slot, (cache.pos + j) as i32)?;
                    }
                    run_chain(e, d, &mut batch_d, &mut p_d, &g_seed, &pos_slots, 0.0)?;
                    let etoks = e.dtoh_u32(&batch_d)?;
                    if gtoks[..=kr] != etoks[..=kr] {
                        eprintln!(
                            "[draft-graph] DIVERGE round={rounds} graph={:?} eager={:?}",
                            &gtoks[..=kr],
                            &etoks[..=kr]
                        );
                    }
                    for (j, &t) in gtoks.iter().enumerate().take(kr + 1) {
                        e.u32_set_k(&mut batch_d, t, j)?;
                    }
                }
            } else {
                for (j, slot) in pos_slots.iter_mut().take(kr).enumerate() {
                    e.set_i32_one(slot, (cache.pos + j) as i32)?;
                }
                let ir_now = match pmin_ir_env {
                    Some(p) => p, // explicit pin (0 disables)
                    None if cache.pos >= floor_ctx && !prev_full => PMIN_IR_DEFAULT,
                    None => 0.0,
                };
                kr = run_chain(e, d, &mut batch_d, &mut p_d, &g_seed, &pos_slots, ir_now)?;
            }
            drafted += kr;
            rounds += 1;
            let pos0 = cache.pos;
            // MEMRA_BURST_VCHECK=1: run the STREAM verify first on the same batch/state and
            // diff its argmaxes against the eager verify (bisect harness — the stream append
            // writes the same rows the eager append then overwrites, so state is untouched).
            let vcheck = std::env::var("MEMRA_BURST_VCHECK").as_deref() == Ok("1");
            let kvsum = |e: &Engine,
                         cache: &Cache|
             -> Result<Vec<(u64, u64)>, Box<dyn std::error::Error>> {
                let mut out = Vec::new();
                for kvl in cache.kv.iter().flatten() {
                    let kb = e.dtoh_u8(&kvl.k)?;
                    let vb = e.dtoh_u8(&kvl.v)?;
                    let lo = pos0 * kvl.k_tok_bytes;
                    let hi = (pos0 + kr + 1) * kvl.k_tok_bytes;
                    let lov = pos0 * kvl.v_tok_bytes;
                    let hiv = (pos0 + kr + 1) * kvl.v_tok_bytes;
                    out.push((
                        kb[lo..hi].iter().map(|&b| b as u64).sum(),
                        vb[lov..hiv].iter().map(|&b| b as u64).sum(),
                    ));
                }
                Ok(out)
            };
            let vam_s = if vcheck && !self.is_gemma4_e4b() {
                let mut ctr = e.htod_i32(&[pos0 as i32])?;
                e.set_i32_one(&mut ctr, pos0 as i32)?;
                let mut scr0 = self.verify_stream_scratch(e, kr + 1)?;
                let (vs, vhs) = self.gemma4_verify_t_am_stream(
                    e,
                    &batch_d,
                    kr + 1,
                    &ctr,
                    pos0 + kr + 3,
                    &mut cache,
                    &mut scr0,
                )?;
                let ss = kvsum(e, &cache)?;
                Some((e.dtoh_u32(&vs)?, ss, e.dtoh(&vhs)?))
            } else {
                None
            };
            let (vam_d, vh) = if self.is_gemma4_e4b() {
                self.gemma4_e4b_decode_step_t_am_dev(e, &batch_d, kr + 1, pos0, &mut cache)?
            } else {
                self.gemma4_decode_step_t_am_dev(e, &batch_d, kr + 1, pos0, &mut cache)?
            };
            if let Some((vs, ss, vhs)) = vam_s {
                let vhe = e.dtoh(&vh)?;
                for r in 0..kr + 1 {
                    let md = vhs[r * n_embd..(r + 1) * n_embd]
                        .iter()
                        .zip(&vhe[r * n_embd..(r + 1) * n_embd])
                        .map(|(a, b)| (a - b).abs())
                        .fold(0.0f32, f32::max);
                    if md > 1e-3 {
                        eprintln!("[vcheck-vh] round={rounds} row={r} maxdiff={md:.3e}");
                    }
                }
                let se = kvsum(e, &cache)?;
                for (il, (a, b)) in ss.iter().zip(&se).enumerate() {
                    if a != b {
                        eprintln!("[vcheck-kv] round={rounds} il={il} stream={a:?} eager={b:?}");
                    }
                }
                let ve = e.dtoh_u32(&vam_d)?;
                if vs[..kr + 1] != ve[..kr + 1] {
                    eprintln!(
                        "[vcheck] DIVERGE round={rounds} pos0={pos0} stream={:?} eager={:?}",
                        &vs[..kr + 1],
                        &ve[..kr + 1]
                    );
                } else {
                    eprintln!("[vcheck] match round={rounds} pos0={pos0}");
                }
            }
            e.u32_pack2(&batch_d, 1, kr, &vam_d, kr + 1, &mut packed)?;
            let host = e.dtoh_u32(&packed)?; // the round's ONE sync
            let k = kr;
            let dtoks: Vec<u32> = host[..k].to_vec();
            let vam: Vec<u32> = host[k..2 * k + 1].to_vec();
            // longest accepted prefix: d_i accepted iff d_i == argmax(verify[i-1])
            // (trimmed heads: batch_d slots were d2t-translated in the draft loop, so dtoks
            // are full-vocab ids here — the 2026-07-10 async rewrite silently dropped this
            // and the trim probes read accept=0.000 through it.)
            let mut m = 0usize;
            while m < k {
                if dtoks[m] == vam[m] {
                    m += 1;
                } else {
                    break;
                }
            }
            prev_full = m == k; // feeds the self-keyed in-round cut (miss → next round cuts)
            if std::env::var("MEMRA_DEBUG_SPEC").as_deref() == Ok("1") {
                let l0 = cache
                    .kv
                    .iter()
                    .flatten()
                    .next()
                    .map(|kv| kv.len)
                    .unwrap_or(0);
                let hh = e.dtoh(&h)?;
                let hn: f32 = hh.iter().map(|x| x * x).sum::<f32>().sqrt();
                eprintln!(
                    "[round {rounds}] pos0={pos0} post_pos={} kv0_len={l0} last={last} dtoks={dtoks:?} vam={vam:?} m={m} |h_in|={hn:.3}",
                    cache.pos
                );
            }
            accepted += m;
            for j in 0..k.min(16) {
                pos_att[j] += 1;
                if j < m {
                    pos_acc[j] += 1;
                }
            }
            // emit last + accepted drafts; the correction token comes from verify row m.
            out.push(last);
            if eos.contains(&last) {
                break 'outer;
            }
            for &dt in &dtoks[..m] {
                out.push(dt);
                if eos.contains(&dt) {
                    break 'outer;
                }
                if out.len() >= max_new {
                    break 'outer;
                }
            }
            let next = vam[m];
            // roll back rejected rows: batch appended k+1 rows; keep m+1 (positions of
            // last + accepted drafts). SWA layers cap t_kv by the window view, so a plain
            // len rewind is safe for every layer.
            let keep = m + 1;
            for kvl in cache.kv.iter_mut().flatten() {
                kvl.len -= (k + 1) - keep;
                // keep len_d in lockstep: the drafter's device-len attention arms read it
                // (the gemma round appends via the HOST-len path, which doesn't maintain
                // the counter — stale len_d gutted acceptance to 0.059 on the dc probe).
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            cache.pos -= (k + 1) - keep;
            // h for the next round = main hidden at the LAST KEPT position (verify row m).
            let hv = e.view(&vh, (k + 1) * n_embd);
            let row = hv.slice(m * n_embd..(m + 1) * n_embd);
            let mut hrow = e.uninit(n_embd)?;
            e.copy_view_into(&mut hrow, 0, &row, n_embd)?;
            h = hrow;
            last = next;
            // Adaptive trim, learn point 2: ALL verify argmaxes — vam[m] is the emitted
            // correction (the only emitted token that can sit outside the trim set; accepted
            // drafts are trim members by construction), and vam[i>m] are main-model
            // predictions for positions never reached this round: next round usually wants
            // exactly those tokens, so learning them here lets the draft propose them
            // BEFORE any miss is paid (prose escapes are first-occurrence-dominated —
            // corrections-only learning measured +0.5 acceptance pts, jsonl 2026-07-19).
            trim_adapt_learn(e, d, &vam)?;
            if adapt {
                let fl_now = floor_at(cache.pos);
                kc = (m + 1).clamp(fl_now.min(k_cap), k_cap);
                // confidence cut (MEMRA_SPEC_PMIN > 0): next round drafts no deeper than one
                // past the first low-confidence draft of THIS round (llama's p-min class,
                // one round late — the zero-sync enqueue stays intact). One extra tiny dtoh.
                if pmin > 0.0 {
                    let ph = e.dtoh(&p_d)?;
                    if let Some(fl) = ph[..kr].iter().position(|&p| p < pmin) {
                        kc = kc.min((fl + 1).max(fl_now.min(k_cap)));
                    }
                }
            }
        }
        eprintln!(
            "[gemma-spec] rounds={rounds} drafted={drafted} accepted={accepted}                    accept-rate={:.3} tok/round={:.2}",
            accepted as f64 / drafted.max(1) as f64,
            out.len() as f64 / rounds.max(1) as f64
        );
        if let Some((used, budget)) = d.trim_adapt_stats() {
            eprintln!("[trim-adapt] {used}/{budget} spare slots learned");
            match d.trim_adapt_save() {
                Ok(n) if n > 0 => {
                    eprintln!("[trim-adapt] {n} new ids appended to the .learned sidecar")
                }
                Ok(_) => {}
                Err(err) => eprintln!("[trim-adapt] sidecar save failed: {err}"),
            }
        }
        if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
            let hist: Vec<String> = (0..16)
                .filter(|&j| pos_att[j] > 0)
                .map(|j| format!("p{j}:{}/{}", pos_acc[j], pos_att[j]))
                .collect();
            eprintln!("[gemma-spec] per-position accept: {}", hist.join(" "));
        }
        Ok(out)
    }
}

/// BURST-SCOPED gemma4 spec session (lane/gemma-batched stage 1, 2026-08-16): the serve
/// twin of `generate_spec_gemma`. That function is GENERATION-scoped — it builds its own
/// cache, primes, loops to completion, and its `break 'outer` exits deliberately skip the
/// final round's rollback/h/pending updates (safe only because the cache dies with the
/// call). A served session must instead stop and RESUME across scheduler ticks, so this
/// type carries the exact cross-round state the eager loop threads through its locals:
///
///   * `cache`  — the trunk Cache; rows = `committed` (prompt + emitted, INCL. overshoot).
///   * `h`      — post-output_norm hidden of the LAST committed row (device; draft seed).
///   * `pending`— the `last` local: the predicted next token. Emitted as the FIRST token
///     of the next round and appended as verify col 0 there; it has NO cache
///     row while parked here (the Q38 `next_pred` convention).
///   * `kc_next`/`prev_full` — the adaptive-depth + self-keyed in-round-cut carries.
///
/// BOUNDARY LAW (the Q38 pending-carry/empty-suffix bug class, banked as gate cases in
/// gemma-spec-session-gate before this was written): a burst NEVER exits mid-round.
/// Every round runs to completion — emission, rollback to the accepted prefix, h/pending
/// update, trim-adapt learn — and only then does the burst-target check run. Overshoot
/// past `target` is committed and returned (the caller clamps VISIBLE emission; state
/// counts every row, exactly like Q38's `SpecSession::committed`). EOS ends the burst at
/// its round boundary with the same complete-state guarantee.
///
/// V1 scope (greedy serve): EAGER round arm only — the round-graph / burst-ring arms are
/// generation-scoped perf doors (their ring/pos-counter state does not checkpoint at
/// round boundaries) and the shipping bench receipts (154.9/176-179, ASSISTANT-ARM-
/// RESULTS.md) were measured on this same eager arm. Dense gemma4 only (E4B refused).
/// Fresh session per request: no prefix reuse, no multi-turn suffix — continuation
/// bursts are always empty-suffix by construction.
pub struct GemmaSpecSession {
    pub cache: Cache,
    /// Every token whose rows the cache holds, in order (prompt + emitted, incl. overshoot).
    pub committed: Vec<u32>,
    h: CudaSlice<f32>,
    pending: u32,
    kc_next: usize,
    prev_full: bool,
    pub prompt_len: usize,
    /// Session-lifetime spec telemetry (rounds / drafted / accepted).
    pub rounds: usize,
    pub drafted: usize,
    pub accepted: usize,
}

impl GemmaSpecSession {
    /// Tokens the session has emitted (committed past the prompt). The pending token is
    /// NOT included — it has no cache row and the next burst emits it first.
    pub fn emitted_len(&self) -> usize {
        self.committed.len() - self.prompt_len
    }
    /// Context capacity of the session's cache (the server's ContextFull guard).
    pub fn cache_max_ctx(&self) -> usize {
        self.cache.max_ctx
    }
    /// DEMOTE HANDOFF (stage-2 seam, gated by the session gate's demote case): hand the
    /// trunk cache to the plain path. The cache rows are exactly `committed` (boundary
    /// law), and the pending token is returned as the plain path's device_next-equivalent
    /// — the plain loop feeds it as its first decode input. The draft side holds no
    /// per-session state (the assistant drafter reads the TRUNK's KV; trim-adapt is
    /// model-lifetime, not session), so dropping self is the whole handoff.
    pub fn into_demoted(self) -> (Cache, u32, Vec<u32>) {
        (self.cache, self.pending, self.committed)
    }
}

impl HybridModel {
    /// Open a burst-scoped gemma spec session: prime the prompt, park the first predicted
    /// token as `pending`. Mirrors `generate_spec_gemma`'s entry verbatim (trim-adapt
    /// learn point 1, the PRIME_MIN_T split, the post-norm h convention).
    pub fn gemma_spec_session_new(
        &self,
        e: &Engine,
        d: &mut GemmaDraft,
        prompt: &[u32],
        max_ctx: usize,
    ) -> Result<GemmaSpecSession, Box<dyn std::error::Error>> {
        if self.is_gemma4_e4b() || !self.uses_gemma_program() {
            return Err(
                "gemma_spec_session_new: dense gemma4 only (E4B keeps its own arms)".into(),
            );
        }
        let n_embd = self.cfg.n_embd as usize;
        let eps = self.cfg.rms_eps;
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;
        trim_adapt_learn(e, d, prompt)?;
        let (pl, h_seed, post_norm) = if prompt.len() >= crate::hybrid_forward::PRIME_MIN_T {
            let (l, hs, _hh) = self.prime_cache(e, prompt, &mut cache, 0)?;
            (l, hs, false)
        } else {
            let n_vocab = self.output.out_features();
            let (lv, hv) = self.gemma4_decode_step_t_h(e, prompt, 0, &mut cache)?;
            let t = prompt.len();
            let last = lv[(t - 1) * n_vocab..t * n_vocab].to_vec();
            let hvv = e.view(&hv, t * n_embd);
            let row = hvv.slice((t - 1) * n_embd..t * n_embd);
            let mut hrow = e.uninit(n_embd)?;
            e.copy_view_into(&mut hrow, 0, &row, n_embd)?;
            (last, hrow, true)
        };
        // drafter h = POST-output_norm hidden; the prime returns PRE-norm h_seed.
        let h = if post_norm {
            h_seed
        } else {
            let mut hh = e.uninit(n_embd)?;
            e.rms_norm(
                &h_seed,
                self.output_norm.float_data(),
                &mut hh,
                n_embd,
                1,
                eps,
            )?;
            hh
        };
        let pending = crate::forward::argmax(&pl) as u32;
        Ok(GemmaSpecSession {
            cache,
            committed: prompt.to_vec(),
            h,
            pending,
            kc_next: usize::MAX, // clamped to the burst's k_cap at entry (one-shot: kc = k_cap)
            prev_full: true,     // round 1: no miss evidence yet
            prompt_len: prompt.len(),
            rounds: 0,
            drafted: 0,
            accepted: 0,
        })
    }

    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18): open a gemma spec
    /// session over a trunk cache the worker already restored from a WHOLE prefix-cache
    /// entry (rows `[0..prefix.len())` == `prefix`, `cache.pos == prefix.len()`), feeding
    /// only the prompt SUFFIX. The assistant drafter holds no per-session KV of its own —
    /// it attends the TRUNK's cache — so the restored rows already ARE the draft state;
    /// the only products a fresh prime supplied were the boundary logits (-> `pending`)
    /// and the post-norm hidden of the last prompt row (-> `h`, the drafter seed), and a
    /// non-empty suffix feed regenerates both. An empty suffix therefore REFUSES: there is
    /// no drafter seed hidden without feeding at least one row (the plain path serves that
    /// shape from the entry's boundary logits, as before).
    ///
    /// PROGRAM CHOICE (the splitiso two-programs law, 0b0ffa13c6): gemma4's monolithic
    /// prime refuses pos > 0, so the suffix rides `gemma4_decode_step_t_h` — the SAME
    /// verify-trunk program every spec round runs and the same arm the cold
    /// sub-PRIME_MIN_T session prime uses (banked byte-identical by
    /// gemma-spec-session-gate). The restored bytes below `prefix.len()` are never
    /// recomputed by construction.
    pub fn gemma_spec_session_from_restored(
        &self,
        e: &Engine,
        d: &mut GemmaDraft,
        mut cache: Cache,
        prefix: &[u32],
        suffix: &[u32],
    ) -> Result<GemmaSpecSession, Box<dyn std::error::Error>> {
        if self.is_gemma4_e4b() || !self.uses_gemma_program() {
            return Err(
                "gemma_spec_session_from_restored: dense gemma4 only (E4B keeps its own arms)"
                    .into(),
            );
        }
        if prefix.is_empty() {
            return Err("gemma_spec_session_from_restored: empty restored prefix".into());
        }
        if suffix.is_empty() {
            return Err(
                "gemma_spec_session_from_restored: empty suffix — the drafter seed \
                        hidden only exists after feeding at least one row (plain path owns \
                        the whole-prompt hit)"
                    .into(),
            );
        }
        if cache.pos != prefix.len() {
            return Err(format!(
                "gemma_spec_session_from_restored: restored cache pos {} != prefix len {}",
                cache.pos,
                prefix.len(),
            )
            .into());
        }
        let n_embd = self.cfg.n_embd as usize;
        let n_vocab = self.output.out_features();
        // trim-adapt learning is model-lifetime (not session state); feed the full logical
        // prompt so restored traffic teaches the head trim exactly what cold traffic does.
        let full: Vec<u32> = prefix.iter().chain(suffix.iter()).copied().collect();
        trim_adapt_learn(e, d, &full)?;
        let base = cache.pos;
        let (lv, hv) = self.gemma4_decode_step_t_h(e, suffix, base, &mut cache)?;
        let t = suffix.len();
        let last = lv[(t - 1) * n_vocab..t * n_vocab].to_vec();
        // gemma4_decode_step_t_h returns POST-output_norm hiddens (the drafter's h
        // convention — same arm gemma_spec_session_new uses below PRIME_MIN_T).
        let hvv = e.view(&hv, t * n_embd);
        let row = hvv.slice((t - 1) * n_embd..t * n_embd);
        let mut h = e.uninit(n_embd)?;
        e.copy_view_into(&mut h, 0, &row, n_embd)?;
        let pending = crate::forward::argmax(&last) as u32;
        let prompt_len = full.len();
        Ok(GemmaSpecSession {
            cache,
            committed: full,
            h,
            pending,
            kc_next: usize::MAX,
            prev_full: true,
            prompt_len,
            rounds: 0,
            drafted: 0,
            accepted: 0,
        })
    }

    /// One serve burst: run complete spec rounds until >= `target` NEW tokens have been
    /// emitted this burst (overshoot committed and returned) or EOS lands. Returns
    /// (tokens emitted this burst in order, drafted, accepted). The round body is the
    /// EAGER arm of `generate_spec_gemma`, kept behaviorally identical under default env
    /// (adapt/floor/pmin/in-round-cut logic verbatim) — gemma-spec-session-gate enforces
    /// byte-equality of the emitted stream against the one-shot at every burst width.
    pub fn gemma_spec_session_burst(
        &self,
        e: &Engine,
        d: &mut GemmaDraft,
        sess: &mut GemmaSpecSession,
        target: usize,
        k: usize,
        eos: &[u32],
    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
        let n_embd = self.cfg.n_embd as usize;
        if target == 0 {
            return Ok((Vec::new(), 0, 0));
        }
        let pmin: f32 = std::env::var("MEMRA_SPEC_PMIN")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0.0);
        let pmin_ir_env: Option<f32> = std::env::var("MEMRA_SPEC_PMIN_INROUND")
            .ok()
            .and_then(|v| v.parse().ok());
        const PMIN_IR_DEFAULT: f32 = 0.7;
        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() != Ok("0");
        let adapt_floor_default: usize = if self.cfg.n_embd >= 3500 {
            4
        } else if self.cfg.n_embd >= 2500 {
            2
        } else {
            1
        };
        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
            .ok()
            .and_then(|v| v.parse().ok());
        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(1024);
        let floor_at = |pos: usize| -> usize {
            if adapt_floor_env.is_some() || pos < floor_ctx {
                adapt_floor
            } else if adapt_floor >= 4 {
                1
            } else {
                adapt_floor
            }
        };
        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(7);
        let k_cap = k.min(cap_max).max(1);
        let mut kc = sess.kc_next.min(k_cap);
        let mut prev_full = sess.prev_full;

        // per-burst device scratch (the one-shot allocates these per generation; per-burst
        // re-allocation is micro against a >= (K+1)-token round).
        let mut batch_d = e.stream().alloc_zeros::<u32>(k_cap + 1)?;
        let mut packed = e.stream().alloc_zeros::<u32>(2 * k_cap + 1)?;
        let mut p_d = e.stream().alloc_zeros::<f32>(k_cap.max(1))?;
        let mut pos_slots: Vec<CudaSlice<i32>> = (0..k_cap.max(1))
            .map(|_| e.htod_i32(&[0]))
            .collect::<Result<_, _>>()?;
        let mut g_seed = e.zeros(n_embd)?;
        // len_d lockstep at burst entry: the drafter's device-len arms read it, and the
        // previous burst's rollback set it — a fresh session's prime went host-len.
        for kvl in sess.cache.kv.iter_mut().flatten() {
            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
        }

        let mut burst_out: Vec<u32> = Vec::with_capacity(target + k_cap + 1);
        let (mut drafted, mut accepted) = (0usize, 0usize);
        let mut ended = false;
        while burst_out.len() < target && !ended {
            let mut kr = if adapt { kc } else { k_cap };
            let dc_bucket: Option<usize> = {
                static DC: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
                if *DC.get_or_init(|| std::env::var("MEMRA_GEMMA_DRAFT_DC").as_deref() != Ok("0")) {
                    let ml = sess
                        .cache
                        .kv
                        .iter()
                        .flatten()
                        .map(|kv| kv.len)
                        .max()
                        .unwrap_or(1);
                    Some((ml + k_cap + 2).next_power_of_two().max(512))
                } else {
                    None
                }
            };
            e.u32_set_k(&mut batch_d, sess.pending, 0)?;
            e.copy_into(&mut g_seed, 0, &sess.h, n_embd)?;
            for (j, slot) in pos_slots.iter_mut().take(kr).enumerate() {
                e.set_i32_one(slot, (sess.cache.pos + j) as i32)?;
            }
            let ir_now = match pmin_ir_env {
                Some(p) => p,
                None if sess.cache.pos >= floor_ctx && !prev_full => PMIN_IR_DEFAULT,
                None => 0.0,
            };
            // draft chain (the one-shot's run_chain, eager): reads g_seed, seeds batch_d.
            {
                let mut hc = e.uninit(n_embd)?;
                e.copy_into(&mut hc, 0, &g_seed, n_embd)?;
                let mut j = 0usize;
                while j < kr {
                    let tv = batch_d.slice(j..j + 1);
                    let (hn, h_next) = self.gemma4_draft_trunk_dev(
                        e,
                        d,
                        &tv,
                        &hc,
                        &pos_slots[j],
                        &sess.cache,
                        dc_bucket,
                    )?;
                    let ld = e.matmul(&d.head, &hn, 1)?;
                    e.argmax_token_device_col(&ld, 0, d.head.out_features(), &mut batch_d, j + 1)?;
                    if pmin > 0.0 || ir_now > 0.0 {
                        e.prob_of_token_device_col(
                            &ld,
                            &batch_d,
                            j + 1,
                            &mut p_d,
                            j,
                            d.head.out_features(),
                        )?;
                    }
                    if let Some(map) = &d.d2t_dev {
                        e.u32_map_k(&mut batch_d, map, j + 1)?;
                    }
                    hc = h_next;
                    if ir_now > 0.0 && j + 1 < kr {
                        let ph = e.dtoh(&p_d)?;
                        if ph[j] < ir_now {
                            kr = j + 1;
                            break;
                        }
                    }
                    j += 1;
                }
            }
            drafted += kr;
            sess.rounds += 1;
            let pos0 = sess.cache.pos;
            let (vam_d, vh) =
                self.gemma4_decode_step_t_am_dev(e, &batch_d, kr + 1, pos0, &mut sess.cache)?;
            e.u32_pack2(&batch_d, 1, kr, &vam_d, kr + 1, &mut packed)?;
            let host = e.dtoh_u32(&packed)?; // the round's ONE sync
            let dtoks: Vec<u32> = host[..kr].to_vec();
            let vam: Vec<u32> = host[kr..2 * kr + 1].to_vec();
            let mut m = 0usize;
            while m < kr {
                if dtoks[m] == vam[m] {
                    m += 1;
                } else {
                    break;
                }
            }
            prev_full = m == kr;
            accepted += m;
            // ---- ROUND COMPLETES UNCONDITIONALLY (the boundary law) ----
            // rollback rejected rows FIRST, then emit — an EOS mid-emission must still
            // leave cache rows == committed tokens.
            let keep = m + 1;
            for kvl in sess.cache.kv.iter_mut().flatten() {
                kvl.len -= (kr + 1) - keep;
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            sess.cache.pos -= (kr + 1) - keep;
            // h for the next round = main hidden at the LAST KEPT position (verify row m).
            let hv2 = e.view(&vh, (kr + 1) * n_embd);
            let row = hv2.slice(m * n_embd..(m + 1) * n_embd);
            let mut hrow = e.uninit(n_embd)?;
            e.copy_view_into(&mut hrow, 0, &row, n_embd)?;
            sess.h = hrow;
            // emit: pending + accepted drafts. VISIBLE emission stops at the first EOS
            // (the one-shot's exact stream); COMMIT accounting continues — every kept row
            // must appear in `committed` or the cache-rows == committed invariant breaks.
            sess.committed.push(sess.pending);
            burst_out.push(sess.pending);
            if eos.contains(&sess.pending) {
                ended = true;
            }
            for &dt in &dtoks[..m] {
                sess.committed.push(dt);
                if !ended {
                    burst_out.push(dt);
                    if eos.contains(&dt) {
                        ended = true;
                    }
                }
            }
            sess.pending = vam[m];
            trim_adapt_learn(e, d, &vam)?;
            if adapt {
                let fl_now = floor_at(sess.cache.pos);
                kc = (m + 1).clamp(fl_now.min(k_cap), k_cap);
                if pmin > 0.0 {
                    let ph = e.dtoh(&p_d)?;
                    if let Some(fl) = ph[..kr].iter().position(|&p| p < pmin) {
                        kc = kc.min((fl + 1).max(fl_now.min(k_cap)));
                    }
                }
            }
        }
        sess.kc_next = kc;
        sess.prev_full = prev_full;
        sess.drafted += drafted;
        sess.accepted += accepted;
        Ok((burst_out, drafted, accepted))
    }
}

impl HybridModel {
    /// PLAIN-DECODE CUDA-GRAPH loop (gemma4, greedy): one captured verify-trunk step
    /// (t=1, device tokens/pos/lens) replayed per token — the launch-gap eraser the
    /// decode decomposition demanded (2026-07-23: ~2.3ms/token idle at 128 launches).
    /// Self-feeding: argmax -> tok_d -> next embed; counters advance in-graph via
    /// spec_rollback_stream(base=1, acc=0). Tokens land in a device ring; ONE host sync
    /// per drain window. Captures are keyed on the (rung, window-side, f512-side) regime
    /// (the round-graph hint law); regime-crossing stretches run the same body eagerly.
    /// Caller guarantees: gemma4, greedy, shared_kv_layers == 0, prompt already primed
    /// (cache.pos = prompt len, host kvl.len mirrors set).
    pub fn gemma4_generate_plain_graph(
        &self,
        e: &Engine,
        cache: &mut Cache,
        last: u32,
        max_new: usize,
        eos: &[u32],
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        const RING: usize = 64;
        const DRAIN: usize = 32; // replays per host sync
        let win_main = self
            .cfg
            .gemma4
            .as_ref()
            .map(|g| g.sliding_window as usize)
            .unwrap_or(0);
        let n_rows = cache.kv.len() + 1;

        let was_tracking = e.ctx().is_event_tracking();
        if was_tracking {
            unsafe {
                e.ctx().disable_event_tracking();
            }
        }
        let r = self
            .gemma4_plain_graph_inner(e, cache, last, max_new, eos, RING, DRAIN, win_main, n_rows);
        if was_tracking {
            unsafe {
                e.ctx().enable_event_tracking();
            }
        }
        r
    }

    #[allow(clippy::too_many_arguments)]
    #[allow(clippy::map_entry)] // allow: the init body is fallible (`?`); Entry::or_insert_with cannot propagate errors
    fn gemma4_plain_graph_inner(
        &self,
        e: &Engine,
        cache: &mut Cache,
        last: u32,
        max_new: usize,
        eos: &[u32],
        ring_cap: usize,
        drain: usize,
        win_main: usize,
        n_rows: usize,
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        let mut scr = self.verify_stream_scratch(e, 1)?;
        let mut tok_d = e.stream().alloc_zeros::<u32>(1)?;
        e.u32_set_k(&mut tok_d, last, 0)?;
        let pos_ctr = e.htod_i32(&[cache.pos as i32])?;
        let mut pos_start_d = e.htod_i32(&[cache.pos as i32])?;
        let acc0 = e.stream().alloc_zeros::<u32>(2)?; // acc[0] = 0 -> counters +1
        let mut ring = e.stream().alloc_zeros::<u32>(ring_cap)?;
        let ptrs = crate::round_stream::kv_len_ptr_table(e, cache, Some(&pos_ctr))?;
        for kvl in cache.kv.iter_mut().flatten() {
            e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
        }
        let ring_base = cache.pos; // baked into every capture

        #[allow(clippy::type_complexity)]
        // allow: one-shot composite type; naming it would hide the shape that matters at the call site
        let mut graphs: std::collections::HashMap<
            (usize, bool, bool),
            (
                cudarc::driver::CudaGraph,
                Vec<Box<dyn std::any::Any + Send>>,
            ),
        > = Default::default();

        let mut out: Vec<u32> = Vec::with_capacity(max_new);
        let mut drained = 0usize; // tokens read off the ring

        // hint law (round-graph): the arm-gating bound must sit on the SAME side of every
        // crossover as the live lengths this capture serves, with the arms' own margins.
        let hint_for = |pos: usize| -> usize {
            if pos > win_main {
                pos + drain + 2
            } else if pos + 1 >= crate::fa512_min_tkv() {
                win_main.saturating_sub(2)
            } else {
                crate::fa512_min_tkv().saturating_sub(5)
            }
        };
        let regime_key = |pos: usize| -> (usize, bool, bool) {
            let rung = (pos + drain + 2).next_power_of_two().max(512);
            (rung, pos > win_main, pos + 1 >= crate::fa512_min_tkv())
        };
        // the whole [pos, pos+n) stretch must share one regime for a captured replay run.
        let stable_for = |pos: usize, n: usize| -> bool {
            regime_key(pos) == regime_key(pos + n)
                && (pos > win_main || pos + n + 2 < win_main)
                && (pos + 1 >= crate::fa512_min_tkv() || pos + n + 2 < crate::fa512_min_tkv())
        };

        while out.len() < max_new {
            let pos = cache.pos;
            let hint = hint_for(pos);
            let scr_ptr: *mut crate::hybrid_forward::VerifyStreamScratch = &mut scr;
            let cache_ptr: *mut Cache = cache as *mut Cache;
            let tok_ptr: *mut CudaSlice<u32> = &mut tok_d;
            let ring_ptr: *mut CudaSlice<u32> = &mut ring;
            let start_ptr: *mut CudaSlice<i32> = &mut pos_start_d;
            let step = |e: &Engine| -> Result<(), Box<dyn std::error::Error>> {
                // SAFETY: single-threaded body; raw pointers alias the outer &mut only here.
                let (scr, cache, tok_d, ring, pos_start_d) = unsafe {
                    (
                        &mut *scr_ptr,
                        &mut *cache_ptr,
                        &mut *tok_ptr,
                        &mut *ring_ptr,
                        &mut *start_ptr,
                    )
                };
                e.i32_copy_add(&pos_ctr, pos_start_d, 0)?;
                let (vam, _hn) =
                    self.gemma4_verify_t_am_stream(e, tok_d, 1, &pos_ctr, hint, cache, scr)?;
                e.u32_copy(&vam, tok_d)?;
                e.plain_tok_ring(&vam, pos_start_d, ring_base, ring)?;
                e.spec_rollback_stream(&ptrs, pos_start_d, &acc0, 1, n_rows)?;
                Ok(())
            };

            let n_left = max_new - out.len();
            let burst = drain.min(n_left);
            // MEMRA_G4PLAIN_EAGER=1: run the body eagerly every step (no capture/replay) —
            // splits "body semantics wrong" from "replay mechanics wrong" (round-graph law).
            let force_eager = std::env::var("MEMRA_G4PLAIN_EAGER").as_deref() == Ok("1");
            let steps_done = if !force_eager && burst >= 4 && stable_for(pos, burst + 3) {
                let key = regime_key(pos);
                if !graphs.contains_key(&key) {
                    // capture cost = 3 SERVED steps (2 warmups + the captured run itself):
                    // the loop is self-feeding, so they are real tokens in the ring.
                    let g = e.capture_graph_retained(step)?;
                    graphs.insert(key, g);
                    3
                } else {
                    let (g, _keep) = graphs.get(&key).unwrap();
                    for _ in 0..burst {
                        g.launch()?;
                    }
                    burst
                }
            } else {
                step(e)?; // eager fallback (same body)
                1
            };

            // host mirrors + drain
            cache.pos += steps_done;
            for kvl in cache.kv.iter_mut().flatten() {
                kvl.len = cache.pos;
            }
            e.stream().synchronize()?;
            let ringh = e.dtoh_u32(&ring)?;
            let total = cache.pos - ring_base;
            while drained < total && out.len() < max_new {
                let t = ringh[drained % ring_cap];
                out.push(t);
                drained += 1;
                if eos.contains(&t) {
                    return Ok(out);
                }
            }
        }
        Ok(out)
    }
}