memra-engine 0.98.0

//! 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>,
    /// Checkpoint training-strategy census (`dspark_strategy_census` over the raw
    /// config.json): true = a SpecForge DSPARK-strategy export (shifted labels, ALL rows
    /// supervised — the q38 arm-a family). Keys the HARVEST DEFAULT strategy-keyed,
    /// never env-keyed (owner-ratified 2026-08-20 after B1 confirmed H1 ×5;
    /// DSPARK-POSTMORTEM-20260820.md B0 default-flip plan).
    pub strategy_dspark: 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). sglang's DSPARK planner
    /// consumes it to SIZE VERIFY WINDOWS (cumprod survival — v0.5.16 headline; the
    /// earlier "reference serving loop never consumes it" note matched SpecForge's
    /// legacy spec_generate only). memra schedules with it under
    /// `MEMRA_DSPARK_VT=confidence` (the H4 fix, DSPARK-POSTMORTEM-20260820.md:
    /// per-round verify window from cumprod survival, `dspark_confidence_vt`) and
    /// keeps it census+parity-only under the default ladder. 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,
}

impl ConfidenceHead {
    /// Host dot: the PRE-sigmoid accept score for one draft slot. `hidden` = the
    /// drafter output row the slot is harvested from (the same row its logits use);
    /// `emb` = the markov `w1` row of the slot's PREVIOUS chain token (required iff
    /// `with_markov`) — the exact input contract the parity gate pins (prev ids =
    /// `[anchor, chain[..nd-1]]`, dspark_q38_parity.rs stage 5).
    pub fn raw_score(&self, hidden: &[f32], emb: Option<&[f32]>) -> f32 {
        let mut acc = self.b;
        for (w, x) in self.w.iter().zip(hidden) {
            acc += w * x;
        }
        if self.with_markov {
            let emb = emb.expect("with_markov confidence head scored without the markov embedding");
            debug_assert_eq!(hidden.len() + emb.len(), self.in_dim);
            for (w, x) in self.w[hidden.len()..].iter().zip(emb) {
                acc += w * x;
            }
        } else {
            debug_assert_eq!(hidden.len(), self.in_dim);
        }
        acc
    }
}

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,
}

/// Draft-row harvest convention for DFlash-family block drafters
/// (darklanes research/deepseek-flash-20260818/DSPARK-POSTMORTEM-20260820.md).
///
/// The DFlash and DSpark SpecForge training strategies supervise DIFFERENT rows of the
/// same `[anchor, MASK x b-1]` block, so the row -> trunk-position mapping is a property
/// of the CHECKPOINT's training strategy, not of the loader:
///
/// - **Dflash** (mask-fill; z-lab dflash / SpecForge `OnlineDFlashModel`): row k is
///   trained to predict the token AT position anchor+k — "Labels: same-position
///   prediction", `weight_mask *= (pos_in_block > 0)` excludes the anchor row
///   (SpecForge `specforge/algorithms/common/dflash_family_model.py:453-472`).
///   Drafts = rows 1..b-1; the anchor row's output is untrained.
/// - **Dspark** (shifted; SpecForge `OnlineDSparkModel`, `training.strategy: dspark` —
///   the q38 arm-a export): row k is trained to predict the token at anchor+k+1, ALL
///   rows supervised INCLUDING the anchor row (`label_offsets = arange(1,
///   block_size+1)`, `dflash_family_model.py:816`). sglang's DSPARK worker — the stack
///   every arm-a bank number was measured on — harvests gamma = block_size drafts with
///   the anchor row's output as draft 1 (verified on the v0.5.17 eval-pin tag:
///   `dspark_components/dspark_draft.py:248,260,318`; `dspark_config.py:269`).
///
/// Mismatching the convention verifies every slot against a position the row was never
/// trained for — the q38 accept collapse (2.9 -> 1.43) in the postmortem.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum DsparkHarvest {
    /// mask-fill: drafts = rows 1..b-1, row k fills position anchor+k.
    Dflash,
    /// shifted: drafts = rows 0..b-1, row k predicts position anchor+k+1.
    Dspark,
}

impl DsparkHarvest {
    /// The served resolution: explicit `MEMRA_DSPARK_HARVEST={dflash|dspark}` wins
    /// (unknown values REFUSE loudly — a typo silently reverting the convention would
    /// re-open the postmortem's misalignment); UNSET defers to the CHECKPOINT's own
    /// training-strategy census — the owner-ratified default flip (2026-08-20, after
    /// B1 confirmed H1 interleaved ×5 on serving-class hardware: accept 1.38→2.41
    /// agentic / 1.53→3.66 math, E2E ALL EXACT both arms). Strategy-keyed, not
    /// env-keyed, per the B0 plan: a DSPARK-strategy export harvests shifted
    /// (all-rows), a mask-fill export keeps the historical dflash arm byte-identical.
    pub fn resolve(cfg: &DflashCfg) -> Self {
        Self::resolve_value(
            std::env::var("MEMRA_DSPARK_HARVEST").ok().as_deref(),
            cfg.strategy_dspark,
        )
    }

    pub fn resolve_value(v: Option<&str>, strategy_dspark: bool) -> Self {
        match v {
            None | Some("") => {
                if strategy_dspark {
                    DsparkHarvest::Dspark
                } else {
                    DsparkHarvest::Dflash
                }
            }
            set => Self::from_env_value(set),
        }
    }

    /// ENV-ONLY parser (no checkpoint census): unset = `Dflash`, the historical arm.
    /// Kept for the explicit-value path of [`Self::resolve_value`] and the seam tests;
    /// round arms resolve through [`Self::resolve`] so the default stays strategy-keyed.
    pub fn from_env_value(v: Option<&str>) -> Self {
        match v {
            None | Some("") | Some("dflash") => DsparkHarvest::Dflash,
            Some("dspark") => DsparkHarvest::Dspark,
            Some(other) => panic!(
                "MEMRA_DSPARK_HARVEST={other}: unknown harvest convention (dflash|dspark); \
                 refusing — a wrong convention verifies every draft slot against a position \
                 the drafter row was not trained for (DSPARK-POSTMORTEM-20260820.md)"
            ),
        }
    }

    /// Manifest/serialized name (the oracle geometry manifest's `harvest` field).
    pub fn name(self) -> &'static str {
        match self {
            DsparkHarvest::Dflash => "dflash",
            DsparkHarvest::Dspark => "dspark",
        }
    }

    pub fn from_name(v: &str) -> Option<Self> {
        match v {
            "dflash" => Some(DsparkHarvest::Dflash),
            "dspark" => Some(DsparkHarvest::Dspark),
            _ => None,
        }
    }

    /// First drafter OUTPUT row consumed as a draft candidate.
    pub fn first_row(self) -> usize {
        match self {
            DsparkHarvest::Dflash => 1,
            DsparkHarvest::Dspark => 0,
        }
    }

    /// Drafted tokens harvested per round from a `b`-row block.
    pub fn n_drafts(self, b: usize) -> usize {
        match self {
            DsparkHarvest::Dflash => b - 1,
            DsparkHarvest::Dspark => b,
        }
    }

    /// The position offset (relative to the round anchor at the block's row 0) that
    /// drafter output row `row` is TRAINED to predict under this convention.
    pub fn trained_offset_of_row(self, row: usize) -> usize {
        match self {
            DsparkHarvest::Dflash => row,
            DsparkHarvest::Dspark => row + 1,
        }
    }
}

/// Checkpoint training-strategy census over the raw config.json text (the loader's
/// minimal-extractor idiom — no json dep in-tree). TRUE iff the export declares the
/// DSPARK strategy: `architectures` naming a DSpark model class (`Qwen3DSparkModel`,
/// the SpecForge OnlineDSparkModel export form) or `dflash_config.projector_type ==
/// "dspark"`. z-lab / OnlineDFlashModel mask-fill exports carry neither signal. Pure,
/// so the census is testable against config fragments without files.
pub fn dspark_strategy_census(txt: &str) -> bool {
    let arch = txt
        .find("\"architectures\"")
        .and_then(|i| {
            let rest = &txt[i..];
            let a = rest.find('[')?;
            let b = rest.find(']')?;
            Some(rest[a..b].contains("DSpark"))
        })
        .unwrap_or(false);
    let proj = txt
        .find("\"projector_type\"")
        .map(|i| {
            let rest = &txt[i..];
            let after = rest.find(':').map(|c| &rest[c + 1..]).unwrap_or("");
            after.trim_start().starts_with("\"dspark\"")
        })
        .unwrap_or(false);
    arch || proj
}

/// Accepted-prefix length of a round's candidates against the trunk's verify argmaxes:
/// `cand[0]` = the round anchor (already decided), `cand[1..]` = the drafts;
/// `vam[j]` = the trunk's argmax prediction for position anchor+j+1. Returns m =
/// number of accepted drafts (`cand[1..=m]` committed, `vam[m]` becomes the next
/// anchor). Pure so the harvest-alignment fixture can exercise it CPU-side.
pub fn dspark_accept_prefix(cand: &[u32], vam: &[u32], vt: usize) -> usize {
    let mut m = 0usize;
    while m < vt - 1 && cand[m + 1] == vam[m] {
        m += 1;
    }
    m
}

/// Verify-window policy for the dspark round (H4, DSPARK-POSTMORTEM-20260820.md §3).
///
/// B2 measured the structural fork: the fixed full-block window (vt=8) buys 95–100%
/// of the sglang accept bank but LOSES wall speed to the reactive ladder everywhere
/// except math — at 0.2–0.5 slot rates, full-block verify pays 5–6 empty rows per
/// round. The confidence policy is the mechanism both leading engines schedule with
/// (sglang v0.5.16 `dspark_planner.py` cumprod survival; vLLM #47808): size EACH
/// round's window from the drafter's own trained accept-rate head, so windows open
/// on confident streaks (math/code) and shrink on bursty text without a 4-round
/// ladder climb.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum DsparkVtPolicy {
    /// The shipped reactive ladder: `vt = (m+2).clamp(3, vt_cap)` per round
    /// (`MEMRA_DFLASH_ADAPT=0` pins vt at `vt_cap` = the fixed-window arm).
    Ladder,
    /// `MEMRA_DSPARK_VT=confidence`: per-round window from cumprod survival of the
    /// confidence head's sigmoid scores, thresholded at `tau`
    /// (`MEMRA_DSPARK_VT_TAU`, default 0.5). Raw sigmoid — no STS sidecar
    /// calibration exists for this export; the postmortem names this the starting
    /// policy.
    Confidence { tau: f32 },
    /// `MEMRA_DSPARK_VT=confidence-slot` (owner directive, 2026-08-20: "take only
    /// high confidence offers"): submit only the longest draft PREFIX whose every
    /// slot clears `tau` on its own sigmoid — the low-confidence tail never enters
    /// verify. Same tau env. vs `Confidence`: if the head's per-row score is the
    /// MARGINAL accept probability (it already sinks with depth), cumprod survival
    /// double-counts the decay and over-truncates; if it is the CONDITIONAL,
    /// per-slot under-truncates. Which statistic the q38 head emits is empirical —
    /// both arms ride the A/B.
    ConfidenceSlot { tau: f32 },
}

impl DsparkVtPolicy {
    /// The served resolution: explicit `MEMRA_DSPARK_VT={ladder|confidence|
    /// confidence-slot}` wins (unknown values REFUSE loudly — a typo silently
    /// reverting the window policy would invalidate an A/B without a trace); UNSET
    /// defaults to **`confidence-slot` at τ = `MEMRA_DSPARK_VT_TAU` (default 0.5)** —
    /// the owner-ratified H4 flip (2026-08-20; cell 2's 4-arm A/B ×5 + cell 3's tau
    /// ladder put the knee at τ=.5 for the slot arm: 94–98% of the fixed-8 accept bank
    /// at wall ≥ the reactive ladder, exactness 11/11 ALL EXACT). Census-keyed per the
    /// capacity-keyed-defaults law: a checkpoint WITHOUT an accept-rate head has no
    /// signal to schedule with, so unset-env resolves to the ladder there (loudly, at
    /// load) instead of panicking on a default; `MEMRA_DFLASH_ADAPT=0` (an explicit
    /// fixed-window request) also keeps the ladder-family arm.
    pub fn resolve(has_confidence_head: bool) -> Self {
        Self::resolve_value(
            std::env::var("MEMRA_DSPARK_VT").ok().as_deref(),
            std::env::var("MEMRA_DSPARK_VT_TAU").ok().as_deref(),
            std::env::var("MEMRA_DFLASH_ADAPT").ok().as_deref(),
            has_confidence_head,
        )
    }

    pub fn resolve_value(
        vt: Option<&str>,
        tau: Option<&str>,
        adapt: Option<&str>,
        has_confidence_head: bool,
    ) -> Self {
        match vt {
            None | Some("") => {
                if adapt == Some("0") || !has_confidence_head {
                    DsparkVtPolicy::Ladder
                } else {
                    // The ratified default rides the SAME tau parse as the explicit
                    // arm (a bad MEMRA_DSPARK_VT_TAU refuses, never silently ignored).
                    Self::from_env_value(Some("confidence-slot"), tau, adapt)
                }
            }
            set => Self::from_env_value(set, tau, adapt),
        }
    }

    /// ENV-ONLY parser (no head census): unset = `Ladder`. Kept for the explicit-value
    /// path of [`Self::resolve_value`] and the policy-gate tests; round arms resolve
    /// through [`Self::resolve`] so the default stays head-census-keyed.
    pub fn from_env_value(vt: Option<&str>, tau: Option<&str>, adapt: Option<&str>) -> Self {
        match vt {
            None | Some("") | Some("ladder") => DsparkVtPolicy::Ladder,
            Some(mode @ ("confidence" | "confidence-slot")) => {
                if adapt == Some("0") {
                    panic!(
                        "MEMRA_DSPARK_VT={mode} together with MEMRA_DFLASH_ADAPT=0 is \
                         contradictory (a pinned fixed window vs a per-round confidence \
                         window); unset one — refuse-on-ambiguity"
                    );
                }
                let tau = tau
                    .map(|t| {
                        t.parse::<f32>()
                            .unwrap_or_else(|_| panic!("MEMRA_DSPARK_VT_TAU={t}: not a float"))
                    })
                    .unwrap_or(0.5);
                assert!(
                    tau > 0.0 && tau < 1.0,
                    "MEMRA_DSPARK_VT_TAU={tau}: confidence threshold must be in (0,1)"
                );
                if mode == "confidence" {
                    DsparkVtPolicy::Confidence { tau }
                } else {
                    DsparkVtPolicy::ConfidenceSlot { tau }
                }
            }
            Some(other) => panic!(
                "MEMRA_DSPARK_VT={other}: unknown verify-window policy \
                 (ladder|confidence|confidence-slot); refusing — a wrong policy \
                 silently reverts the H4 arm (DSPARK-POSTMORTEM-20260820.md)"
            ),
        }
    }

    /// True for every head-scheduled arm (the loops gate the head requirement and
    /// the embedding stash on this).
    pub fn is_confidence(&self) -> bool {
        !matches!(self, DsparkVtPolicy::Ladder)
    }

    /// Size this round's verify window from the head's pre-sigmoid slot scores.
    /// `None` under the ladder (the caller keeps its carried vt).
    pub fn size_window(&self, raws: &[f32], vt_cap: usize) -> Option<usize> {
        match *self {
            DsparkVtPolicy::Ladder => None,
            DsparkVtPolicy::Confidence { tau } => Some(dspark_confidence_vt(raws, tau, vt_cap)),
            DsparkVtPolicy::ConfidenceSlot { tau } => {
                Some(dspark_slot_confidence_vt(raws, tau, vt_cap))
            }
        }
    }
}

/// H4 window sizing (the sglang-planner/vLLM-#47808 mechanism, thresholded): `raws[k]`
/// = the accept-rate head's PRE-sigmoid score for draft slot k+1; survival
/// `S_k = prod_{j<=k} sigmoid(raws[j])`; the window keeps leading slots while
/// `S_k >= tau`. Returns `vt` = 1 (anchor) + kept drafts, clamped to `[2, vt_cap]`:
/// the draft forward is already paid, so at least one draft rides every verify — one
/// extra verify row costs less than a guaranteed empty round. Pure, so the policy's
/// knee is testable CPU-side like `dspark_accept_prefix`.
pub fn dspark_confidence_vt(raws: &[f32], tau: f32, vt_cap: usize) -> usize {
    let mut surv = 1.0f32;
    let mut kept = 0usize;
    for &r in raws {
        surv *= 1.0 / (1.0 + (-r).exp());
        if surv < tau {
            break;
        }
        kept += 1;
    }
    (1 + kept).clamp(2, vt_cap.max(2))
}

/// Owner-directive arm (2026-08-20, "take only high confidence offers"): keep the
/// longest draft PREFIX whose EVERY slot clears `tau` on its own sigmoid — truncate
/// at the first sub-threshold slot, so the low-confidence tail (B2 measured 0.2–0.5
/// slot rates at depth) never enters verify. Prefix truncation is forced by the
/// accept rule anyway (`dspark_accept_prefix` stops at the first miss — a kept slot
/// after a dropped one could never commit); the policy fork vs `dspark_confidence_vt`
/// is only the stopping statistic (per-slot marginal vs cumulative survival). Same
/// floor/cap contract.
pub fn dspark_slot_confidence_vt(raws: &[f32], tau: f32, vt_cap: usize) -> usize {
    let mut kept = 0usize;
    for &r in raws {
        let p = 1.0 / (1.0 + (-r).exp());
        if p < tau {
            break;
        }
        kept += 1;
    }
    (1 + kept).clamp(2, vt_cap.max(2))
}

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,
            strategy_dspark: dspark_strategy_census(&txt),
        };
        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
            };
        // Ratified-default receipts (capacity-keyed-defaults law: the active program is
        // NAMED at load, never inferred from silence). The boot output-sample gate greps
        // these two lines; a run whose log lacks them did not load this code.
        eprintln!(
            "[dspark] harvest={} (checkpoint census strategy_dspark={}, MEMRA_DSPARK_HARVEST {})",
            DsparkHarvest::resolve_value(
                std::env::var("MEMRA_DSPARK_HARVEST").ok().as_deref(),
                cfg.strategy_dspark,
            )
            .name(),
            cfg.strategy_dspark,
            match std::env::var("MEMRA_DSPARK_HARVEST") {
                Ok(v) if !v.is_empty() => "set",
                _ => "unset",
            },
        );
        eprintln!(
            "[dspark] verify-window={:?} (accept-rate head {}, MEMRA_DSPARK_VT {})",
            DsparkVtPolicy::resolve(confidence.is_some()),
            if confidence.is_some() {
                "present"
            } else {
                "ABSENT -> ladder"
            },
            match std::env::var("MEMRA_DSPARK_VT") {
                Ok(v) if !v.is_empty() => "set",
                _ => "unset",
            },
        );
        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)
    }
}

// ================= Engine-bundle slice 1: batched GDN state snapshot ====================
// DSF-ROUNDCOST-20260820 §1.1 measured the dspark round's `cache.snapshot(e)` at 0.67 ms
// native wall — 48 linear layers x {conv, ssm} x (alloc_zeros + memcpy_dtod) of pure
// dispatch serialization, zero kernels. This batcher holds ONE persistent CacheSnapshot
// (buffers allocated on round 1, reused every round — kills the per-round alloc/memset
// churn) plus device pointer tables, so a round's snap is one small H2D table refresh
// (the ssm handles ping-pong per verify row, so live pointers are re-read each round;
// conv handles are rolled in place and never move) + TWO `copy_batch_uniform_f32`
// launches. Bytes, buffers and stream order are identical to `Cache::snapshot`; only the
// dispatch count changes, so acceptance and streams stay bit-identical (E2E-gated).
// `MEMRA_STATE_COPY_BATCH=0` reverts to the legacy per-layer snapshot.

pub(crate) struct DsparkSnapBatch {
    pub(crate) snap: crate::cache::CacheSnapshot,
    /// Linear-attention layer indices, in `conv_table`/`ssm_table` order.
    lin: Vec<usize>,
    /// [src_0..src_{n-1}, dst_0..dst_{n-1}] — live conv states -> snapshot conv buffers.
    conv_table: CudaSlice<u64>,
    ssm_table: CudaSlice<u64>,
    host_ssm: Vec<u64>,
    conv_words: usize,
    ssm_words: usize,
}

impl DsparkSnapBatch {
    /// Build from a fresh full snapshot (this IS round 1's snap — the caller uses
    /// `self.snap` directly after `new`). Returns None when the cache has no linear
    /// layers or their state sizes are non-uniform (a future hybrid shape) — the caller
    /// then stays on the legacy per-layer snapshot rather than copying wrong byte counts.
    pub(crate) fn new(
        e: &Engine,
        cache: &crate::cache::Cache,
    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
        use cudarc::driver::DevicePtr;
        let snap = cache.snapshot(e)?;
        let lin: Vec<usize> = (0..cache.recur.len())
            .filter(|&il| cache.recur[il].is_some())
            .collect();
        if lin.is_empty() {
            return Ok(None);
        }
        let first = cache.recur[lin[0]].as_ref().unwrap();
        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
        for &il in &lin {
            let rl = cache.recur[il].as_ref().unwrap();
            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
                return Ok(None);
            }
        }
        let n = lin.len();
        let mut host_conv = vec![0u64; 2 * n];
        let mut host_ssm = vec![0u64; 2 * n];
        {
            let s = &e.gpu.stream();
            for (k, &il) in lin.iter().enumerate() {
                let rl = cache.recur[il].as_ref().unwrap();
                let (pc, _g0) = rl.conv_state.device_ptr(s);
                let (ps, _g1) = rl.ssm_state.device_ptr(s);
                let (dc, _g2) = snap.conv[il].as_ref().unwrap().device_ptr(s);
                let (ds, _g3) = snap.ssm[il].as_ref().unwrap().device_ptr(s);
                host_conv[k] = pc as u64;
                host_conv[n + k] = dc as u64;
                host_ssm[k] = ps as u64;
                host_ssm[n + k] = ds as u64;
            }
        }
        let conv_table = e.htod_u64(&host_conv)?;
        let ssm_table = e.htod_u64(&host_ssm)?;
        Ok(Some(Self {
            snap,
            lin,
            conv_table,
            ssm_table,
            host_ssm,
            conv_words,
            ssm_words,
        }))
    }

    /// The per-round snap: refresh kv lens/pos host-side (as `snapshot_into` does),
    /// re-read the live ssm handles into the table (gdn ping-pong moves them; the conv
    /// handles and every snapshot dst are stable), then two batched-copy launches.
    pub(crate) fn refresh(
        &mut self,
        e: &Engine,
        cache: &crate::cache::Cache,
    ) -> Result<(), Box<dyn std::error::Error>> {
        use cudarc::driver::DevicePtr;
        for il in 0..cache.kv.len() {
            self.snap.kv_len[il] = cache.kv[il].as_ref().map(|kvl| kvl.len);
        }
        self.snap.pos = cache.pos;
        let n = self.lin.len();
        {
            let s = &e.gpu.stream();
            for (k, &il) in self.lin.iter().enumerate() {
                let rl = cache.recur[il].as_ref().unwrap();
                let (ps, _g) = rl.ssm_state.device_ptr(s);
                self.host_ssm[k] = ps as u64;
            }
        }
        e.htod_u64_into(&self.host_ssm, &mut self.ssm_table)?;
        e.copy_batch_uniform_f32(&self.conv_table, n, self.conv_words)?;
        e.copy_batch_uniform_f32(&self.ssm_table, n, self.ssm_words)?;
        Ok(())
    }
}

// ================= 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();
        // Harvest convention (DSPARK-POSTMORTEM-20260820.md): which drafter output rows
        // become draft candidates. nd = drafts/round; verify carries [anchor, drafts]
        // = up to nd+1 rows. Default = the CHECKPOINT's own strategy census
        // (owner-ratified flip, 2026-08-20); explicit env still wins.
        let harvest = DsparkHarvest::resolve(&draft.cfg);
        let nd = harvest.n_drafts(b);
        let r0 = harvest.first_row();
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(nd + 1)
            .clamp(2, nd + 1);
        let adapt = std::env::var("MEMRA_DFLASH_ADAPT").as_deref() != Ok("0");
        // Verify-window policy (H4, DSPARK-POSTMORTEM-20260820.md): default =
        // confidence-slot tau=.5 when the checkpoint carries an accept-rate head
        // (owner-ratified flip 2026-08-20; cell-3 tau ladder knee) — each round's
        // window is sized from the head's own slot scores, post-draft pre-verify.
        // Head-less checkpoints and MEMRA_DFLASH_ADAPT=0 keep the reactive ladder.
        let vt_policy = DsparkVtPolicy::resolve(draft.confidence.is_some());
        if vt_policy.is_confidence() {
            assert!(
                draft.confidence.is_some(),
                "MEMRA_DSPARK_VT={vt_policy:?} needs a checkpoint with an accept-rate \
                 head (confidence_head.* absent in this export)"
            );
        }
        let mut vt = vt_cap;
        let mut attempted = 0usize;
        let mut accepted = 0usize;
        // Engine-bundle slice 1: persistent batched snapshot (None until round 1; stays
        // None — legacy per-layer snapshot — under MEMRA_STATE_COPY_BATCH=0 or when the
        // batcher declines the cache shape).
        let mut snapb: Option<DsparkSnapBatch> = None;
        let mut snapb_off = !crate::spec::state_copy_batch_on();
        // Engine-bundle slice 2: deferred chain readback needs the resident embed table
        // (verify then embeds chain_d directly). Ladder/stash arms only — the confidence
        // policies size vt from a pre-verify head readback and keep the legacy order.
        let defer_rb = crate::spec::dspark_defer_readback_on() && !vt_policy.is_confidence();
        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
        let embd_gpu = if !defer_rb || crate::spec::spec_host_embd() {
            None
        } else {
            Some(
                self.embd_gpu
                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
            )
        };
        // Engine-bundle slice 3: per-(segment, vt) verify graphs for the linear-layer runs
        // (rides the slice-2 deferred path only — device tokens keep the whole verify off
        // the host). PERSISTENT across generations on the model (rebuilding per call
        // re-captured ~80 graphs per prompt — measured 97.8 -> 79.1 tok/s e2e); the
        // captured bodies are cache-independent: all state reads go through per-round
        // refreshed pointer tables and ctx-owned slabs. None = eager walk, byte-identical.
        let mut vg_guard = self.dspark_vgraphs.lock().unwrap();
        if vg_guard.is_none() && embd_gpu.is_some() && crate::spec::dspark_verify_graph_on() {
            *vg_guard = crate::spec::DsparkVerifyGraphs::new(e, &cache, vt_cap, n_embd)?;
        }
        let vgraphs: &mut Option<crate::spec::DsparkVerifyGraphs> = &mut vg_guard;
        // 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)?;
            // Harvest: logits over rows r0..r0+nd (Dflash: mask rows 1..b-1, fill
            // semantics; Dspark: ALL b rows, shifted semantics — row k predicts
            // anchor+k+1, so col k of `dl` is the draft for position start+k+1).
            let mut rows = e.uninit(nd * n_embd)?;
            {
                let dv = e.view(&dh, b * n_embd);
                let src = dv.slice(r0 * n_embd..(r0 + nd) * n_embd);
                e.copy_view_into(&mut rows, 0, &src, nd * n_embd)?;
            }
            let mut dl = e.matmul(&self.output, &rows, nd)?;
            let markov_on = std::env::var("MEMRA_DFLASH_MARKOV").as_deref() != Ok("0");
            let mut chain_d = e.stream().alloc_zeros::<u32>(nd + 1)?;
            // Confidence policy: stash each slot's markov prev-token embedding (the
            // exact `w1` row the chain gathers) into a [nd, rank] buffer — d2d async,
            // read back beside `rows` in one host sync after the chain.
            let want_conf_emb = vt_policy.is_confidence()
                && draft.confidence.as_ref().is_some_and(|ch| ch.with_markov);
            let mut conf_emb: Option<CudaSlice<f32>> = match (&draft.markov, want_conf_emb) {
                (Some(mk), true) => Some(e.uninit(nd * mk.rank)?),
                (None, true) => unreachable!(
                    "with_markov confidence head without a markov table — the loader forbids it"
                ),
                _ => None,
            };
            if let (Some(mk), true) = (&draft.markov, markov_on) {
                e.set_u32_one(&mut chain_d, last)?;
                for k in 0..nd {
                    let mut f = e.uninit(mk.rank)?;
                    e.gather_row_bf16(&mk.w1_bf16, &chain_d, k, &mut f, mk.rank)?;
                    if let Some(ce) = conf_emb.as_mut() {
                        let fv = e.view(&f, mk.rank);
                        e.copy_view_into(ce, k * mk.rank, &fv, 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 {
                if want_conf_emb {
                    // chain_d[0] must carry the anchor — slot 0's prev token.
                    e.set_u32_one(&mut chain_d, last)?;
                }
                for i in 0..nd {
                    if let (Some(ce), Some(mk)) = (conf_emb.as_mut(), &draft.markov) {
                        let mut f = e.uninit(mk.rank)?;
                        e.gather_row_bf16(&mk.w1_bf16, &chain_d, i, &mut f, mk.rank)?;
                        let fv = e.view(&f, mk.rank);
                        e.copy_view_into(ce, i * mk.rank, &fv, mk.rank)?;
                    }
                    e.argmax_token_device_col(&dl, i, n_vocab, &mut chain_d, i + 1)?;
                }
            }
            e.set_verify_exact(false);
            // 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).
            // Read here (was at the verify site) — slice 2's deferral needs the arm
            // choice before deciding whether the chain readback can move past verify.
            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");
            // Slice 2: under the stash/gate arms with a resident embed table, the chain
            // readback is DEFERRED past verify dispatch and merged with the argmax
            // readback into one sync. The replay arm (CKPT=0) verifies host tokens and
            // keeps the legacy order.
            let deferred = embd_gpu.is_some() && (ckpt_on || ckpt_gate);
            // ---- H4 confidence window: size THIS round's verify from the head ----
            if vt_policy.is_confidence() {
                let ch = draft.confidence.as_ref().expect("asserted at loop entry");
                let (rows_h, emb_h) = match conf_emb.as_ref() {
                    Some(ce) => {
                        let (a, b2) = e.dtoh_pair(&rows, ce)?;
                        (a, Some(b2))
                    }
                    None => (e.dtoh(&rows)?, None),
                };
                let rank = draft.markov.as_ref().map(|m| m.rank).unwrap_or(0);
                let mut raws = Vec::with_capacity(nd);
                for k in 0..nd {
                    let hrow = &rows_h[k * n_embd..(k + 1) * n_embd];
                    let emb = emb_h.as_ref().map(|eh| &eh[k * rank..(k + 1) * rank]);
                    raws.push(ch.raw_score(hrow, emb));
                }
                vt = vt_policy
                    .size_window(&raws, vt_cap)
                    .expect("confidence policies always size the window");
            }
            // Verify candidates: [anchor, draft 1..nd]. Under Dflash this is the
            // historical `block` content; under Dspark it is one longer than the
            // drafter's input block (nd = b drafts + the anchor). Deferred rounds build
            // this after the merged readback — the bytes are identical (chain_d is
            // written before either sync).
            let mut cand: Vec<u32> = Vec::with_capacity(nd + 1);
            if !deferred {
                let chain = e.dtoh_u32(&chain_d)?;
                cand.push(last);
                cand.extend_from_slice(&chain[1..]);
            }
            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();
            // Slice 1: batched snap (one table refresh + two copy launches) with the
            // legacy per-layer snapshot as the kill-switch / non-uniform fallback.
            let mut snap_legacy: Option<crate::cache::CacheSnapshot> = None;
            if !snapb_off && snapb.is_none() {
                snapb = DsparkSnapBatch::new(e, &cache)?;
                snapb_off = snapb.is_none();
            } else if let Some(sb) = snapb.as_mut() {
                sb.refresh(e, &cache)?;
            }
            let snap: &crate::cache::CacheSnapshot = match snapb.as_ref() {
                Some(sb) => &sb.snap,
                None => {
                    snap_legacy = Some(cache.snapshot(e)?);
                    snap_legacy.as_ref().unwrap()
                }
            };
            let _ = &snap_legacy;
            ns_snap += clock(stats, e).duration_since(t1).as_nanos() as u64;
            let t2 = std::time::Instant::now();
            // Slice 3: the tap-sink buffer is persistent per vt in the graphs ctx
            // (captured segments bake its address); fully rewritten by every verify.
            let tap_buf = match vgraphs.as_mut().and_then(|g| g.tap_bufs.remove(&vt)) {
                Some(buf) => buf,
                None => e.uninit(vt * n_taps * n_embd)?,
            };
            cache.dflash_taps = Some(DflashTapSink {
                layer_ids: c.target_layer_ids.clone(),
                buf: tap_buf,
                hidden: n_embd,
                t: vt,
                base: 0,
            });
            // The whole fallible verify window runs inside a closure so the Err path can
            // return the sink buffer to the ctx pool before propagating (v0.98 review
            // carry-over): five `?`s span the window, and an early return would drop
            // `cache.dflash_taps` — freeing the buffer whose ADDRESS the model-persistent
            // captured graphs bake, so the next generation's replayed tap copies would
            // write freed memory. The never-orphan invariant below now holds on EVERY
            // exit, not just the EOS/budget break.
            let verify_res = (|cache: &mut crate::cache::Cache,
                               cand: &mut Vec<u32>,
                               vgraphs: &mut Option<crate::spec::DsparkVerifyGraphs>|
             -> Result<
                (Vec<u32>, Option<crate::spec::DsparkVerifyCkpt>),
                Box<dyn std::error::Error>,
            > {
                if deferred {
                    // Slice 2: verify embeds the DEVICE chain (cand layout by construction:
                    // chain_d[0] = anchor, chain_d[1..] = drafts), then ONE host sync reads
                    // chain + verify argmaxes together — the host dispatched snap + all of
                    // verify while the draft was still executing.
                    let g = embd_gpu.expect("deferred implies resident embed");
                    let (am_d, vck) = self.dspark_verify_t_am_ckpt_dev(
                        e,
                        &chain_d,
                        vt,
                        start,
                        cache,
                        (g, embd_qt, embd_rb),
                        vgraphs.as_mut(),
                    )?;
                    let ch = e.stream().clone_dtoh(&chain_d)?;
                    let am = e.stream().clone_dtoh(&am_d)?;
                    e.stream().synchronize()?;
                    cand.push(last);
                    cand.extend_from_slice(&ch[1..]);
                    Ok((am, Some(vck)))
                } else if ckpt_on || ckpt_gate {
                    let (vam, vck) = self.dspark_verify_t_am_ckpt(e, &cand[..vt], start, cache)?;
                    Ok((vam, Some(vck)))
                } else {
                    Ok((self.dspark_verify_t_am(e, &cand[..vt], start, cache)?, None))
                }
            })(&mut cache, &mut cand, vgraphs);
            let (vam, vck) = match verify_res {
                Ok(v) => v,
                Err(err) => {
                    if let (Some(g), Some(taps)) = (vgraphs.as_mut(), cache.dflash_taps.take()) {
                        g.tap_bufs.insert(vt, taps.buf);
                    }
                    return Err(err);
                }
            };
            let taps = cache.dflash_taps.take().unwrap();
            // Return the tap buffer to the ctx pool IMMEDIATELY — an EOS/budget break
            // between accept and ingest must never orphan an address the captured
            // graphs bake (the next generation would alloc a fresh buffer and the
            // replayed tap copies would write freed memory). Ingest reads it borrowed.
            let tap_local: Option<CudaSlice<f32>> = match vgraphs.as_mut() {
                Some(g) => {
                    g.tap_bufs.insert(vt, taps.buf);
                    None
                }
                None => Some(taps.buf),
            };
            let tap_ref: &CudaSlice<f32> = match &tap_local {
                Some(b) => b,
                None => &vgraphs.as_ref().expect("ctx present above").tap_bufs[&vt],
            };
            ns_verify += clock(stats, e).duration_since(t2).as_nanos() as u64;

            // ---- accept ----
            let m = dspark_accept_prefix(&cand, &vam, vt);
            attempted += vt - 1;
            accepted += m;
            out.push(last);
            if eos.contains(&last) {
                break 'outer;
            }
            for &dt in &cand[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();
            // Slice 3: rounds whose linear column stash lives in the graphs ctx's slabs
            // commit through the slab twin (same semantics, slab-addressed sources).
            let slab_commit = vgraphs.as_ref().map(|g| g.round_slab).unwrap_or(false);
            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).
                    if slab_commit {
                        self.dspark_commit_prefix_slab(
                            e,
                            &mut cache,
                            snap,
                            vgraphs.as_ref().expect("slab_commit implies ctx"),
                            keep,
                        )?;
                    } else {
                        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, &cand[..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 slab_commit {
                    // STASH ARM, slab twin (slice 3): same restore, slab-addressed.
                    self.dspark_commit_prefix_slab(
                        e,
                        &mut cache,
                        snap,
                        vgraphs.as_ref().expect("slab_commit implies ctx"),
                        keep,
                    )?;
                } 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, &cand[..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(tap_ref, 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;
            // Ladder update only — under the confidence policies vt is recomputed
            // from the head every round, post-draft pre-verify.
            if !vt_policy.is_confidence() && 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,
    /// Engine-bundle slice 1: persistent batched snapshot (buffers + pointer tables live
    /// with the session so bursts reuse them). None until the first round; stays None —
    /// legacy per-layer snapshot — when `snapb_off`.
    snapb: Option<DsparkSnapBatch>,
    snapb_off: 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()?;
        // Verify carries [anchor, drafts] = up to n_drafts+1 rows (harvest-dependent;
        // DSPARK-POSTMORTEM-20260820.md; default = checkpoint strategy census).
        let nd = DsparkHarvest::resolve(&draft.cfg).n_drafts(b);
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(nd + 1)
            .clamp(2, nd + 1);
        Ok(DsparkSpecSession {
            cache,
            dkv,
            last,
            ctx_len: tp,
            vt: vt_cap,
            rounds: 0,
            max_ctx,
            done: false,
            snapb: None,
            snapb_off: !crate::spec::state_copy_batch_on(),
        })
    }

    /// 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();
        // Harvest convention (DSPARK-POSTMORTEM-20260820.md) — identical to the bin arm
        // (default = checkpoint strategy census; owner-ratified flip 2026-08-20).
        let harvest = DsparkHarvest::resolve(&draft.cfg);
        let nd = harvest.n_drafts(b);
        let r0 = harvest.first_row();
        let vt_cap: usize = std::env::var("MEMRA_DFLASH_VERIFY_T")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(nd + 1)
            .clamp(2, nd + 1);
        let adapt = std::env::var("MEMRA_DFLASH_ADAPT").as_deref() != Ok("0");
        // Verify-window policy (H4, DSPARK-POSTMORTEM-20260820.md) — identical to the
        // bin arm: default = confidence-slot tau=.5 on a head-carrying checkpoint
        // (owner-ratified flip 2026-08-20); head-less/ADAPT=0 keep the ladder.
        let vt_policy = DsparkVtPolicy::resolve(draft.confidence.is_some());
        if vt_policy.is_confidence() {
            assert!(
                draft.confidence.is_some(),
                "MEMRA_DSPARK_VT={vt_policy:?} needs a checkpoint with an accept-rate \
                 head (confidence_head.* absent in this export)"
            );
        }
        let mut out: Vec<u32> = Vec::with_capacity(burst_target + b);
        let mut drafted = 0usize;
        let mut accepted_n = 0usize;
        // Engine-bundle slice 2 — identical to the bin arm: deferred chain readback under
        // the stash arm with a resident embed table (ladder policy only).
        let defer_rb = crate::spec::dspark_defer_readback_on() && !vt_policy.is_confidence();
        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
        let embd_gpu = if !defer_rb || crate::spec::spec_host_embd() {
            None
        } else {
            Some(
                self.embd_gpu
                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
            )
        };
        'outer: while out.len() < burst_target && !sess.done {
            let start = sess.cache.pos;
            if start + nd + 1 > sess.max_ctx {
                sess.done = true;
                break;
            }
            sess.rounds += 1;
            let mut 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)?;
            // Harvest: logits over rows r0..r0+nd (see the bin arm / the postmortem).
            let mut rows = e.uninit(nd * n_embd)?;
            {
                let dv = e.view(&dh, b * n_embd);
                let src = dv.slice(r0 * n_embd..(r0 + nd) * n_embd);
                e.copy_view_into(&mut rows, 0, &src, nd * n_embd)?;
            }
            let mut dl = e.matmul(&self.output, &rows, nd)?;
            let markov_on = std::env::var("MEMRA_DFLASH_MARKOV").as_deref() != Ok("0");
            let mut chain_d = e.stream().alloc_zeros::<u32>(nd + 1)?;
            // Confidence policy: stash markov prev-token embeddings d2d during the
            // chain, one host readback after — identical to the bin arm.
            let want_conf_emb = vt_policy.is_confidence()
                && draft.confidence.as_ref().is_some_and(|ch| ch.with_markov);
            let mut conf_emb: Option<CudaSlice<f32>> = match (&draft.markov, want_conf_emb) {
                (Some(mk), true) => Some(e.uninit(nd * mk.rank)?),
                (None, true) => unreachable!(
                    "with_markov confidence head without a markov table — the loader forbids it"
                ),
                _ => None,
            };
            if let (Some(mk), true) = (&draft.markov, markov_on) {
                e.set_u32_one(&mut chain_d, sess.last)?;
                for k in 0..nd {
                    let mut f = e.uninit(mk.rank)?;
                    e.gather_row_bf16(&mk.w1_bf16, &chain_d, k, &mut f, mk.rank)?;
                    if let Some(ce) = conf_emb.as_mut() {
                        let fv = e.view(&f, mk.rank);
                        e.copy_view_into(ce, k * mk.rank, &fv, 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 {
                if want_conf_emb {
                    // chain_d[0] must carry the anchor — slot 0's prev token.
                    e.set_u32_one(&mut chain_d, sess.last)?;
                }
                for i in 0..nd {
                    if let (Some(ce), Some(mk)) = (conf_emb.as_mut(), &draft.markov) {
                        let mut f = e.uninit(mk.rank)?;
                        e.gather_row_bf16(&mk.w1_bf16, &chain_d, i, &mut f, mk.rank)?;
                        let fv = e.view(&f, mk.rank);
                        e.copy_view_into(ce, i * mk.rank, &fv, mk.rank)?;
                    }
                    e.argmax_token_device_col(&dl, i, n_vocab, &mut chain_d, i + 1)?;
                }
            }
            e.set_verify_exact(false);
            // Slice 2: arm choice read before the chain readback (see the bin arm; the
            // serve arm has no CKPT_GATE oracle — the bin arm carries it).
            let ckpt_on = std::env::var("MEMRA_DSPARK_CKPT").as_deref() != Ok("0");
            let deferred = embd_gpu.is_some() && ckpt_on;
            // ---- H4 confidence window: size THIS round's verify from the head ----
            if vt_policy.is_confidence() {
                let ch = draft.confidence.as_ref().expect("asserted at burst entry");
                let (rows_h, emb_h) = match conf_emb.as_ref() {
                    Some(ce) => {
                        let (a, b2) = e.dtoh_pair(&rows, ce)?;
                        (a, Some(b2))
                    }
                    None => (e.dtoh(&rows)?, None),
                };
                let rank = draft.markov.as_ref().map(|m| m.rank).unwrap_or(0);
                let mut raws = Vec::with_capacity(nd);
                for k in 0..nd {
                    let hrow = &rows_h[k * n_embd..(k + 1) * n_embd];
                    let emb = emb_h.as_ref().map(|eh| &eh[k * rank..(k + 1) * rank]);
                    raws.push(ch.raw_score(hrow, emb));
                }
                vt = vt_policy
                    .size_window(&raws, vt_cap)
                    .expect("confidence policies always size the window");
            }
            let mut cand: Vec<u32> = Vec::with_capacity(nd + 1);
            if !deferred {
                let chain = e.dtoh_u32(&chain_d)?;
                cand.push(sess.last);
                cand.extend_from_slice(&chain[1..]);
            }

            // ---- snapshot, then verify t=vt (ckpt stash default; oracle arms kept) ----
            // Slice 1: batched snap (see DsparkSnapBatch) with the legacy per-layer
            // snapshot as the kill-switch / non-uniform fallback.
            let mut snap_legacy: Option<crate::cache::CacheSnapshot> = None;
            if !sess.snapb_off && sess.snapb.is_none() {
                sess.snapb = DsparkSnapBatch::new(e, &sess.cache)?;
                sess.snapb_off = sess.snapb.is_none();
            } else if let Some(sb) = sess.snapb.as_mut() {
                sb.refresh(e, &sess.cache)?;
            }
            let snap: &crate::cache::CacheSnapshot = match sess.snapb.as_ref() {
                Some(sb) => &sb.snap,
                None => {
                    snap_legacy = Some(sess.cache.snapshot(e)?);
                    snap_legacy.as_ref().unwrap()
                }
            };
            let _ = &snap_legacy;
            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 (vam, vck) = if deferred {
                // Slice 2: device-token verify + ONE merged readback (see the bin arm).
                let g = embd_gpu.expect("deferred implies resident embed");
                // Slice 3 stays bin-arm-only for now: session lifetime (per-request
                // caches, capture storms) needs the cache-reuse-pool design first.
                let (am_d, vck) = self.dspark_verify_t_am_ckpt_dev(
                    e,
                    &chain_d,
                    vt,
                    start,
                    &mut sess.cache,
                    (g, embd_qt, embd_rb),
                    None,
                )?;
                let ch = e.stream().clone_dtoh(&chain_d)?;
                let am = e.stream().clone_dtoh(&am_d)?;
                e.stream().synchronize()?;
                cand.push(sess.last);
                cand.extend_from_slice(&ch[1..]);
                (am, Some(vck))
            } else if ckpt_on {
                let (vam, vck) =
                    self.dspark_verify_t_am_ckpt(e, &cand[..vt], start, &mut sess.cache)?;
                (vam, Some(vck))
            } else {
                (
                    self.dspark_verify_t_am(e, &cand[..vt], start, &mut sess.cache)?,
                    None,
                )
            };
            let taps = sess.cache.dflash_taps.take().unwrap();

            // ---- accept ----
            let m = dspark_accept_prefix(&cand, &vam, vt);
            drafted += vt - 1;
            accepted_n += m;
            out.push(sess.last);
            if eos.contains(&sess.last) {
                sess.done = true;
                break 'outer;
            }
            for &dt in &cand[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, &cand[..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;
            // Ladder update only — the confidence policies recompute vt from the
            // head every round, post-draft pre-verify; their carry just keeps
            // observability (sess.vt = the last confidence-sized window).
            if vt_policy.is_confidence() {
                sess.vt = vt;
            } else if adapt {
                sess.vt = (m + 2).clamp(3, vt_cap);
            }
        }
        Ok((out, drafted, accepted_n))
    }
}

// ================= Harvest-convention gate (CPU; DSPARK-POSTMORTEM-20260820.md) =========
// The parity oracle is row-count-agnostic (it reproduces the markov MODULE on whatever
// rows it is fed) and the E2E gate is harvest-independent (verify-side truth), so
// NEITHER can catch a wrong row->position mapping — that blindness is how the q38
// misalignment shipped. These tests pin the convention itself as logic the round
// consumes, so a mutation back to the mask-fill harvest under the Dspark variant fails
// HERE, naming the convention.
#[cfg(test)]
mod dspark_harvest_tests {
    use super::{DsparkHarvest, DsparkVtPolicy, dspark_accept_prefix, dspark_strategy_census};

    const B: usize = 7; // q38 arm-a block_size

    #[test]
    fn dspark_strategy_requires_shifted_harvest() {
        let h = DsparkHarvest::Dspark;
        assert_eq!(
            h.first_row(),
            0,
            "DSPARK-strategy checkpoints (SpecForge OnlineDSparkModel, \
             training.strategy=dspark — the q38 arm-a export) supervise ALL rows with \
             SHIFTED labels: label_offsets = arange(1, block_size+1), i.e. the ANCHOR \
             row's output is draft 1 (specforge/algorithms/common/\
             dflash_family_model.py:816; sglang v0.5.17 dspark_draft.py:248,260). \
             Harvesting from row 1 re-opens the DSPARK-POSTMORTEM-20260820 slot \
             misalignment (accept 2.9 -> 1.43)."
        );
        assert_eq!(
            h.n_drafts(B),
            B,
            "DSpark harvests gamma = block_size drafts per round (sglang \
             dspark_config.py:269, verify_num_draft_tokens = gamma+1); b-1 is the \
             DFlash mask-fill count and drops the best-trained slot \
             (DSPARK-POSTMORTEM-20260820.md §3-H1)."
        );
        for row in 0..B {
            assert_eq!(
                h.trained_offset_of_row(row),
                row + 1,
                "OnlineDSparkModel trains row k to predict anchor+k+1 \
                 (dflash_family_model.py:816); a same-position (mask-fill) mapping \
                 here verifies every slot one position early — the postmortem's \
                 collapse."
            );
        }
    }

    #[test]
    fn dflash_strategy_keeps_mask_fill_harvest() {
        // Guards the reverse mutation: z-lab dflash checkpoints (the gemma arm) are
        // mask-fill — row k FILLS anchor+k, the anchor row is loss-excluded
        // (dflash_family_model.py:453-472). Shifting THEM would break the gemma arm.
        let h = DsparkHarvest::Dflash;
        assert_eq!(h.first_row(), 1, "DFlash drafts start at mask row 1");
        assert_eq!(h.n_drafts(B), B - 1, "DFlash harvests block_size-1 drafts");
        for row in 1..B {
            assert_eq!(h.trained_offset_of_row(row), row);
        }
    }

    #[test]
    fn every_candidate_verifies_the_position_its_row_was_trained_for() {
        // The round's invariant: draft candidate i (1-based; verified against the
        // trunk's prediction for anchor+i) is filled from drafter output row
        // first_row + i - 1. Alignment == that row was TRAINED for offset i.
        for h in [DsparkHarvest::Dflash, DsparkHarvest::Dspark] {
            for i in 1..=h.n_drafts(B) {
                let row = h.first_row() + i - 1;
                assert_eq!(
                    h.trained_offset_of_row(row),
                    i,
                    "{h:?}: candidate {i} rides row {row}, which is trained for \
                     offset {} — harvest misaligned",
                    h.trained_offset_of_row(row)
                );
            }
        }
    }

    #[test]
    fn env_seam_parses_and_refuses() {
        assert_eq!(
            DsparkHarvest::from_env_value(None),
            DsparkHarvest::Dflash,
            "the ENV-ONLY parser keeps the historical arm; the ratified strategy-keyed \
             default lives in resolve_value (checkpoint census), not here"
        );
        assert_eq!(
            DsparkHarvest::from_env_value(Some("dspark")),
            DsparkHarvest::Dspark
        );
        assert_eq!(
            DsparkHarvest::from_env_value(Some("dflash")),
            DsparkHarvest::Dflash
        );
        assert!(
            std::panic::catch_unwind(|| DsparkHarvest::from_env_value(Some("shifted"))).is_err(),
            "unknown harvest values must REFUSE, not default"
        );
        assert_eq!(
            DsparkHarvest::from_name("dspark"),
            Some(DsparkHarvest::Dspark)
        );
        assert_eq!(
            DsparkHarvest::from_name("dflash"),
            Some(DsparkHarvest::Dflash)
        );
        assert_eq!(DsparkHarvest::from_name("mask-fill"), None);
    }

    /// The owner-ratified default flips (2026-08-20). Each assertion names its
    /// evidence; mutating either resolve back to the old default fails these.
    #[test]
    fn ratified_default_harvest_is_strategy_keyed() {
        // DSPARK-strategy checkpoint + unset env = the shifted harvest (B1: accept
        // 1.38->2.41 agentic / 1.53->3.66 math, E2E ALL EXACT x5, interleaved x5).
        assert_eq!(
            DsparkHarvest::resolve_value(None, true),
            DsparkHarvest::Dspark,
            "owner-ratified 2026-08-20: unset env defaults a DSPARK-strategy \
             checkpoint to the shifted harvest (DSPARK-POSTMORTEM-20260820.md B1)"
        );
        // mask-fill checkpoint + unset env = the historical arm, byte-identical.
        assert_eq!(
            DsparkHarvest::resolve_value(None, false),
            DsparkHarvest::Dflash
        );
        assert_eq!(
            DsparkHarvest::resolve_value(Some(""), false),
            DsparkHarvest::Dflash
        );
        // Explicit env overrides the census in BOTH directions (the A/B seam).
        assert_eq!(
            DsparkHarvest::resolve_value(Some("dflash"), true),
            DsparkHarvest::Dflash
        );
        assert_eq!(
            DsparkHarvest::resolve_value(Some("dspark"), false),
            DsparkHarvest::Dspark
        );
        // Unknown values still REFUSE through the resolve path.
        assert!(
            std::panic::catch_unwind(|| DsparkHarvest::resolve_value(Some("shifted"), true))
                .is_err()
        );
    }

    #[test]
    fn strategy_census_reads_the_checkpoint_not_the_env() {
        // The q38 arm-a export shape: both signals present.
        let q38 = r#"{"architectures": ["Qwen3DSparkModel"], "block_size": 7,
            "dflash_config": {"projector_type": "dspark", "markov_rank": 256}}"#;
        assert!(dspark_strategy_census(q38));
        // Either signal alone suffices.
        assert!(dspark_strategy_census(
            r#"{"architectures": ["Qwen3DSparkModel"]}"#
        ));
        assert!(dspark_strategy_census(
            r#"{"dflash_config": {"projector_type": "dspark"}}"#
        ));
        // A mask-fill DFlash export carries neither -> historical default.
        let dflash = r#"{"architectures": ["Qwen3DFlashModel"],
            "dflash_config": {"attention_mode": "gqa"}}"#;
        assert!(!dspark_strategy_census(dflash));
        assert!(!dspark_strategy_census("{}"));
    }

    #[test]
    fn ratified_default_vt_is_confidence_slot_tau_half() {
        // Head-carrying checkpoint + unset env = confidence-slot tau=.5 (H4 cell 3:
        // the tau ladder's knee; cell 2: 93.9%/97.7% of fixed-8 accept at wall >=
        // the reactive ladder, exactness 11/11 ALL EXACT).
        assert_eq!(
            DsparkVtPolicy::resolve_value(None, None, None, true),
            DsparkVtPolicy::ConfidenceSlot { tau: 0.5 },
            "owner-ratified 2026-08-20: unset MEMRA_DSPARK_VT defaults to \
             confidence-slot tau=.5 on a head-carrying checkpoint (H4 cells 2-3)"
        );
        // tau env still steers the default arm (and a bad tau still refuses).
        assert_eq!(
            DsparkVtPolicy::resolve_value(None, Some("0.35"), None, true),
            DsparkVtPolicy::ConfidenceSlot { tau: 0.35 }
        );
        assert!(
            std::panic::catch_unwind(|| DsparkVtPolicy::resolve_value(
                None,
                Some("nan-ish"),
                None,
                true
            ))
            .is_err()
        );
        // Census: no accept-rate head -> nothing to schedule with -> ladder.
        assert_eq!(
            DsparkVtPolicy::resolve_value(None, None, None, false),
            DsparkVtPolicy::Ladder
        );
        // MEMRA_DFLASH_ADAPT=0 is an explicit fixed-window request: honored.
        assert_eq!(
            DsparkVtPolicy::resolve_value(None, None, Some("0"), true),
            DsparkVtPolicy::Ladder
        );
        // Explicit values keep their exact prior semantics through resolve.
        assert_eq!(
            DsparkVtPolicy::resolve_value(Some("ladder"), None, None, true),
            DsparkVtPolicy::Ladder
        );
        assert_eq!(
            DsparkVtPolicy::resolve_value(Some("confidence"), Some("0.35"), None, true),
            DsparkVtPolicy::Confidence { tau: 0.35 }
        );
        // Explicit confidence mode with ADAPT=0 stays a refusal.
        assert!(
            std::panic::catch_unwind(|| DsparkVtPolicy::resolve_value(
                Some("confidence-slot"),
                None,
                Some("0"),
                true
            ))
            .is_err()
        );
    }

    /// End-to-end alignment fixture in miniature: a mock drafter whose row r argmaxes
    /// to token BASE + (its trained offset under the DSPARK strategy), and a mock trunk
    /// whose prediction for anchor+j is BASE + j. The DSpark harvest accepts the whole
    /// block; feeding the same drafter through the mask-fill harvest accepts ZERO —
    /// the postmortem's collapse reproduced as pure logic.
    #[test]
    fn dspark_trained_rows_through_mask_fill_harvest_accept_nothing() {
        const BASE: u32 = 1000;
        let anchor: u32 = BASE; // token at the round anchor position (offset 0)
        // trunk verify argmaxes: vam[j] = prediction for anchor offset j+1
        let vam: Vec<u32> = (1..=B as u32 + 1).map(|j| BASE + j).collect();
        // drafter rows trained under the DSPARK strategy: row r predicts offset r+1
        let dspark_trained_row_argmax =
            |r: usize| BASE + DsparkHarvest::Dspark.trained_offset_of_row(r) as u32;

        // Correct (shifted) harvest: candidate i <- row i-1.
        let h = DsparkHarvest::Dspark;
        let mut cand = vec![anchor];
        for i in 1..=h.n_drafts(B) {
            cand.push(dspark_trained_row_argmax(h.first_row() + i - 1));
        }
        let vt = h.n_drafts(B) + 1;
        assert_eq!(
            dspark_accept_prefix(&cand, &vam, vt),
            vt - 1,
            "aligned harvest must accept the full block"
        );

        // Mask-fill harvest of the SAME dspark-trained drafter: candidate i <- row i,
        // which was trained for offset i+1 — every slot one position late.
        let wrong = DsparkHarvest::Dflash;
        let mut cand_wrong = vec![anchor];
        for i in 1..=wrong.n_drafts(B) {
            cand_wrong.push(dspark_trained_row_argmax(wrong.first_row() + i - 1));
        }
        let vt_wrong = wrong.n_drafts(B) + 1;
        assert_eq!(
            dspark_accept_prefix(&cand_wrong, &vam, vt_wrong),
            0,
            "mask-fill harvest of a dspark-trained drafter verifies every slot against \
             a position the row was not trained for (DSPARK-POSTMORTEM-20260820.md)"
        );
    }
}

// ================= Verify-window policy gate (CPU; H4, DSPARK-POSTMORTEM-20260820.md) ===
// Pins the confidence-vt semantics as logic the round consumes: cumprod survival over
// sigmoid scores, thresholded, anchor + kept drafts, floor 2 / cap vt_cap — and the env
// seam's refuse-on-ambiguity. Mutating the policy (per-slot threshold instead of
// survival, off-by-one on the anchor, silent unknown-value fallback) fails HERE.
#[cfg(test)]
mod dspark_vt_tests {
    use super::{ConfidenceHead, DsparkVtPolicy, dspark_confidence_vt, dspark_slot_confidence_vt};

    /// Pre-sigmoid logit for a target probability: sigmoid(logit(p)) == p.
    fn logit(p: f32) -> f32 {
        (p / (1.0 - p)).ln()
    }

    #[test]
    fn confidence_vt_is_cumprod_survival_not_per_slot_threshold() {
        // sigmoids = [0.9, 0.8, 0.9, ...]: every PER-SLOT score clears tau=0.5, but
        // cumulative survival sinks below it at slot 6 (0.9, 0.72, 0.648, 0.583,
        // 0.525, then 0.472 < 0.5) — the window must stop where the EXPECTED
        // accepted-prefix stops paying, not where a slot looks locally fine.
        let raws: Vec<f32> = [0.9, 0.8, 0.9, 0.9, 0.9, 0.9, 0.9]
            .iter()
            .map(|&p| logit(p))
            .collect();
        assert_eq!(
            dspark_confidence_vt(&raws, 0.5, 8),
            6,
            "keeps 5 drafts + anchor"
        );
        // Tighter threshold closes the window sooner; looser opens it to the cap.
        assert_eq!(
            dspark_confidence_vt(&raws, 0.7, 8),
            3,
            "tau=0.7 keeps 2 drafts"
        );
        assert_eq!(
            dspark_confidence_vt(&raws, 0.05, 8),
            8,
            "tau→0 = full block"
        );
    }

    #[test]
    fn slot_arm_truncates_at_first_low_confidence_slot() {
        // Owner directive (2026-08-20): submit only the longest prefix whose EVERY
        // slot clears tau on its own sigmoid. On the survival test's raws
        // ([0.9, 0.8, 0.9 x5], tau=0.5) every slot clears per-slot, so the slot arm
        // opens the full block where survival stopped at 6 — the two stopping
        // statistics must stay distinct arms.
        let raws: Vec<f32> = [0.9, 0.8, 0.9, 0.9, 0.9, 0.9, 0.9]
            .iter()
            .map(|&p| logit(p))
            .collect();
        assert_eq!(dspark_slot_confidence_vt(&raws, 0.5, 8), 8);
        assert_eq!(dspark_confidence_vt(&raws, 0.5, 8), 6);
        // A low-confidence tail never enters verify: [0.9, 0.9, 0.3, 0.9, ...]
        // truncates at slot 3 REGARDLESS of the confident slots behind it — a kept
        // slot after a dropped one could never commit (prefix accept rule).
        let tail: Vec<f32> = [0.9, 0.9, 0.3, 0.9, 0.9, 0.9, 0.9]
            .iter()
            .map(|&p| logit(p))
            .collect();
        assert_eq!(
            dspark_slot_confidence_vt(&tail, 0.5, 8),
            3,
            "2 drafts + anchor"
        );
        // Tighter tau keeps less.
        assert_eq!(
            dspark_slot_confidence_vt(&tail, 0.95, 8),
            2,
            "floor at tau=0.95"
        );
    }

    #[test]
    fn confidence_vt_floor_and_cap() {
        // A hopeless round still verifies ONE draft (the draft forward is paid;
        // vt=1 would guarantee an empty round at the same cost class).
        let cold: Vec<f32> = [0.1f32, 0.1, 0.1].iter().map(|&p| logit(p)).collect();
        assert_eq!(
            dspark_confidence_vt(&cold, 0.5, 8),
            2,
            "floor = anchor + 1 draft"
        );
        assert_eq!(
            dspark_slot_confidence_vt(&cold, 0.5, 8),
            2,
            "slot arm same floor"
        );
        // The MEMRA_DFLASH_VERIFY_T cap still binds a confident round.
        let hot: Vec<f32> = vec![logit(0.99); 7];
        assert_eq!(dspark_confidence_vt(&hot, 0.5, 5), 5, "vt_cap binds");
        assert_eq!(
            dspark_confidence_vt(&hot, 0.5, 8),
            8,
            "full block when confident"
        );
        assert_eq!(
            dspark_slot_confidence_vt(&hot, 0.5, 5),
            5,
            "slot arm same cap"
        );
        // No scores (defensive): floor.
        assert_eq!(dspark_confidence_vt(&[], 0.5, 8), 2);
        assert_eq!(dspark_slot_confidence_vt(&[], 0.5, 8), 2);
    }

    #[test]
    fn vt_policy_env_seam_parses_and_refuses() {
        assert_eq!(
            DsparkVtPolicy::from_env_value(None, None, None),
            DsparkVtPolicy::Ladder,
            "default stays the shipped ladder — the H4 arm is opt-in"
        );
        assert_eq!(
            DsparkVtPolicy::from_env_value(Some(""), None, None),
            DsparkVtPolicy::Ladder
        );
        assert_eq!(
            DsparkVtPolicy::from_env_value(Some("ladder"), None, Some("0")),
            DsparkVtPolicy::Ladder,
            "ladder + ADAPT=0 = the fixed-window arm, untouched"
        );
        assert_eq!(
            DsparkVtPolicy::from_env_value(Some("confidence"), None, None),
            DsparkVtPolicy::Confidence { tau: 0.5 },
            "tau defaults to 0.5 (raw sigmoid, no STS sidecar — postmortem §3-H4)"
        );
        assert_eq!(
            DsparkVtPolicy::from_env_value(Some("confidence"), Some("0.35"), Some("1")),
            DsparkVtPolicy::Confidence { tau: 0.35 }
        );
        assert_eq!(
            DsparkVtPolicy::from_env_value(Some("confidence-slot"), Some("0.6"), None),
            DsparkVtPolicy::ConfidenceSlot { tau: 0.6 },
            "the owner-directive per-slot arm parses with the same tau env"
        );
        assert!(
            std::panic::catch_unwind(|| DsparkVtPolicy::from_env_value(
                Some("confidence-slot"),
                None,
                Some("0")
            ))
            .is_err(),
            "confidence-slot + MEMRA_DFLASH_ADAPT=0 must REFUSE like confidence"
        );
        assert!(
            std::panic::catch_unwind(|| DsparkVtPolicy::from_env_value(Some("static"), None, None))
                .is_err(),
            "unknown policy values must REFUSE, not default — a typo silently \
             reverting the window policy invalidates an A/B"
        );
        assert!(
            std::panic::catch_unwind(|| DsparkVtPolicy::from_env_value(
                Some("confidence"),
                None,
                Some("0")
            ))
            .is_err(),
            "confidence + MEMRA_DFLASH_ADAPT=0 is contradictory and must REFUSE"
        );
        for bad in ["0", "1", "1.5", "-0.1", "nan"] {
            assert!(
                std::panic::catch_unwind(|| DsparkVtPolicy::from_env_value(
                    Some("confidence"),
                    Some(bad),
                    None
                ))
                .is_err(),
                "tau={bad} must REFUSE (survival threshold lives in (0,1))"
            );
        }
    }

    #[test]
    fn raw_score_matches_the_parity_gate_dot() {
        // The head is a raw linear proj over [hidden ; markov_prev_embedding] + b —
        // the exact stage-5 contract in dspark_q38_parity.rs.
        let ch = ConfidenceHead {
            w: vec![0.5, -1.0, 2.0, 0.25, -0.5],
            b: 0.125,
            in_dim: 5,
            with_markov: true,
        };
        let hidden = [1.0f32, 2.0, 3.0];
        let emb = [4.0f32, 8.0];
        let want = 0.125 + 0.5 * 1.0 - 1.0 * 2.0 + 2.0 * 3.0 + 0.25 * 4.0 - 0.5 * 8.0;
        assert_eq!(ch.raw_score(&hidden, Some(&emb)), want);
        let ch_plain = ConfidenceHead {
            w: vec![0.5, -1.0, 2.0],
            b: -0.25,
            in_dim: 3,
            with_markov: false,
        };
        let want_plain = -0.25 + 0.5 * 1.0 - 1.0 * 2.0 + 2.0 * 3.0;
        assert_eq!(ch_plain.raw_score(&hidden, None), want_plain);
    }
}