memra_engine/

dflash.rs

//! DFlash block-diffusion drafter (DFLASH-BRINGUP-PLAN.md, 2026-07-13).
//!
//! 5-layer qwen3-class mini-transformer that drafts a 16-token block in ONE non-causal
//! forward, conditioned on the TARGET's hidden states at 6 tapped layers (concatenated
//! through `fc` + `hidden_norm`). No embed / lm_head of its own — the round reuses the
//! target's. Reference: z-lab/dflash `dflash/model.py` (semantics frozen in the plan doc);
//! oracle: tools/dflash_oracle.py -> /data/cache/dflash-oracle.npz.
//!
//! FIRST LIGHT = f32-resident weights + fresh full-context forward (no draft KV cache) —
//! correctness vs the oracle, then the cache/quant/window arms land measurement-gated.

use crate::Engine;
use crate::model::GpuTensor;
use cudarc::driver::CudaSlice;

pub struct DflashCfg {
    pub hidden: usize,                // 5376
    pub n_head: usize,                // 64
    pub n_kv: usize,                  // 8
    pub head_dim: usize,              // 128
    pub n_ff: usize,                  // 10752
    pub n_layer: usize,               // 5
    pub eps: f32,                     // 1e-6
    pub rope_theta: f32,              // 1e6
    pub block_size: usize,            // 16
    pub mask_token_id: u32,           // 4
    pub target_layer_ids: Vec<usize>, // [1,12,23,35,46,57]
    pub sliding_window: usize,        // 2048
    /// true = sliding_attention for that layer (4x true + 1x false on the 31B draft).
    pub layer_sliding: Vec<bool>,
}

pub struct DflashLayer {
    pub wq: GpuTensor,           // [nh*hd, hidden] row-major (out_f rows)
    pub wk: GpuTensor,           // [nkv*hd, hidden]
    pub wv: GpuTensor,           // [nkv*hd, hidden]
    pub wo: GpuTensor,           // [hidden, nh*hd]
    pub w_gate: GpuTensor,       // [n_ff, hidden]
    pub w_up: GpuTensor,         // [n_ff, hidden]
    pub w_down: GpuTensor,       // [hidden, n_ff]
    pub ln_in: CudaSlice<f32>,   // [hidden]
    pub ln_post: CudaSlice<f32>, // [hidden]
    pub q_norm: CudaSlice<f32>,  // [hd]
    pub k_norm: CudaSlice<f32>,  // [hd]
}

pub struct DflashDraft {
    pub cfg: DflashCfg,
    pub layers: Vec<DflashLayer>,
    pub fc: GpuTensor,               // [hidden, n_taps*hidden]
    pub hidden_norm: CudaSlice<f32>, // [hidden]
    pub norm: CudaSlice<f32>,        // [hidden]
    /// DSpark semi-AR markov head (present in the repo-root checkpoint variant):
    /// draft logits at position k get + W2(W1[prev_realized_token]) — left-to-right
    /// within the block (the patch's _markov_semiar_sample_block semantics, greedy).
    /// w1 = raw bf16 [V, rank] (row-gathered by device token id); w2 = q8_0 [rank->V].
    pub markov: Option<MarkovHead>,
    /// DSpark accept-rate head (trained with confidence loss; the reference serving
    /// loop — SpecForge spec_generate — never consumes it, so serving ignores it too.
    /// Loaded for census completeness + the parity gate; host-resident (5k floats).
    pub confidence: Option<ConfidenceHead>,
    /// YaRN rope (q38 arm-a inherits the target's rope_parameters: rope_type yarn,
    /// factor 32, original 8192, beta 32/1). ff = per-dim divisors for rope_neox_ff
    /// (effective inv_freq_j = base^(-2j/d)/ff[j] = the HF-yarn remapped frequency,
    /// verified vs Qwen3RotaryEmbedding to 1.6e-7), mscale = attention_scaling
    /// (0.1*ln(factor)+1) applied to q/k post-rope — cos/sin scaling distributes onto
    /// the rotated vector exactly. None = plain rope (gemma/z-lab drafters).
    pub rope_yarn: Option<(CudaSlice<f32>, f32)>,
}

/// AcceptRatePredictor: raw linear proj over [hidden ; markov_prev_embedding(rank)]
/// (with_markov=true on the q38 arm-a export) — output is the PRE-sigmoid scalar.
pub struct ConfidenceHead {
    pub w: Vec<f32>, // [in_dim]
    pub b: f32,
    pub in_dim: usize,
    pub with_markov: bool,
}

pub struct MarkovHead {
    pub w1_bf16: CudaSlice<u8>, // [V, rank] bf16 raw
    pub w2: GpuTensor,          // [rank -> V] q8_0
    pub rank: usize,
    pub vocab: usize,
}

fn bf16_to_f32(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(2)
        .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
        .collect()
}

/// Host q8_0 encode (ggml block layout: [d f16][32 x i8] = 34B/32 vals). The drafter's
/// weights ride the dp4a fast path at 1.6GB resident (bf16 3.1GB + the 31B trunk OOM'd
/// 24GB; f32 6.2GB worse). Drafter quantization moves ACCEPTANCE only — verify exactness
/// is structural.
fn encode_q8_0(vals: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(vals.len() / 32 * 34);
    for blk in vals.chunks_exact(32) {
        let amax = blk.iter().fold(0f32, |a, v| a.max(v.abs()));
        let d = amax / 127.0;
        let id = if d > 0.0 { 1.0 / d } else { 0.0 };
        let dh = half_from_f32(d);
        out.extend_from_slice(&dh.to_le_bytes());
        for &v in blk {
            out.push(((v * id).round().clamp(-127.0, 127.0)) as i8 as u8);
        }
    }
    out
}

/// Host q4_0 encode (ggml: [d f16][16B packed nibbles] = 18B/32 vals; q = round(v/d)+8,
/// d = amax/-7 sign trick NOT used — plain amax/7? ggml uses d = max/-8 .. follow ggml:
/// d = amax / -8 when the max is negative-dominant; reference quantize_row_q4_0: d =
/// max(|v|)/-8 signed-max form). Implemented to match ggml quantize_row_q4_0_ref.
fn encode_q4_0(vals: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(vals.len() / 32 * 18);
    for blk in vals.chunks_exact(32) {
        // ggml ref: pick the value with the LARGEST |v| (keeping sign), d = that / -8
        let mut amax = 0f32;
        let mut mx = 0f32;
        for &v in blk {
            if v.abs() > amax {
                amax = v.abs();
                mx = v;
            }
        }
        let d = mx / -8.0;
        let id = if d != 0.0 { 1.0 / d } else { 0.0 };
        out.extend_from_slice(&half_from_f32(d).to_le_bytes());
        for j in 0..16 {
            let x0 = (blk[j] * id + 8.5).clamp(0.0, 15.0) as u8;
            let x1 = (blk[j + 16] * id + 8.5).clamp(0.0, 15.0) as u8;
            out.push(x0 | (x1 << 4));
        }
    }
    out
}

fn half_from_f32(v: f32) -> u16 {
    // f32 -> IEEE f16 (round-to-nearest-even; range of q8_0 d values is tame)
    let b = v.to_bits();
    let sign = ((b >> 16) & 0x8000) as u16;
    let exp = ((b >> 23) & 0xff) as i32 - 127 + 15;
    let man = b & 0x7fffff;
    if exp <= 0 {
        return sign;
    } // flush tiny d to zero
    if exp >= 31 {
        return sign | 0x7c00;
    } // inf (unreachable for sane d)
    let mut h = sign | ((exp as u16) << 10) | ((man >> 13) as u16);
    // round to nearest even on the truncated 13 bits
    let rem = man & 0x1fff;
    if rem > 0x1000 || (rem == 0x1000 && (h & 1) == 1) {
        h += 1;
    }
    h
}

impl DflashDraft {
    /// Load the backbone-only checkpoint dir (config.json + model.safetensors, bf16).
    /// Config scalars ride a minimal extractor (no json dep in-tree — HfConfig precedent).
    pub fn load(e: &Engine, dir: &std::path::Path) -> Result<Self, Box<dyn std::error::Error>> {
        let txt = std::fs::read_to_string(dir.join("config.json"))?;
        fn num(txt: &str, key: &str) -> Option<f64> {
            let i = txt.find(&format!("\"{key}\""))?;
            let rest = &txt[i..];
            let colon = rest.find(':')?;
            let val: String = rest[colon + 1..]
                .trim_start()
                .chars()
                .take_while(|c| {
                    c.is_ascii_digit()
                        || *c == '.'
                        || *c == '-'
                        || *c == 'e'
                        || *c == 'E'
                        || *c == '+'
                })
                .collect();
            val.parse().ok()
        }
        fn num_list(txt: &str, key: &str) -> Vec<usize> {
            let Some(i) = txt.find(&format!("\"{key}\"")) else {
                return Vec::new();
            };
            let rest = &txt[i..];
            let (Some(a), Some(b)) = (rest.find('['), rest.find(']')) else {
                return Vec::new();
            };
            rest[a + 1..b]
                .split(',')
                .filter_map(|s| s.trim().parse().ok())
                .collect()
        }
        let g = |k: &str| num(&txt, k).unwrap_or_else(|| panic!("config missing {k}")) as usize;
        // layer_types order: count entries, mark sliding ones
        let layer_sliding: Vec<bool> = {
            let i = txt.find("\"layer_types\"").expect("layer_types");
            let rest = &txt[i..];
            let (a, b) = (rest.find('[').unwrap(), rest.find(']').unwrap());
            rest[a + 1..b]
                .split(',')
                .map(|s| s.contains("sliding_attention"))
                .collect()
        };
        // sliding_window is null on all-full-attention exports (q38 arm-a); the window
        // only constrains rounds when a sliding layer exists (reference: resolve_dflash_
        // attention_layout returns None when no layer slides).
        let sliding_window = if layer_sliding.iter().any(|&s| s) {
            g("sliding_window")
        } else {
            num(&txt, "sliding_window")
                .map(|v| v as usize)
                .unwrap_or(usize::MAX)
        };
        let cfg = DflashCfg {
            hidden: g("hidden_size"),
            n_head: g("num_attention_heads"),
            n_kv: g("num_key_value_heads"),
            head_dim: g("head_dim"),
            n_ff: g("intermediate_size"),
            n_layer: g("num_hidden_layers"),
            eps: num(&txt, "rms_norm_eps").expect("rms_norm_eps") as f32,
            rope_theta: num(&txt, "rope_theta").expect("rope_theta") as f32,
            block_size: g("block_size"),
            mask_token_id: g("mask_token_id") as u32,
            target_layer_ids: num_list(&txt, "target_layer_ids"),
            sliding_window,
            layer_sliding,
        };
        let st = memra_gguf::safetensors::StModel::open(&dir.join("model.safetensors"))?;
        // 1D norm weights ride raw slices; 2D matmul weights ride GpuTensor::Float
        // (cuBLASLt f32 arm — the Stage-A numeric class, right for oracle parity).
        let up = |name: &str| -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
            let (_info, bytes) = st
                .raw(name)
                .ok_or_else(|| format!("missing tensor {name}"))?;
            Ok(e.htod(&bf16_to_f32(bytes))?)
        };
        // Precision policy (MEMRA_DFLASH_PREC seam): "q8" = all q8_0 (1.6GB, default);
        // "mixed" = bf16 attn+fc (the ctx-conditioning path) + q8_0 ffn (~2.2GB — fits the
        // ~2.8GB headroom beside the 31B trunk); "bf16" = all bf16 (parity runs, no target).
        let prec = std::env::var("MEMRA_DFLASH_PREC").unwrap_or_else(|_| "q8".into());
        let upw = |name: &str| -> Result<GpuTensor, Box<dyn std::error::Error>> {
            let (info, bytes) = st
                .raw(name)
                .ok_or_else(|| format!("missing tensor {name}"))?;
            let shape = info.ne(); // ggml order: ne[0]=in_f, ne[1]=out_f
            let in_f = shape[0] as usize;
            let is_ffn = name.contains(".mlp.");
            let bf16 = prec == "bf16"
                || (prec == "mixed" && !is_ffn)
                || (prec == "fc" && name == "fc.weight");
            if bf16 {
                return Ok(GpuTensor::FloatBf16 {
                    data: e.upload_u8(bytes)?,
                    ne: shape.to_vec(),
                });
            }
            let f32s = bf16_to_f32(bytes);
            if prec == "q4" {
                let q = encode_q4_0(&f32s);
                return Ok(GpuTensor::Quant {
                    bytes: e.upload_u8(&q)?,
                    qtype: crate::QT_Q4_0,
                    row_bytes: in_f / 32 * 18,
                    ne: shape.to_vec(),
                    scale: 1.0,
                    rp: false,
                    #[cfg(memra_cutlass)]
                    cutlass: None,
                    fp8: None,
                    blk: None,
                    rp4: None,
                    f16: None,
                });
            }
            let q = encode_q8_0(&f32s);
            Ok(GpuTensor::Quant {
                bytes: e.upload_u8(&q)?,
                qtype: crate::QT_Q8_0,
                row_bytes: in_f / 32 * 34,
                ne: shape.to_vec(),
                scale: 1.0,
                rp: false,
                #[cfg(memra_cutlass)]
                cutlass: None,
                fp8: None,
                blk: None,
                rp4: None,
                f16: None,
            })
        };
        let mut layers = Vec::with_capacity(cfg.n_layer);
        for i in 0..cfg.n_layer {
            let p = |s: &str| format!("layers.{i}.{s}");
            layers.push(DflashLayer {
                wq: upw(&p("self_attn.q_proj.weight"))?,
                wk: upw(&p("self_attn.k_proj.weight"))?,
                wv: upw(&p("self_attn.v_proj.weight"))?,
                wo: upw(&p("self_attn.o_proj.weight"))?,
                w_gate: upw(&p("mlp.gate_proj.weight"))?,
                w_up: upw(&p("mlp.up_proj.weight"))?,
                w_down: upw(&p("mlp.down_proj.weight"))?,
                ln_in: up(&p("input_layernorm.weight"))?,
                ln_post: up(&p("post_attention_layernorm.weight"))?,
                q_norm: up(&p("self_attn.q_norm.weight"))?,
                k_norm: up(&p("self_attn.k_norm.weight"))?,
            });
        }
        let markov = if let Some((info, bytes)) = st.raw("markov_head.markov_w1.weight") {
            let sh = info.ne(); // [rank, vocab] in ggml order (safetensors [V, rank] reversed)
            let (rank, vocab) = (sh[0] as usize, sh[1] as usize);
            let (i2, b2) = st
                .raw("markov_head.markov_w2.weight")
                .ok_or("markov_w2 missing beside markov_w1")?;
            // w2 follows the precision seam: bf16 for parity runs (the q8_0 encode is a
            // serving-size choice and would put quant error inside the markov-logits gate),
            // q8_0 otherwise (acceptance-only impact, like the trunk weights).
            let w2 = if prec == "bf16" {
                GpuTensor::FloatBf16 {
                    data: e.upload_u8(b2)?,
                    ne: i2.ne().to_vec(),
                }
            } else {
                let w2f = bf16_to_f32(b2);
                let w2q = encode_q8_0(&w2f);
                GpuTensor::Quant {
                    bytes: e.upload_u8(&w2q)?,
                    qtype: crate::QT_Q8_0,
                    row_bytes: rank / 32 * 34,
                    ne: vec![rank as u64, vocab as u64],
                    scale: 1.0,
                    rp: false,
                    #[cfg(memra_cutlass)]
                    cutlass: None,
                    fp8: None,
                    blk: None,
                    rp4: None,
                    f16: None,
                }
            };
            Some(MarkovHead {
                w1_bf16: e.upload_u8(bytes)?,
                w2,
                rank,
                vocab,
            })
        } else {
            None
        };
        let confidence = if let Some((info, bytes)) = st.raw("confidence_head.proj.weight") {
            let sh = info.ne(); // ggml order: ne[0]=in_dim, ne[1]=1
            let in_dim = sh[0] as usize;
            let (_bi, bb) = st
                .raw("confidence_head.proj.bias")
                .ok_or("confidence bias missing beside weight")?;
            let with_markov = markov
                .as_ref()
                .map(|m| in_dim == cfg.hidden + m.rank)
                .unwrap_or(false);
            if !with_markov && in_dim != cfg.hidden {
                panic!(
                    "confidence_head in_dim {in_dim} matches neither hidden {} nor hidden+rank",
                    cfg.hidden
                );
            }
            Some(ConfidenceHead {
                w: bf16_to_f32(bytes),
                b: bf16_to_f32(bb)[0],
                in_dim,
                with_markov,
            })
        } else {
            None
        };
        // CENSUS GATE: every tensor in the export must be consumed by the map above.
        // DSpark-class checkpoints (markov head present) REFUSE on unrecognized names —
        // an unmapped tensor is a semantic program we would silently drop (house law).
        // Plain dflash checkpoints keep the historical warn-only behavior.
        {
            let mut consumed: std::collections::HashSet<String> = std::collections::HashSet::new();
            for i in 0..cfg.n_layer {
                for s in [
                    "self_attn.q_proj.weight",
                    "self_attn.k_proj.weight",
                    "self_attn.v_proj.weight",
                    "self_attn.o_proj.weight",
                    "self_attn.q_norm.weight",
                    "self_attn.k_norm.weight",
                    "input_layernorm.weight",
                    "post_attention_layernorm.weight",
                    "mlp.gate_proj.weight",
                    "mlp.up_proj.weight",
                    "mlp.down_proj.weight",
                ] {
                    consumed.insert(format!("layers.{i}.{s}"));
                }
            }
            for s in [
                "fc.weight",
                "hidden_norm.weight",
                "norm.weight",
                "markov_head.markov_w1.weight",
                "markov_head.markov_w2.weight",
                "confidence_head.proj.weight",
                "confidence_head.proj.bias",
            ] {
                consumed.insert(s.into());
            }
            let leftovers: Vec<&String> = st.names().filter(|n| !consumed.contains(*n)).collect();
            if !leftovers.is_empty() {
                if markov.is_some() {
                    panic!("dspark census: unrecognized tensors {leftovers:?}");
                }
                eprintln!("[dflash census] unmapped tensors (ignored): {leftovers:?}");
            }
        }
        // YaRN rope from config rope_parameters (HF _compute_yarn_parameters, verified
        // numerically vs Qwen3RotaryEmbedding on the arm-a export).
        let rope_yarn =
            if txt.contains("\"rope_type\": \"yarn\"") || txt.contains("\"rope_type\":\"yarn\"") {
                let factor = num(&txt, "factor").expect("yarn factor") as f64;
                let orig = num(&txt, "original_max_position_embeddings").expect("yarn orig");
                let beta_fast = num(&txt, "beta_fast").expect("beta_fast");
                let beta_slow = num(&txt, "beta_slow").expect("beta_slow");
                let base = cfg.rope_theta as f64;
                let d = cfg.head_dim as f64;
                let corr =
                    |r: f64| d * (orig / (r * 2.0 * std::f64::consts::PI)).ln() / (2.0 * base.ln());
                let low = corr(beta_fast).floor().max(0.0);
                let high = corr(beta_slow).ceil().min(d - 1.0);
                let half = cfg.head_dim / 2;
                let mut ff = Vec::with_capacity(half);
                for j in 0..half {
                    let base_inv = base.powf(-2.0 * j as f64 / d);
                    let ramp = (((j as f64) - low) / (high - low)).clamp(0.0, 1.0);
                    let ex = 1.0 - ramp; // extrapolation share
                    let yarn_inv = (base_inv / factor) * (1.0 - ex) + base_inv * ex;
                    ff.push((base_inv / yarn_inv) as f32);
                }
                let mscale = (0.1 * factor.ln() + 1.0) as f32;
                Some((e.htod(&ff)?, mscale))
            } else {
                None
            };
        Ok(Self {
            fc: upw("fc.weight")?,
            hidden_norm: up("hidden_norm.weight")?,
            norm: up("norm.weight")?,
            cfg,
            layers,
            markov,
            confidence,
            rope_yarn,
        })
    }

    /// Rope q or k rows in place: yarn (ff divisors + post-rope mscale) when the config
    /// carries it, plain neox otherwise. One primitive for all five drafter rope sites.
    fn rope_rows(
        &self,
        e: &Engine,
        x: &mut CudaSlice<f32>,
        pos_d: &CudaSlice<i32>,
        n_heads: usize,
        n_tokens: usize,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let c = &self.cfg;
        match &self.rope_yarn {
            Some((ff, mscale)) => {
                e.rope_neox_ff(
                    x,
                    pos_d,
                    c.head_dim,
                    c.head_dim,
                    n_heads,
                    n_tokens,
                    c.rope_theta,
                    1.0,
                    ff,
                )?;
                e.scale_inplace(x, *mscale, n_tokens * n_heads * c.head_dim)?;
            }
            None => {
                e.rope_neox(
                    x,
                    pos_d,
                    c.head_dim,
                    c.head_dim,
                    n_heads,
                    n_tokens,
                    c.rope_theta,
                    1.0,
                )?;
            }
        }
        Ok(())
    }

    /// f32 GEMM helper via the engine Float arm (cuBLASLt): y[t, out_f].
    fn mm(
        &self,
        e: &Engine,
        w: &GpuTensor,
        x: &CudaSlice<f32>,
        t: usize,
        _in_f: usize,
        _out_f: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        Ok(e.matmul(w, x, t)?)
    }

    /// FIRST-LIGHT forward (oracle contract): full non-causal attention over
    /// [ctx_features ; block], NO draft KV cache, NO sliding window (the oracle bypasses
    /// the reference mask machinery the same way — window/caching land in the round arm).
    ///
    /// `target_hidden`: [ctx, n_taps*hidden] (f32, device)  — raw tapped states.
    /// `noise_emb`:     [block, hidden] — target embed rows for [accepted, MASK x b-1].
    /// `pos`:           absolute positions for ctx rows THEN block rows (ctx+block i32).
    /// Returns final normed hidden [block, hidden] (feed target lm_head for draft logits).
    /// ctx features for `t` tapped rows: hidden_norm(fc(taps)) — the drafter's context
    /// representation, cacheable across rounds (append-only in committed-token order).
    pub fn ctx_features(
        &self,
        e: &Engine,
        taps: &CudaSlice<f32>,
        t: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let c = &self.cfg;
        let n_taps = c.target_layer_ids.len();
        let fc_out = self.mm(e, &self.fc, taps, t, n_taps * c.hidden, c.hidden)?;
        let mut out = e.uninit(t * c.hidden)?;
        e.rms_norm(&fc_out, &self.hidden_norm, &mut out, c.hidden, t, c.eps)?;
        Ok(out)
    }

    pub fn forward(
        &self,
        e: &Engine,
        target_hidden: &CudaSlice<f32>,
        noise_emb: &CudaSlice<f32>,
        pos: &[i32],
        ctx: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let ctx_f = self.ctx_features(e, target_hidden, ctx)?;
        if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
            let v = e.dtoh(&ctx_f)?;
            let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
            std::fs::write(format!("{dir}/memra-ctx_features.f32"), bytes)?;
        }
        self.forward_block(e, &ctx_f, noise_emb, pos, ctx)
    }

    /// Block forward over PRECOMPUTED ctx features (the round arm's entry: features are
    /// cached across rounds; only the block work repeats).
    pub fn forward_block(
        &self,
        e: &Engine,
        ctx_f: &CudaSlice<f32>,
        noise_emb: &CudaSlice<f32>,
        pos: &[i32],
        ctx: usize,
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let c = &self.cfg;
        let (h, nh, nkv, hd) = (c.hidden, c.n_head, c.n_kv, c.head_dim);
        let b = c.block_size;
        assert_eq!(pos.len(), ctx + b, "pos covers ctx rows then block rows");

        let pos_blk = e.htod_i32(&pos[ctx..])?;

        let mut x = e.clone_dtod(noise_emb)?; // [b, hidden] residual stream
        for (li, l) in self.layers.iter().enumerate() {
            let _ = li;
            // input_layernorm on the block rows only (ctx features are norm-free per ref:
            // k/v project the SAME ctx_f every layer, un-layernormed).
            let mut xn = e.uninit(b * h)?;
            e.rms_norm(&x, &l.ln_in, &mut xn, h, b, c.eps)?;

            // q from block; k/v from [ctx_f ; block-normed]
            let q0 = self.mm(e, &l.wq, &xn, b, h, nh * hd)?;
            let k0c = self.mm(e, &l.wk, ctx_f, ctx, h, nkv * hd)?;
            let v0c = self.mm(e, &l.wv, ctx_f, ctx, h, nkv * hd)?;
            let k0b = self.mm(e, &l.wk, &xn, b, h, nkv * hd)?;
            let v0b = self.mm(e, &l.wv, &xn, b, h, nkv * hd)?;

            // per-head q/k rms norm (v passes through: ones weight trick not needed — the
            // qkv kernel norms rq+rk rows; concatenate k first).
            let mut k0 = e.uninit((ctx + b) * nkv * hd)?;
            e.copy_into(&mut k0, 0, &k0c, ctx * nkv * hd)?;
            e.copy_into(&mut k0, ctx * nkv * hd, &k0b, b * nkv * hd)?;
            let mut v = e.uninit((ctx + b) * nkv * hd)?;
            e.copy_into(&mut v, 0, &v0c, ctx * nkv * hd)?;
            e.copy_into(&mut v, ctx * nkv * hd, &v0b, b * nkv * hd)?;

            if li == 0 {
                if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
                    let v = e.dtoh(&q0)?;
                    let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
                    std::fs::write(format!("{dir}/memra-l0_q0.f32"), bytes)?;
                }
            }
            let mut q = e.uninit(b * nh * hd)?;
            let mut k = e.uninit((ctx + b) * nkv * hd)?;
            // rms over head_dim rows: q has b*nh rows, k has (ctx+b)*nkv rows.
            e.rms_norm(&q0, &l.q_norm, &mut q, hd, b * nh, c.eps)?;
            if li == 0 {
                if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
                    let v = e.dtoh(&q)?;
                    let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
                    std::fs::write(format!("{dir}/memra-l0_qn.f32"), bytes)?;
                }
            }
            e.rms_norm(&k0, &l.k_norm, &mut k, hd, (ctx + b) * nkv, c.eps)?;

            // rope: q at block positions, k at ctx-then-block positions (absolute).
            let norope = std::env::var("MEMRA_DFLASH_NOROPE").is_ok();
            if !norope {
                self.rope_rows(e, &mut q, &pos_blk, nh, b)?;
            }
            if li == 0 {
                if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
                    let dump = |name: &str,
                                t: &cudarc::driver::CudaSlice<f32>|
                     -> Result<(), Box<dyn std::error::Error>> {
                        let v = e.dtoh(t)?;
                        let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
                        std::fs::write(format!("{dir}/memra-l0_{name}.f32"), bytes)?;
                        Ok(())
                    };
                    dump("xn", &xn)?;
                    dump("q_prerope", &q)?;
                }
            }
            // k rows are laid out [row, nkv, hd] with row-major tokens — rope_neox expects
            // (n_heads, n_tokens); ctx and block ropes run as one call over ctx+b tokens.
            let pos_all = e.htod_i32(pos)?;
            if !norope {
                self.rope_rows(e, &mut k, &pos_all, nkv, ctx + b)?;
            }

            // full non-causal attention: every block query sees all ctx+b keys.
            let mut attn = e.uninit(b * nh * hd)?;
            let scale = 1.0f32 / (hd as f32).sqrt();
            // NAIVE SDPA for first light: fa_prefill's NON-CAUSAL arm with T != T_kv is
            // BROKEN (attn maxdiff 0.34 vs the torch oracle; q/k inputs bit-close — no
            // existing caller exercises that shape class, jsonl 2026-07-13). The 16 x
            // (ctx+16) block attention is tiny; the fa arm returns behind this seam once
            // its kernel is fixed + parity-gated.
            if std::env::var("MEMRA_DFLASH_FA").is_ok() {
                e.fa_prefill(&q, &k, &v, &mut attn, hd, nh, nkv, b, ctx + b, scale, false)?;
            } else {
                e.sdpa_naive(&q, &k, &v, &mut attn, hd, nh, nkv, b, ctx + b, scale, false)?;
            }

            let o = self.mm(e, &l.wo, &attn, b, nh * hd, h)?;
            let mut x1 = e.uninit(b * h)?;
            e.add(&o, &x, &mut x1, b * h)?;
            if li == 0 {
                if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
                    let dump = |name: &str,
                                t: &cudarc::driver::CudaSlice<f32>|
                     -> Result<(), Box<dyn std::error::Error>> {
                        let v = e.dtoh(t)?;
                        let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
                        std::fs::write(format!("{dir}/memra-l0_{name}.f32"), bytes)?;
                        Ok(())
                    };
                    dump("q", &q)?;
                    dump("k", &k)?;
                    dump("attn", &attn)?;
                    dump("x1", &x1)?;
                }
            }

            // mlp
            let mut x1n = e.uninit(b * h)?;
            e.rms_norm(&x1, &l.ln_post, &mut x1n, h, b, c.eps)?;
            let gate = self.mm(e, &l.w_gate, &x1n, b, h, c.n_ff)?;
            let up_ = self.mm(e, &l.w_up, &x1n, b, h, c.n_ff)?;
            let mut act = e.uninit(b * c.n_ff)?;
            e.silu_mul(&gate, &up_, &mut act, b * c.n_ff)?;
            let down = self.mm(e, &l.w_down, &act, b, c.n_ff, h)?;
            let mut x2 = e.uninit(b * h)?;
            e.add(&down, &x1, &mut x2, b * h)?;
            x = x2;
            if let Ok(dir) = std::env::var("MEMRA_DFLASH_DUMP") {
                let v = e.dtoh(&x)?;
                let bytes: Vec<u8> = v.iter().flat_map(|f| f.to_le_bytes()).collect();
                std::fs::write(format!("{dir}/memra-layer{li}_out.f32"), bytes)?;
            }
        }
        let mut out = e.uninit(b * h)?;
        e.rms_norm(&x, &self.norm, &mut out, h, b, c.eps)?;
        Ok(out)
    }
}

/// Draft KV cache (round-cost fix, 2026-07-13): per-layer normed+roped ctx K and raw ctx V,
/// append-only in committed order. Block K/V land TRANSIENTLY at [len..len+b] each round
/// (never committed — the reference crops them identically). Kills the per-round full-ctx
/// projection recompute (first light was O(ctx)/round -> 7 tok/s).
pub struct DflashKv {
    pub k: Vec<CudaSlice<f32>>, // per layer [cap + block, nkv*hd]
    pub v: Vec<CudaSlice<f32>>,
    pub len: usize,
    pub cap: usize,
}

impl DflashKv {
    pub fn new(
        e: &Engine,
        cfg: &DflashCfg,
        cap: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let rowsz = cfg.n_kv * cfg.head_dim;
        let mut k = Vec::with_capacity(cfg.n_layer);
        let mut v = Vec::with_capacity(cfg.n_layer);
        for _ in 0..cfg.n_layer {
            k.push(e.uninit((cap + cfg.block_size) * rowsz)?);
            v.push(e.uninit((cap + cfg.block_size) * rowsz)?);
        }
        Ok(Self { k, v, len: 0, cap })
    }
}

impl DflashDraft {
    /// Ingest `t` NEW ctx-feature rows (committed order, absolute positions `pos_new`) into
    /// the draft KV: per layer k/v projections + k head-norm + rope, appended at kv.len.
    pub fn ingest_ctx(
        &self,
        e: &Engine,
        kv: &mut DflashKv,
        feats: &CudaSlice<f32>,
        pos_new: &[i32],
        t: usize,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let c = &self.cfg;
        let (h, nkv, hd) = (c.hidden, c.n_kv, c.head_dim);
        assert!(kv.len + t <= kv.cap, "draft kv overflow");
        let pos_d = e.htod_i32(pos_new)?;
        for (li, l) in self.layers.iter().enumerate() {
            let k0 = self.mm(e, &l.wk, feats, t, h, nkv * hd)?;
            let v0 = self.mm(e, &l.wv, feats, t, h, nkv * hd)?;
            let mut kn = e.uninit(t * nkv * hd)?;
            e.rms_norm(&k0, &l.k_norm, &mut kn, hd, t * nkv, c.eps)?;
            self.rope_rows(e, &mut kn, &pos_d, nkv, t)?;
            e.copy_into(&mut kv.k[li], kv.len * nkv * hd, &kn, t * nkv * hd)?;
            e.copy_into(&mut kv.v[li], kv.len * nkv * hd, &v0, t * nkv * hd)?;
        }
        kv.len += t;
        Ok(())
    }

    /// Block forward over the CACHED ctx KV: only the 16 block rows are projected per layer;
    /// block K/V land transiently at kv[len..len+b]. Bit-class-identical to forward_block
    /// (same kernels, same per-row programs; ONLY the ctx K/V recompute is cached).
    pub fn forward_round(
        &self,
        e: &Engine,
        kv: &mut DflashKv,
        noise_emb: &CudaSlice<f32>,
        pos_block: &[i32],
    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
        let c = &self.cfg;
        let (h, nh, nkv, hd) = (c.hidden, c.n_head, c.n_kv, c.head_dim);
        let b = c.block_size;
        assert_eq!(pos_block.len(), b);
        let ctx = kv.len;
        let pos_blk = e.htod_i32(pos_block)?;
        let mut x = e.clone_dtod(noise_emb)?;
        for (li, l) in self.layers.iter().enumerate() {
            let mut xn = e.uninit(b * h)?;
            e.rms_norm(&x, &l.ln_in, &mut xn, h, b, c.eps)?;
            let q0 = self.mm(e, &l.wq, &xn, b, h, nh * hd)?;
            let k0b = self.mm(e, &l.wk, &xn, b, h, nkv * hd)?;
            let v0b = self.mm(e, &l.wv, &xn, b, h, nkv * hd)?;
            let mut q = e.uninit(b * nh * hd)?;
            let mut kb = e.uninit(b * nkv * hd)?;
            e.rms_norm(&q0, &l.q_norm, &mut q, hd, b * nh, c.eps)?;
            e.rms_norm(&k0b, &l.k_norm, &mut kb, hd, b * nkv, c.eps)?;
            self.rope_rows(e, &mut q, &pos_blk, nh, b)?;
            self.rope_rows(e, &mut kb, &pos_blk, nkv, b)?;
            e.copy_into(&mut kv.k[li], ctx * nkv * hd, &kb, b * nkv * hd)?;
            e.copy_into(&mut kv.v[li], ctx * nkv * hd, &v0b, b * nkv * hd)?;
            let mut attn = e.uninit(b * nh * hd)?;
            let scale = 1.0f32 / (hd as f32).sqrt();
            if std::env::var("MEMRA_DFLASH_FA").is_ok() {
                e.fa_prefill(
                    &q,
                    &kv.k[li],
                    &kv.v[li],
                    &mut attn,
                    hd,
                    nh,
                    nkv,
                    b,
                    ctx + b,
                    scale,
                    false,
                )?;
            } else {
                e.sdpa_naive(
                    &q,
                    &kv.k[li],
                    &kv.v[li],
                    &mut attn,
                    hd,
                    nh,
                    nkv,
                    b,
                    ctx + b,
                    scale,
                    false,
                )?;
            }
            let o = self.mm(e, &l.wo, &attn, b, nh * hd, h)?;
            let mut x1 = e.uninit(b * h)?;
            e.add(&o, &x, &mut x1, b * h)?;
            let mut x1n = e.uninit(b * h)?;
            e.rms_norm(&x1, &l.ln_post, &mut x1n, h, b, c.eps)?;
            let gate = self.mm(e, &l.w_gate, &x1n, b, h, c.n_ff)?;
            let up_ = self.mm(e, &l.w_up, &x1n, b, h, c.n_ff)?;
            let mut act = e.uninit(b * c.n_ff)?;
            e.silu_mul(&gate, &up_, &mut act, b * c.n_ff)?;
            let down = self.mm(e, &l.w_down, &act, b, c.n_ff, h)?;
            let mut x2 = e.uninit(b * h)?;
            e.add(&down, &x1, &mut x2, b * h)?;
            x = x2;
        }
        let mut out = e.uninit(b * h)?;
        e.rms_norm(&x, &self.norm, &mut out, h, b, c.eps)?;
        Ok(out)
    }
}

// ================= DFlash spec round (greedy, first light) =================
// Exact contract: identical output stream to plain greedy decode BY CONSTRUCTION — the
// target's batched verify argmax decides every committed token; the drafter only proposes.
// (Same verify+rewind pattern as generate_spec_gemma's eager round; t=16 verify rides the
// straddle-split-safe fa_decode_rows.)
impl crate::hybrid::HybridModel {
    pub fn generate_spec_dflash(
        &self,
        e: &Engine,
        draft: &DflashDraft,
        prompt: &[u32],
        max_new: usize,
        eos: &[u32],
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        use crate::cache::{Cache, DflashTapSink};
        let n_embd = self.cfg.n_embd as usize;
        let c = &draft.cfg;
        assert_eq!(n_embd, c.hidden, "draft hidden must match target n_embd");
        let b = c.block_size;
        let n_taps = c.target_layer_ids.len();
        let max_ctx = prompt.len() + max_new + b + 8;
        // First light holds ctx <= sliding_window: the draft was trained with 4 sliding
        // layers (window 2048) and the first-light attention is windowless full — inside
        // the window the two are identical. The depth cell (1736 + 128) fits.
        assert!(
            max_ctx <= c.sliding_window,
            "first-light dflash round is windowless — ctx cap {} exceeds the draft window {}",
            max_ctx,
            c.sliding_window
        );
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;

        // ---- prime with taps armed ----
        let tp = prompt.len();
        cache.dflash_taps = Some(DflashTapSink {
            layer_ids: c.target_layer_ids.clone(),
            buf: e.uninit(tp * n_taps * n_embd)?,
            hidden: n_embd,
            t: tp,
            base: 0,
        });
        let t_prime = std::time::Instant::now();
        let (logits, _h_seed, _hiddens) = self.prime_cache(e, prompt, &mut cache, 0)?;
        let mut last = crate::forward::argmax(&logits) as u32;
        // draft KV cache: ingest the prompt's ctx features once; per round only the kept
        // rows ingest + the block projects (round cost O(block), not O(ctx)).
        let mut dkv = DflashKv::new(e, &draft.cfg, max_ctx)?;
        {
            // CHUNKED ingest (depth OOM fix): the 1736-row prompt tap buffer is ~224MB f32;
            // running fc + 5-layer k/v projection over it in one shot stacks another
            // ~300MB of transients on the ~21.3GB trunk peak. 256-row windows bound the
            // transient set; identical values (row-independent ops).
            let taps = cache.dflash_taps.take().unwrap();
            let n_taps_h = n_taps * n_embd;
            let mut r0 = 0usize;
            while r0 < tp {
                let t_c = (tp - r0).min(256);
                let tv = e.view(&taps.buf, tp * n_taps_h);
                let win = tv.slice(r0 * n_taps_h..(r0 + t_c) * n_taps_h);
                let mut chunk = e.uninit(t_c * n_taps_h)?;
                e.copy_view_into(&mut chunk, 0, &win, t_c * n_taps_h)?;
                let f = draft.ctx_features(e, &chunk, t_c)?;
                let pos_c: Vec<i32> = ((r0 as i32)..(r0 + t_c) as i32).collect();
                draft.ingest_ctx(e, &mut dkv, &f, &pos_c, t_c)?;
                r0 += t_c;
            }
        }
        let mut ctx_len = tp;
        e.stream().synchronize()?;
        // published prime wall (the run-spec/gemma-gate timing contract subtracts it)
        crate::PRIME_NANOS.store(
            t_prime.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );

        // embed-scale seam (MEMRA_DFLASH_EMB_SCALE): gemma trunks scale embeddings by
        // sqrt(n_embd) INSIDE the forward; whether the z-lab gemma4 training fed the
        // drafter scaled or raw embed rows is not visible from the reference (qwen path
        // uses raw embed_tokens). Acceptance arbitrates; default raw.
        let emb_scale = if std::env::var("MEMRA_DFLASH_EMB_SCALE").as_deref() == Ok("1") {
            (n_embd as f32).sqrt()
        } else {
            1.0
        };

        let mut out = Vec::with_capacity(max_new);
        let n_vocab = self.output.out_features();
        // VERIFY WIDTH (MEMRA_DFLASH_VERIFY_T, default 8): the drafter always drafts a full
        // block (its trained mask pattern) but only the first vt rows go through the target
        // verify — the t=16 verify rides the untuned b16 tier at ~32% of the byte wall
        // (65ms/verify) while b8 rides the tuned r2 tier; with ~2.7 committed/round the
        // deep block positions almost never survive anyway. Exactness unaffected (verify
        // still decides every committed token).
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(8)
            .clamp(2, b);
        // adaptive verify width (MEMRA_DFLASH_ADAPT!=0, MTP accepted+1 recipe): next round
        // verifies one past this round's accepted run, clamped [3, cap].
        let adapt = std::env::var("MEMRA_DFLASH_ADAPT").as_deref() != Ok("0");
        let mut vt = vt_cap;
        let mut attempted = 0usize;
        let mut accepted = 0usize;
        // The whole round runs in the decode-exact matmul scope: the m=16 draft mms were
        // otherwise falling into the prefill-GEMM class (770us/matmul, 17% of the depth
        // round). Prime (before this loop) keeps the prefill GEMM path.
        e.set_verify_exact(true);
        'outer: while out.len() < max_new {
            let start = cache.pos; // committed length
            // ---- draft: block = [last, MASK x b-1] ----
            let mut block: Vec<u32> = vec![c.mask_token_id; b];
            block[0] = last;
            let mut noise = e.htod(&self.embd.gather(n_embd, &block))?;
            if emb_scale != 1.0 {
                e.scale_inplace(&mut noise, emb_scale, b * n_embd)?;
            }
            if std::env::var("MEMRA_DFLASH_DEBUG").as_deref() == Ok("1") && start == cache.pos {
                let nv = e.dtoh(&noise)?;
                let r0: f32 = nv[..n_embd].iter().map(|x| x * x).sum::<f32>().sqrt();
                let r1: f32 = nv[n_embd..2 * n_embd]
                    .iter()
                    .map(|x| x * x)
                    .sum::<f32>()
                    .sqrt();
                eprintln!(
                    "[dflash noise] |row0(last)|={r0:.3} |row1(MASK id {})|={r1:.3}",
                    c.mask_token_id
                );
            }
            let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
            let dh = draft.forward_round(e, &mut dkv, &noise, &pos_block)?;
            // draft tokens = argmax(lm_head(h rows 1..b))
            let mut rows = e.uninit((b - 1) * n_embd)?;
            {
                let dv = e.view(&dh, b * n_embd);
                let tail = dv.slice(n_embd..b * n_embd);
                e.copy_view_into(&mut rows, 0, &tail, (b - 1) * n_embd)?;
            }
            let mut dl = e.matmul(&self.output, &rows, b - 1)?;
            // SEMI-AR MARKOV CHAIN (DSpark head, when present + MEMRA_DFLASH_MARKOV!=0):
            // left-to-right, logits_k += W2(W1[prev realized token]) — the whole chain
            // stays on-device (chain_d[0] = the pending token; argmax k writes
            // chain_d[k+1], the k+1 bias gathers from it). Greedy mirror of the patch's
            // _markov_semiar_sample_block.
            let markov_on = std::env::var("MEMRA_DFLASH_MARKOV").as_deref() != Ok("0");
            let mut chain_d = e.stream().alloc_zeros::<u32>(b)?;
            if let (Some(mk), true) = (&draft.markov, markov_on) {
                e.set_u32_one(&mut chain_d, last)?;
                for k in 0..(b - 1) {
                    let mut f = e.uninit(mk.rank)?;
                    e.gather_row_bf16(&mk.w1_bf16, &chain_d, k, &mut f, mk.rank)?;
                    let bias = e.matmul(&mk.w2, &f, 1)?;
                    e.add_row_inplace(&mut dl, &bias, n_vocab, k * n_vocab)?;
                    e.argmax_token_device_col(&dl, k, n_vocab, &mut chain_d, k + 1)?;
                }
            } else {
                for i in 0..(b - 1) {
                    e.argmax_token_device_col(&dl, i, n_vocab, &mut chain_d, i + 1)?;
                }
            }
            let chain = e.dtoh_u32(&chain_d)?;
            let dtoks = &chain[1..];
            for (i, &dt) in dtoks.iter().enumerate() {
                block[i + 1] = dt;
            }
            let dbg = std::env::var("MEMRA_DFLASH_DEBUG").as_deref() == Ok("1");

            // ---- verify: one t=vt target forward with taps armed ----
            let vblock = &block[..vt];
            cache.dflash_taps = Some(DflashTapSink {
                layer_ids: c.target_layer_ids.clone(),
                buf: e.uninit(vt * n_taps * n_embd)?,
                hidden: n_embd,
                t: vt,
                base: 0,
            });
            let (vam, _vh) = self.gemma4_decode_step_t_am(e, vblock, start, &mut cache)?;
            let taps = cache.dflash_taps.take().unwrap();
            if dbg {
                eprintln!(
                    "[dflash r] start={start} last={last}\n  draft={:?}\n  vam  ={:?}",
                    &block[1..],
                    &vam
                );
            }

            // ---- accept ----
            let mut m = 0usize;
            while m < vt - 1 && block[m + 1] as usize == vam[m] as usize {
                m += 1;
            }
            attempted += vt - 1;
            accepted += m;
            out.push(last);
            if eos.contains(&last) {
                break 'outer;
            }
            for &dt in &block[1..=m] {
                out.push(dt);
                if eos.contains(&dt) {
                    break 'outer;
                }
                if out.len() >= max_new {
                    break 'outer;
                }
            }
            let next = vam[m] as u32;

            // ---- commit/rollback: keep m+1 of the b appended rows ----
            let keep = m + 1;
            for kvl in cache.kv.iter_mut().flatten() {
                kvl.len -= vt - keep;
                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
            }
            cache.pos -= vt - keep;

            // ---- ingest the kept rows' ctx features into the draft KV ----
            {
                let tv = e.view(&taps.buf, vt * n_taps * n_embd);
                let keep_view = tv.slice(0..keep * n_taps * n_embd);
                let mut kept = e.uninit(keep * n_taps * n_embd)?;
                e.copy_view_into(&mut kept, 0, &keep_view, keep * n_taps * n_embd)?;
                let f = draft.ctx_features(e, &kept, keep)?;
                let pos_k: Vec<i32> = ((ctx_len as i32)..(ctx_len + keep) as i32).collect();
                draft.ingest_ctx(e, &mut dkv, &f, &pos_k, keep)?;
                ctx_len += keep;
            }
            last = next;
            if adapt {
                vt = (m + 2).clamp(3, vt_cap);
            }
        }
        e.set_verify_exact(false);
        if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
            eprintln!(
                "[dflash] acceptance {accepted}/{attempted} = {:.3}",
                accepted as f64 / attempted.max(1) as f64
            );
        }
        Ok(out)
    }
}

// ================= DSpark spec round, QWEN-HYBRID target (lane/dspark-q38-recover) =====
// The q38 twin of generate_spec_dflash. Same drafter machinery (rounds, markov chain,
// draft KV, adaptive verify width); the TARGET side swaps gemma4's dense verify for the
// qwen serving-class verify funnel (dspark_verify_t_am) + snapshot/rollback, because the
// hybrid GDN conv/ssm state mutates in place — dense KV truncation cannot roll it back.
// Exactness contract unchanged: identical stream to plain greedy BY CONSTRUCTION (the
// target's verify argmax decides every committed token).
impl crate::hybrid::HybridModel {
    pub fn generate_spec_dspark(
        &self,
        e: &Engine,
        draft: &DflashDraft,
        prompt: &[u32],
        max_new: usize,
        eos: &[u32],
    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
        use crate::cache::{Cache, DflashTapSink};
        assert!(
            self.cfg.gemma4.is_none(),
            "gemma4 targets use generate_spec_dflash; this is the qwen-hybrid arm"
        );
        let n_embd = self.cfg.n_embd as usize;
        let c = &draft.cfg;
        assert_eq!(n_embd, c.hidden, "draft hidden must match target n_embd");
        let b = c.block_size;
        let n_taps = c.target_layer_ids.len();
        let max_ctx = prompt.len() + max_new + b + 8;
        assert!(
            max_ctx <= c.sliding_window,
            "dspark round is windowless — ctx cap {} exceeds the draft window {}",
            max_ctx,
            c.sliding_window
        );
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;

        // ---- prime with taps armed (chunked prime writes at chunk offsets via sink.base) ----
        let tp = prompt.len();
        cache.dflash_taps = Some(DflashTapSink {
            layer_ids: c.target_layer_ids.clone(),
            buf: e.uninit(tp * n_taps * n_embd)?,
            hidden: n_embd,
            t: tp,
            base: 0,
        });
        let t_prime = std::time::Instant::now();
        let (logits, _h_seed, _hiddens) = self.prime_cache(e, prompt, &mut cache, 0)?;
        let mut last = crate::forward::argmax(&logits) as u32;
        let mut dkv = DflashKv::new(e, &draft.cfg, max_ctx)?;
        {
            let taps = cache.dflash_taps.take().unwrap();
            let n_taps_h = n_taps * n_embd;
            let mut r0 = 0usize;
            while r0 < tp {
                let t_c = (tp - r0).min(256);
                let tv = e.view(&taps.buf, tp * n_taps_h);
                let win = tv.slice(r0 * n_taps_h..(r0 + t_c) * n_taps_h);
                let mut chunk = e.uninit(t_c * n_taps_h)?;
                e.copy_view_into(&mut chunk, 0, &win, t_c * n_taps_h)?;
                let f = draft.ctx_features(e, &chunk, t_c)?;
                let pos_c: Vec<i32> = ((r0 as i32)..(r0 + t_c) as i32).collect();
                draft.ingest_ctx(e, &mut dkv, &f, &pos_c, t_c)?;
                r0 += t_c;
            }
        }
        let mut ctx_len = tp;
        e.stream().synchronize()?;
        crate::PRIME_NANOS.store(
            t_prime.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );

        let mut out = Vec::with_capacity(max_new);
        let n_vocab = self.output.out_features();
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(b)
            .clamp(2, b);
        let adapt = std::env::var("MEMRA_DFLASH_ADAPT").as_deref() != Ok("0");
        let mut vt = vt_cap;
        let mut attempted = 0usize;
        let mut accepted = 0usize;
        // per-phase economics counters (ns) — the verify-toll dataset
        let (mut ns_draft, mut ns_snap, mut ns_verify, mut ns_roll, mut ns_ingest) =
            (0u64, 0u64, 0u64, 0u64, 0u64);
        let mut rounds = 0usize;
        let stats = std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1");
        let clock = |on: bool, e: &Engine| -> std::time::Instant {
            if on {
                let _ = e.stream().synchronize();
            }
            std::time::Instant::now()
        };
        'outer: while out.len() < max_new {
            rounds += 1;
            let start = cache.pos; // committed length
            // ---- draft: block = [last, MASK x b-1] (decode-exact class for the m=b mms) ----
            let t0 = clock(stats, e);
            e.set_verify_exact(true);
            let mut block: Vec<u32> = vec![c.mask_token_id; b];
            block[0] = last;
            let noise = e.htod(&self.embd.gather(n_embd, &block))?;
            let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
            let dh = draft.forward_round(e, &mut dkv, &noise, &pos_block)?;
            let mut rows = e.uninit((b - 1) * n_embd)?;
            {
                let dv = e.view(&dh, b * n_embd);
                let tail = dv.slice(n_embd..b * n_embd);
                e.copy_view_into(&mut rows, 0, &tail, (b - 1) * n_embd)?;
            }
            let mut dl = e.matmul(&self.output, &rows, b - 1)?;
            let markov_on = std::env::var("MEMRA_DFLASH_MARKOV").as_deref() != Ok("0");
            let mut chain_d = e.stream().alloc_zeros::<u32>(b)?;
            if let (Some(mk), true) = (&draft.markov, markov_on) {
                e.set_u32_one(&mut chain_d, last)?;
                for k in 0..(b - 1) {
                    let mut f = e.uninit(mk.rank)?;
                    e.gather_row_bf16(&mk.w1_bf16, &chain_d, k, &mut f, mk.rank)?;
                    let bias = e.matmul(&mk.w2, &f, 1)?;
                    e.add_row_inplace(&mut dl, &bias, n_vocab, k * n_vocab)?;
                    e.argmax_token_device_col(&dl, k, n_vocab, &mut chain_d, k + 1)?;
                }
            } else {
                for i in 0..(b - 1) {
                    e.argmax_token_device_col(&dl, i, n_vocab, &mut chain_d, i + 1)?;
                }
            }
            e.set_verify_exact(false);
            let chain = e.dtoh_u32(&chain_d)?;
            for (i, &dt) in chain[1..].iter().enumerate() {
                block[i + 1] = dt;
            }
            ns_draft += clock(stats, e).duration_since(t0).as_nanos() as u64;

            // ---- snapshot (GDN conv/ssm state + KV lens), then verify t=vt ----
            let t1 = std::time::Instant::now();
            let snap = cache.snapshot(e)?;
            ns_snap += clock(stats, e).duration_since(t1).as_nanos() as u64;
            let t2 = std::time::Instant::now();
            let vblock = &block[..vt];
            cache.dflash_taps = Some(DflashTapSink {
                layer_ids: c.target_layer_ids.clone(),
                buf: e.uninit(vt * n_taps * n_embd)?,
                hidden: n_embd,
                t: vt,
                base: 0,
            });
            // MEMRA_DSPARK_CKPT (default 1): verify with the MTP column-stash armed so a
            // partial accept restores state directly. =0 keeps the snapshot+replay arm
            // (the oracle the stash arm is gated against — MEMRA_DSPARK_CKPT_GATE=1 runs
            // BOTH per partial round and byte-compares the resulting cache state).
            let ckpt_on = std::env::var("MEMRA_DSPARK_CKPT").as_deref() != Ok("0");
            let ckpt_gate = std::env::var("MEMRA_DSPARK_CKPT_GATE").as_deref() == Ok("1");
            let (vam, vck) = if ckpt_on || ckpt_gate {
                let (vam, vck) = self.dspark_verify_t_am_ckpt(e, vblock, start, &mut cache)?;
                (vam, Some(vck))
            } else {
                (self.dspark_verify_t_am(e, vblock, start, &mut cache)?, None)
            };
            let taps = cache.dflash_taps.take().unwrap();
            ns_verify += clock(stats, e).duration_since(t2).as_nanos() as u64;

            // ---- accept ----
            let mut m = 0usize;
            while m < vt - 1 && block[m + 1] == vam[m] {
                m += 1;
            }
            attempted += vt - 1;
            accepted += m;
            out.push(last);
            if eos.contains(&last) {
                break 'outer;
            }
            for &dt in &block[1..=m] {
                // budget check BEFORE the push: at real acceptance the final round often
                // accepts a draft at the boundary, and push-then-check emitted max_new+1
                // tokens (plain emits exactly max_new — the E2E gate read it as a length
                // divergence at index max_new with the shared prefix byte-identical).
                if out.len() >= max_new {
                    break 'outer;
                }
                out.push(dt);
                if eos.contains(&dt) {
                    break 'outer;
                }
            }
            let next = vam[m];

            // ---- commit/rollback: hybrid state cannot truncate — restore + replay kept ----
            let keep = m + 1;
            let t3 = std::time::Instant::now();
            if keep < vt {
                if ckpt_gate {
                    // GATE ARM: stash-restore, snapshot S1; then the replay oracle, snapshot
                    // S2; the two cache states must match BIT-FOR-BIT (kv lens, pos, every
                    // conv/ssm buffer). Continue from the replay state (proven identical).
                    let vck = vck.as_ref().expect("gate arm always fills the ckpt");
                    self.dspark_commit_prefix(e, &mut cache, &snap, vck, keep)?;
                    // host-side state capture (NO device snapshot copies — two extra
                    // device snapshots per round OOM'd beside the 15GB trunk)
                    let capture = |cache: &Cache| -> Result<
                        (usize, Vec<Option<usize>>, Vec<(Vec<f32>, Vec<f32>)>),
                        Box<dyn std::error::Error>,
                    > {
                        let mut lens = Vec::new();
                        let mut states = Vec::new();
                        for il in 0..cache.kv.len() {
                            lens.push(cache.kv[il].as_ref().map(|k| k.len));
                            if let Some(rl) = &cache.recur[il] {
                                states.push((e.dtoh(&rl.conv_state)?, e.dtoh(&rl.ssm_state)?));
                            }
                        }
                        Ok((cache.pos, lens, states))
                    };
                    let (p1, l1, st1) = capture(&cache)?;
                    crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut cache, &snap)?;
                    let ram = self.dspark_verify_t_am(e, &block[..keep], start, &mut cache)?;
                    assert_eq!(
                        &ram[..],
                        &vam[..keep],
                        "prefix replay must reproduce the verify argmaxes"
                    );
                    let (p2, l2, st2) = capture(&cache)?;
                    assert_eq!(p1, p2, "ckpt-gate: pos mismatch");
                    assert_eq!(l1, l2, "ckpt-gate: kv_len mismatch");
                    for (il, ((c1, s1v), (c2, s2v))) in st1.iter().zip(&st2).enumerate() {
                        let bits = |a: &[f32], b: &[f32]| {
                            a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
                        };
                        assert!(
                            bits(c1, c2),
                            "ckpt-gate: linear layer {il} conv state differs"
                        );
                        assert!(
                            bits(s1v, s2v),
                            "ckpt-gate: linear layer {il} ssm state differs"
                        );
                    }
                } else if let Some(vck) = vck.as_ref() {
                    // STASH ARM (default): column-state restore, no replay forward.
                    self.dspark_commit_prefix(e, &mut cache, &snap, vck, keep)?;
                } else {
                    // REPLAY ARM (MEMRA_DSPARK_CKPT=0): the original snapshot+replay oracle.
                    crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut cache, &snap)?;
                    debug_assert_eq!(cache.pos, start, "rollback landed off the round start");
                    let ram = self.dspark_verify_t_am(e, &block[..keep], start, &mut cache)?;
                    debug_assert_eq!(
                        &ram[..],
                        &vam[..keep],
                        "prefix replay must reproduce the verify argmaxes"
                    );
                }
            }
            ns_roll += clock(stats, e).duration_since(t3).as_nanos() as u64;

            // ---- ingest the kept rows' ctx features into the draft KV ----
            let t4 = std::time::Instant::now();
            {
                let tv = e.view(&taps.buf, vt * n_taps * n_embd);
                let keep_view = tv.slice(0..keep * n_taps * n_embd);
                let mut kept = e.uninit(keep * n_taps * n_embd)?;
                e.copy_view_into(&mut kept, 0, &keep_view, keep * n_taps * n_embd)?;
                let f = draft.ctx_features(e, &kept, keep)?;
                let pos_k: Vec<i32> = ((ctx_len as i32)..(ctx_len + keep) as i32).collect();
                draft.ingest_ctx(e, &mut dkv, &f, &pos_k, keep)?;
                ctx_len += keep;
            }
            ns_ingest += clock(stats, e).duration_since(t4).as_nanos() as u64;
            last = next;
            if adapt {
                vt = (m + 2).clamp(3, vt_cap);
            }
        }
        if stats {
            let ms = |n: u64| n as f64 / 1e6;
            eprintln!(
                "[dspark-q38] acceptance {accepted}/{attempted} = {:.3} rounds={rounds} \
                 draft={:.1}ms snap={:.1}ms verify={:.1}ms rollback+replay={:.1}ms ingest={:.1}ms",
                accepted as f64 / attempted.max(1) as f64,
                ms(ns_draft),
                ms(ns_snap),
                ms(ns_verify),
                ms(ns_roll),
                ms(ns_ingest)
            );
        }
        Ok(out)
    }
}

// ================= DSpark SERVING session (lane/dspark-q38-recover serve route) =========
// Burst-scoped state for the worker's dspark spec arm — the qwen-hybrid twin of
// GemmaSpecSession. Holds the trunk cache + draft KV + the round loop's carry state
// (`last`, ctx_len, adaptive vt) so the scheduler round-robins other sessions between
// bursts. The round body is generate_spec_dspark's loop, hoisted; that bin arm stays the
// banked oracle (E2E gate), and the serve-route smoke gates this twin byte-identical to
// a spec-off boot over the real HTTP surface. Exactness contract unchanged: the target's
// verify argmax decides every committed token, so the stream equals plain greedy BY
// CONSTRUCTION on every accept path (ckpt stash, gate, replay).
pub struct DsparkSpecSession {
    pub cache: crate::cache::Cache,
    dkv: DflashKv,
    last: u32,
    ctx_len: usize,
    vt: usize,
    pub rounds: usize,
    max_ctx: usize,
    done: bool,
}

impl DsparkSpecSession {
    pub fn cache_max_ctx(&self) -> usize {
        self.max_ctx
    }
    pub fn finished(&self) -> bool {
        self.done
    }
    pub fn pos(&self) -> usize {
        self.cache.pos
    }
}

impl crate::hybrid::HybridModel {
    /// Turn-1 prime: trunk prefill with taps armed + chunked ctx ingest into the draft KV.
    /// Mirrors generate_spec_dspark's prime block exactly (chunk offsets via sink.base are
    /// handled inside prime_cache's tick loop; the 256-row ingest chunks match the bin arm).
    pub fn dspark_spec_session_new(
        &self,
        e: &Engine,
        draft: &DflashDraft,
        prompt: &[u32],
        ctx_cap: usize,
    ) -> Result<DsparkSpecSession, Box<dyn std::error::Error>> {
        use crate::cache::{Cache, DflashTapSink};
        assert!(
            self.cfg.gemma4.is_none(),
            "gemma4 targets use the assistant-drafter route; dspark is the qwen-hybrid arm"
        );
        let n_embd = self.cfg.n_embd as usize;
        let c = &draft.cfg;
        assert_eq!(n_embd, c.hidden, "draft hidden must match target n_embd");
        let b = c.block_size;
        let n_taps = c.target_layer_ids.len();
        // The dspark round is windowless: every position the session will ever hold must
        // fit the draft window. Clamp the session ctx to it and refuse prompts that
        // cannot take even one round — admission falls back to the plain path.
        let max_ctx = ctx_cap.min(c.sliding_window);
        if prompt.len() + b + 8 > max_ctx {
            return Err(format!(
                "dspark session needs {} ctx (prompt {} + block {b} + 8), cap {max_ctx}",
                prompt.len() + b + 8,
                prompt.len()
            )
            .into());
        }
        let mut cache = Cache::new(e, &self.cfg, max_ctx)?;
        let tp = prompt.len();
        cache.dflash_taps = Some(DflashTapSink {
            layer_ids: c.target_layer_ids.clone(),
            buf: e.uninit(tp * n_taps * n_embd)?,
            hidden: n_embd,
            t: tp,
            base: 0,
        });
        let (logits, _h_seed, _hiddens) = self.prime_cache(e, prompt, &mut cache, 0)?;
        let last = crate::forward::argmax(&logits) as u32;
        let mut dkv = DflashKv::new(e, &draft.cfg, max_ctx)?;
        {
            let taps = cache.dflash_taps.take().unwrap();
            let n_taps_h = n_taps * n_embd;
            let mut r0 = 0usize;
            while r0 < tp {
                let t_c = (tp - r0).min(256);
                let tv = e.view(&taps.buf, tp * n_taps_h);
                let win = tv.slice(r0 * n_taps_h..(r0 + t_c) * n_taps_h);
                let mut chunk = e.uninit(t_c * n_taps_h)?;
                e.copy_view_into(&mut chunk, 0, &win, t_c * n_taps_h)?;
                let f = draft.ctx_features(e, &chunk, t_c)?;
                let pos_c: Vec<i32> = ((r0 as i32)..(r0 + t_c) as i32).collect();
                draft.ingest_ctx(e, &mut dkv, &f, &pos_c, t_c)?;
                r0 += t_c;
            }
        }
        e.stream().synchronize()?;
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(b)
            .clamp(2, b);
        Ok(DsparkSpecSession {
            cache,
            dkv,
            last,
            ctx_len: tp,
            vt: vt_cap,
            rounds: 0,
            max_ctx,
            done: false,
        })
    }

    /// One scheduler burst: dspark rounds until >= `burst_target` tokens are committed,
    /// EOS lands, or the ctx cap is reached. Returns (tokens, drafted, accepted) for this
    /// burst — the worker clamps the public slice (engine overshoot within a round stays
    /// in the session cache, exactly the gemma-burst contract).
    pub fn dspark_spec_session_burst(
        &self,
        e: &Engine,
        draft: &DflashDraft,
        sess: &mut DsparkSpecSession,
        burst_target: usize,
        eos: &[u32],
    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
        use crate::cache::DflashTapSink;
        let n_embd = self.cfg.n_embd as usize;
        let c = &draft.cfg;
        let b = c.block_size;
        let n_taps = c.target_layer_ids.len();
        let n_vocab = self.output.out_features();
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(b)
            .clamp(2, b);
        let adapt = std::env::var("MEMRA_DFLASH_ADAPT").as_deref() != Ok("0");
        let mut out: Vec<u32> = Vec::with_capacity(burst_target + b);
        let mut drafted = 0usize;
        let mut accepted_n = 0usize;
        'outer: while out.len() < burst_target && !sess.done {
            let start = sess.cache.pos;
            if start + b + 1 > sess.max_ctx {
                sess.done = true;
                break;
            }
            sess.rounds += 1;
            let vt = sess.vt;
            // ---- draft: block = [last, MASK x b-1] (identical to the bin arm) ----
            e.set_verify_exact(true);
            let mut block: Vec<u32> = vec![c.mask_token_id; b];
            block[0] = sess.last;
            let noise = e.htod(&self.embd.gather(n_embd, &block))?;
            let pos_block: Vec<i32> = ((start as i32)..(start + b) as i32).collect();
            let dh = draft.forward_round(e, &mut sess.dkv, &noise, &pos_block)?;
            let mut rows = e.uninit((b - 1) * n_embd)?;
            {
                let dv = e.view(&dh, b * n_embd);
                let tail = dv.slice(n_embd..b * n_embd);
                e.copy_view_into(&mut rows, 0, &tail, (b - 1) * n_embd)?;
            }
            let mut dl = e.matmul(&self.output, &rows, b - 1)?;
            let markov_on = std::env::var("MEMRA_DFLASH_MARKOV").as_deref() != Ok("0");
            let mut chain_d = e.stream().alloc_zeros::<u32>(b)?;
            if let (Some(mk), true) = (&draft.markov, markov_on) {
                e.set_u32_one(&mut chain_d, sess.last)?;
                for k in 0..(b - 1) {
                    let mut f = e.uninit(mk.rank)?;
                    e.gather_row_bf16(&mk.w1_bf16, &chain_d, k, &mut f, mk.rank)?;
                    let bias = e.matmul(&mk.w2, &f, 1)?;
                    e.add_row_inplace(&mut dl, &bias, n_vocab, k * n_vocab)?;
                    e.argmax_token_device_col(&dl, k, n_vocab, &mut chain_d, k + 1)?;
                }
            } else {
                for i in 0..(b - 1) {
                    e.argmax_token_device_col(&dl, i, n_vocab, &mut chain_d, i + 1)?;
                }
            }
            e.set_verify_exact(false);
            let chain = e.dtoh_u32(&chain_d)?;
            for (i, &dt) in chain[1..].iter().enumerate() {
                block[i + 1] = dt;
            }

            // ---- snapshot, then verify t=vt (ckpt stash default; oracle arms kept) ----
            let snap = sess.cache.snapshot(e)?;
            let vblock = &block[..vt];
            sess.cache.dflash_taps = Some(DflashTapSink {
                layer_ids: c.target_layer_ids.clone(),
                buf: e.uninit(vt * n_taps * n_embd)?,
                hidden: n_embd,
                t: vt,
                base: 0,
            });
            let ckpt_on = std::env::var("MEMRA_DSPARK_CKPT").as_deref() != Ok("0");
            let (vam, vck) = if ckpt_on {
                let (vam, vck) = self.dspark_verify_t_am_ckpt(e, vblock, start, &mut sess.cache)?;
                (vam, Some(vck))
            } else {
                (
                    self.dspark_verify_t_am(e, vblock, start, &mut sess.cache)?,
                    None,
                )
            };
            let taps = sess.cache.dflash_taps.take().unwrap();

            // ---- accept ----
            let mut m = 0usize;
            while m < vt - 1 && block[m + 1] == vam[m] {
                m += 1;
            }
            drafted += vt - 1;
            accepted_n += m;
            out.push(sess.last);
            if eos.contains(&sess.last) {
                sess.done = true;
                break 'outer;
            }
            for &dt in &block[1..=m] {
                out.push(dt);
                if eos.contains(&dt) {
                    sess.done = true;
                    break 'outer;
                }
            }
            let next = vam[m];

            // ---- commit/rollback (stash arm default; replay oracle kept) ----
            let keep = m + 1;
            if keep < vt {
                if let Some(vck) = vck.as_ref() {
                    self.dspark_commit_prefix(e, &mut sess.cache, &snap, vck, keep)?;
                } else {
                    crate::pp::restore_cache_checkpoint(
                        e,
                        &self.cfg,
                        None,
                        &mut sess.cache,
                        &snap,
                    )?;
                    debug_assert_eq!(sess.cache.pos, start, "rollback landed off the round start");
                    let ram = self.dspark_verify_t_am(e, &block[..keep], start, &mut sess.cache)?;
                    debug_assert_eq!(
                        &ram[..],
                        &vam[..keep],
                        "prefix replay must reproduce the verify argmaxes"
                    );
                }
            }

            // ---- ingest the kept rows' ctx features into the draft KV ----
            {
                let tv = e.view(&taps.buf, vt * n_taps * n_embd);
                let keep_view = tv.slice(0..keep * n_taps * n_embd);
                let mut kept = e.uninit(keep * n_taps * n_embd)?;
                e.copy_view_into(&mut kept, 0, &keep_view, keep * n_taps * n_embd)?;
                let f = draft.ctx_features(e, &kept, keep)?;
                let pos_k: Vec<i32> =
                    ((sess.ctx_len as i32)..(sess.ctx_len + keep) as i32).collect();
                draft.ingest_ctx(e, &mut sess.dkv, &f, &pos_k, keep)?;
                sess.ctx_len += keep;
            }
            sess.last = next;
            if adapt {
                sess.vt = (m + 2).clamp(3, vt_cap);
            }
        }
        Ok((out, drafted, accepted_n))
    }
}