Skip to main content

memra_engine/
spec.rs

1//! Qwen3.5 MTP (NextN) greedy speculative decode (research/mtp/MTP-PLAN.md §A/§B/§C/§D).
2//!
3//! Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token
4//! identical to plain greedy `generate`. This module provides:
5//!   - `mtp_head_forward`  (§A, T=1): one NextN draft-token forward.
6//!   - `decode_step_t`     (§D.3, T=K+1): batched target verify forward, all-column logits.
7//!   - `generate_spec`     (§B): the draft/verify/accept/rollback orchestrator.
8//! Cache snapshot/rollback lives in cache.rs (§D.4). The MTP head uses its OWN scratch KV (§D.6),
9//! PERSISTENT over the committed sequence (see `MtpScratch`).
10
11use crate::Engine;
12use crate::cache::{Cache, KvLayer};
13use crate::forward::argmax;
14use crate::hybrid::{FullAttnLayer, HybridModel, LinearAttnLayer, Mixer, MtpHead};
15use cudarc::driver::CudaSlice;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18/// Parse the documented `MEMRA_SPEC_REPLAY=1` rollback seam.
19///
20/// Keep this shared with serving admission so `=0` cannot select replay in one
21/// layer while another layer treats it as disabled.
22pub fn spec_replay_env_on(value: Option<&str>) -> bool {
23    value == Some("1")
24}
25
26pub fn spec_replay_env_enabled() -> bool {
27    let value = std::env::var("MEMRA_SPEC_REPLAY").ok();
28    spec_replay_env_on(value.as_deref())
29}
30
31/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
32/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
33/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
34/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
35/// target arrays are `[gamma, top_k]` in row-major order.
36pub struct DsparkAnchorRecord {
37    pub position: usize,
38    pub hidden: Vec<f32>,
39    pub tokens: Vec<u32>,
40    pub target_top_ids: Vec<u32>,
41    pub target_top_logits: Vec<f32>,
42    pub target_top_probs: Vec<f32>,
43    pub target_tail_probs: Vec<f32>,
44}
45
46fn dspark_sparse_softmax_topk(
47    logits: &[f32],
48    top_k: usize,
49    temperature: f32,
50) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
51    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
52        return Err("invalid DSpark sparse-softmax shape or temperature".into());
53    }
54    if logits.iter().any(|value| !value.is_finite()) {
55        return Err("DSpark target logits contain a non-finite value".into());
56    }
57    let mut ranked: Vec<(u32, f32)> = logits
58        .iter()
59        .copied()
60        .enumerate()
61        .map(|(index, value)| (index as u32, value))
62        .collect();
63    let compare = |left: &(u32, f32), right: &(u32, f32)| {
64        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
65    };
66    ranked.select_nth_unstable_by(top_k - 1, compare);
67    ranked[..top_k].sort_unstable_by(compare);
68
69    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
70    let inv_temperature = 1.0f64 / temperature as f64;
71    let denominator: f64 = logits
72        .iter()
73        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
74        .sum();
75    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
76    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
77    let top_probs: Vec<f32> = top_logits
78        .iter()
79        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
80        .collect();
81    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
82    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
83    Ok((ids, top_logits, top_probs, tail))
84}
85
86fn flatten_dspark_rows<T>(
87    rows: Vec<Option<Vec<T>>>,
88    position: usize,
89    label: &str,
90) -> Result<Vec<T>, Box<dyn std::error::Error>> {
91    let mut flattened = Vec::new();
92    for (slot, row) in rows.into_iter().enumerate() {
93        flattened.extend(
94            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
95        );
96    }
97    Ok(flattened)
98}
99
100/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
101/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
102/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
103/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
104/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
105/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
106/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
107pub(crate) fn spec_hpost() -> bool {
108    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
109    *H.get_or_init(|| {
110        std::env::var("MEMRA_SPEC_HPOST")
111            .map(|v| v != "0")
112            .unwrap_or(false)
113    })
114}
115
116/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
117/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
118/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
119/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
120/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
121/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
122/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
123/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
124/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
125pub(crate) fn spec_lean() -> bool {
126    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
127    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
128    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
129    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
130    *L.get_or_init(|| {
131        std::env::var("MEMRA_SPEC_LEAN")
132            .map(|v| v != "0")
133            .unwrap_or(true)
134    })
135}
136
137/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
138/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
139/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
140/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
141/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
142/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
143///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
144///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
145///     t-loop == chained T=1 steps);
146/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
147///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
148/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
149pub(crate) fn spec_m2() -> bool {
150    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
151    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
152    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
153    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
154    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
155    *M.get_or_init(|| {
156        std::env::var("MEMRA_SPEC_M2")
157            .map(|v| v != "0")
158            .unwrap_or(true)
159    })
160}
161pub(crate) fn spec_stream() -> bool {
162    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
164}
165pub(crate) fn spec_stream_m() -> usize {
166    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
167    *M.get_or_init(|| {
168        std::env::var("MEMRA_SPEC_STREAM_M")
169            .ok()
170            .and_then(|v| v.parse().ok())
171            .unwrap_or(4)
172    })
173}
174pub(crate) fn spec_devacc() -> bool {
175    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
176    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
177}
178
179/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
180///
181/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
182/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
183/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
184/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
185/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
186/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
187/// the flag crashed precisely the regime it exists to investigate.
188///
189/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
190/// indexing (an out-of-range pred there is a real bug and must still be loud).
191fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
192    if base == 0 {
193        return last_pred.to_string();
194    }
195    match preds.get(base - 1) {
196        Some(p) => p.to_string(),
197        // sampled: the greedy per-column argmax was never run for this round.
198        None => {
199            debug_assert!(
200                sampled,
201                "greedy spec: preds[{}] missing at base {base}",
202                base - 1
203            );
204            "n/a".to_string()
205        }
206    }
207}
208
209/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
210///
211/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
212/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
213/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
214/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
215/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
216/// not believe in — and `u * 0 < p` then accepts it unconditionally.
217///
218/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
219/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
220pub(crate) fn skey_probe() -> bool {
221    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
222    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
223}
224
225/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
226/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
227/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
228/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
229/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
230/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
231/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
232/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
233/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
234pub trait SpecConstraint {
235    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
236    /// masked argmax).
237    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
238    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
239    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
240    /// Is `tok` consumable in the CURRENT state?
241    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
242    /// Advance the state with an emitted token.
243    fn consume(&mut self, tok: u32) -> Result<(), String>;
244
245    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
246    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
247    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
248    // loose, research/constrained-full-20260803). These three methods let the engine mask the
249    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
250    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
251    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
252    // stays the correctness backstop and the emitted stream is unchanged by construction
253    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
254    // argmax; a cut slot is recomputed as the masked argmax either way).
255    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
256
257    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
258    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
259    fn draft_mask_enabled(&self) -> bool {
260        false
261    }
262    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
263    /// slot. Called once per spec round, before the first draft position.
264    fn draft_begin(&mut self) -> Result<(), String> {
265        Ok(())
266    }
267    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
268    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
269    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
270        Ok(None)
271    }
272    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
273    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
274    /// engine stops drafting; the token already pushed still goes through verify.
275    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
276        Ok(false)
277    }
278}
279
280/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
281/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
282/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
283/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
284/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
285/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
286/// verify emits the masked argmax as usual).
287fn upload_draft_mask(
288    e: &Engine,
289    c: &mut dyn SpecConstraint,
290    dst: &mut CudaSlice<u32>,
291    d2t: Option<&Vec<u32>>,
292    d_vocab: usize,
293    words: usize,
294) -> Result<bool, Box<dyn std::error::Error>> {
295    let Some(tw) = c
296        .draft_mask_words()
297        .map_err(|e2| format!("constraint: {e2}"))?
298    else {
299        return Ok(false);
300    };
301    let bit = |t: usize| -> bool {
302        let w = t >> 5;
303        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
304    };
305    let mut buf = vec![0u32; words];
306    match d2t {
307        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
308        Some(map) => {
309            for (i, &t) in map.iter().enumerate().take(d_vocab) {
310                if bit(t as usize) {
311                    buf[i >> 5] |= 1u32 << (i & 31);
312                }
313            }
314        }
315        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
316        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
317        None => {
318            let n = tw.len().min(words);
319            buf[..n].copy_from_slice(&tw[..n]);
320        }
321    }
322    if buf.iter().all(|w| *w == 0) {
323        return Ok(false);
324    }
325    e.htod_u32_into(dst, &buf)?;
326    Ok(true)
327}
328
329/// Keep the full token-embedding table in host memory and upload only the rows needed by each
330/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
331/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
332/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
333pub(crate) fn spec_host_embd() -> bool {
334    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
335    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
336}
337
338/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
339/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
340/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
341/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
342/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
343/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
344/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
345/// run-spec K=1..8 + acceptance identity arbitrate e2e).
346pub(crate) fn spec_fused_t() -> bool {
347    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
348    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
349    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
350    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
351    *F.get_or_init(|| {
352        std::env::var("MEMRA_SPEC_FUSED_T")
353            .map(|v| v != "0")
354            .unwrap_or(true)
355    })
356}
357
358/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
359/// Only call this on such buffers — the lean contract is "identical bytes by construction".
360fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
361    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
362}
363
364/// Scratch KV for the MTP block (one full-attn layer).
365///
366/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
367/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
368/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
369/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
370/// engine's "mtp_update" design). Entries come from two sources:
371///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
372///     hidden chain-approximate — the reference engine accepts the same);
373///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
374///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
375/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
376/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
377/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
378/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
379/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
380/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
381/// committed row across turns (the predecessor-pairing seed + fill anchor).
382/// Per-request sampling config for the sampled-spec serve path.
383#[derive(Clone, Copy, Debug)]
384pub struct SpecSampling {
385    pub temp: f32,
386    pub seed: u64,
387    pub top_k: i32,            // 0 = off
388    pub top_p: f32,            // 1.0 = off
389    pub min_p: f32,            // 0.0 = off
390    pub penalty_last_n: usize, // 0 = penalties off
391    pub penalty_repeat: f32,
392    pub penalty_freq: f32,
393    pub penalty_present: f32,
394}
395
396/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
397/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
398pub const SPEC_TELEM_POS: usize = 8;
399
400/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
401/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
402/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
403/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
404/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
405/// in NEITHER drafted nor accepted.
406#[derive(Clone, Copy, Default, Debug)]
407pub struct SpecTelemetry {
408    /// verify rounds completed (a round-stream burst counts each of its M rounds).
409    pub rounds: u64,
410    /// tokens drafted / accepted across all rounds.
411    pub drafted: u64,
412    pub accepted: u64,
413    /// how often draft position j (0-based within a round's chain) was offered / accepted.
414    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
415    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
416    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
417    pub pos_drafted: [u64; SPEC_TELEM_POS],
418    pub pos_accepted: [u64; SPEC_TELEM_POS],
419}
420
421impl SpecTelemetry {
422    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
423    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
424    /// a wrapped counter.
425    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
426        let mut d = SpecTelemetry {
427            rounds: self.rounds.saturating_sub(prev.rounds),
428            drafted: self.drafted.saturating_sub(prev.drafted),
429            accepted: self.accepted.saturating_sub(prev.accepted),
430            ..Default::default()
431        };
432        for j in 0..SPEC_TELEM_POS {
433            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
434            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
435        }
436        d
437    }
438    /// Fieldwise `self += d` — the worker's per-model aggregation.
439    pub fn merge(&mut self, d: &SpecTelemetry) {
440        self.rounds += d.rounds;
441        self.drafted += d.drafted;
442        self.accepted += d.accepted;
443        for j in 0..SPEC_TELEM_POS {
444            self.pos_drafted[j] += d.pos_drafted[j];
445            self.pos_accepted[j] += d.pos_accepted[j];
446        }
447    }
448
449    /// Mean accepted draft-prefix length per verify round (tau).
450    pub fn tau(&self) -> f64 {
451        if self.rounds > 0 {
452            self.accepted as f64 / self.rounds as f64
453        } else {
454            0.0
455        }
456    }
457}
458
459/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
460/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
461/// launch, synchronization, allocation, or ordering dependency to the numeric path.
462struct SpecTelemetryCounters {
463    rounds: AtomicU64,
464    drafted: AtomicU64,
465    accepted: AtomicU64,
466    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
467    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
468}
469
470impl Default for SpecTelemetryCounters {
471    fn default() -> Self {
472        Self {
473            rounds: AtomicU64::new(0),
474            drafted: AtomicU64::new(0),
475            accepted: AtomicU64::new(0),
476            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
477            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
478        }
479    }
480}
481
482impl SpecTelemetryCounters {
483    fn record_round(&self, drafted: usize, accepted: usize) {
484        debug_assert!(accepted <= drafted);
485        self.rounds.fetch_add(1, Ordering::Relaxed);
486        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
487        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
488        for counter in self.pos_drafted.iter().take(drafted) {
489            counter.fetch_add(1, Ordering::Relaxed);
490        }
491        for counter in self.pos_accepted.iter().take(accepted) {
492            counter.fetch_add(1, Ordering::Relaxed);
493        }
494    }
495
496    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
497    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
498    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
499        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
500        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
501        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
502    }
503
504    fn snapshot(&self) -> SpecTelemetry {
505        SpecTelemetry {
506            rounds: self.rounds.load(Ordering::Relaxed),
507            drafted: self.drafted.load(Ordering::Relaxed),
508            accepted: self.accepted.load(Ordering::Relaxed),
509            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
510            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
511        }
512    }
513}
514
515pub struct SpecSession {
516    pub(crate) cache: Cache,
517    pub(crate) scratch: MtpScratch,
518    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
519    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
520    /// session must count them. Callers render output from this, not from their own echo.
521    pub committed: Vec<u32>,
522    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
523    pub(crate) last_h: Option<CudaSlice<f32>>,
524    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
525    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
526    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
527    pub next_pred: Option<u32>,
528    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
529    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
530    pub sctr: u32,
531    pub uctr: u32,
532    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
533    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
534    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
535    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
536    /// research/spec-serving-20260801). None before the first turn; error paths drop it
537    /// (next burst recaptures — serve retires errored sessions anyway).
538    pub(crate) draft_ctx: Option<DraftGraphCtx>,
539    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
540    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
541    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
542    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
543    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
544    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
545    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
546    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
547    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
548    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
549    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
550    pub pending_tok: Option<u32>,
551    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
552    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
553    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
554    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
555    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
556    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
557    /// accounting the loop already does — no syncs, no allocation. NOTE a
558    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
559    /// diff with [`SpecTelemetry::delta_since`] around each burst.
560    telem: SpecTelemetryCounters,
561    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
562    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
563    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
564    /// prime, result lands in `boundary_capture`.
565    pub capture_at: Option<usize>,
566    /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
567    /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
568    /// publication just isn't available for that request.
569    pub boundary_capture: Option<SpecBoundaryCapture>,
570}
571impl SpecSession {
572    /// Context capacity of the session's caches (the server's ContextFull guard).
573    pub fn cache_max_ctx(&self) -> usize {
574        self.cache.max_ctx
575    }
576    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
577    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
578    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
579    /// the prime boundary), so no copy was taken at prime time.
580    pub fn cache_ref(&self) -> &Cache {
581        &self.cache
582    }
583    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
584    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
585    /// like the trunk KV — draft rows below the prompt end are append-only for the
586    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
587    /// committed length, never below the prime boundary, and the true-hidden refresh
588    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
589    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
590    /// prefix-addressable; the prefix cache already refuses that class end to end).
591    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
592        if self.scratch.kv.ring.is_some() {
593            return None;
594        }
595        Some((
596            &self.scratch.kv.k,
597            &self.scratch.kv.v,
598            self.scratch.kv.k_tok_bytes,
599            self.scratch.kv.v_tok_bytes,
600        ))
601    }
602    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
603    pub fn telemetry(&self) -> SpecTelemetry {
604        self.telem.snapshot()
605    }
606    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
607    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
608    /// `spec_rewind_to_checkpoint`.
609    pub fn rewind_pos(&self) -> Option<usize> {
610        self.turn_ckpt.as_ref().map(|c| c.pos)
611    }
612    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
613    pub fn rewind_is_resident(&self) -> bool {
614        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
615            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
616        })
617    }
618    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
619    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
620    /// session has never run a turn and has no prediction to hand over.
621    pub fn demote_ready(&self) -> bool {
622        self.pending_tok.is_none() && self.next_pred.is_some()
623    }
624    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
625    pub fn has_pending(&self) -> bool {
626        self.pending_tok.is_some()
627    }
628    /// Committed row count == cache rows (the session invariant), for the caller's own
629    /// `fed`-length cross-check at a handoff boundary.
630    pub fn committed_len(&self) -> usize {
631        self.committed.len()
632    }
633    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
634    /// cache + next-token prediction to the plain batched-decode path.
635    ///
636    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
637    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
638    /// tokenwise prime of the same `committed` sequence would have left it (that is the
639    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
640    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
641    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
642    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
643    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
644    /// a state indistinguishable from one the batched path produced itself: the batched tick
645    /// emits `next_pred`, feeds it into this same cache, and decodes on.
646    ///
647    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
648    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
649    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
650    /// path would silently skip a token.
651    ///
652    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
653    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
654    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
655    /// would mean an `mtp_kv_fill` over the whole committed history).
656    pub fn into_demoted(self) -> Option<(Cache, u32)> {
657        if self.pending_tok.is_some() {
658            return None;
659        }
660        let np = self.next_pred?;
661        debug_assert_eq!(
662            self.cache.pos,
663            self.committed.len(),
664            "demotion handoff: cache rows != committed tokens"
665        );
666        Some((self.cache, np))
667    }
668    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
669    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
670    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
671    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
672    pub fn reset_graph_fallback_on_resume(&mut self) {
673        if let Some(line) = self
674            .draft_ctx
675            .as_mut()
676            .and_then(|c| c.failed.reset_on_resume())
677        {
678            eprintln!("{line}");
679        }
680    }
681}
682
683/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
684///
685/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
686/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
687/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
688/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
689/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
690/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
691///
692/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
693/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
694/// position index, so it must be a real device COPY — that copy is the entire reason a spec
695/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
696/// below the boundary were written by this turn's fill and are never revisited (the per-round
697/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
698/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
699/// predecessor-pairing anchor the next prime's fill reads for its first row.
700///
701/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
702pub(crate) struct SpecCheckpoint {
703    snap: crate::cache::CacheSnapshot,
704    /// Committed length at the boundary (== cache.pos there, the session invariant).
705    pos: usize,
706    /// Pre-output_norm hidden of row `pos - 1`.
707    last_h: CudaSlice<f32>,
708}
709
710/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
711/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
712/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
713/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
714/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
715/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
716/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
717/// so the worker slices those from the live caches post-burst instead of copying at prime time.
718pub struct SpecBoundaryCapture {
719    pub snap: crate::cache::CacheSnapshot,
720    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
721    pub pos: usize,
722    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
723    pub logits: Vec<f32>,
724    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
725    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
726    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
727    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
728    pub last_h: Vec<f32>,
729}
730
731/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
732/// spec boundary capture carries for later restored-session fills. Failure is silent
733/// (`turn_ckpt` convention): the capture publishes without an anchor.
734fn capture_boundary_hidden(
735    e: &Engine,
736    h_rows: &CudaSlice<f32>,
737    pos: usize,
738    n_embd: usize,
739) -> Vec<f32> {
740    if pos == 0 || h_rows.len() < pos * n_embd {
741        return Vec::new();
742    }
743    let Ok(mut row) = e.uninit(n_embd) else {
744        return Vec::new();
745    };
746    if e.copy_view_into(
747        &mut row,
748        0,
749        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
750        n_embd,
751    )
752    .is_err()
753    {
754        return Vec::new();
755    }
756    e.dtoh(&row).unwrap_or_default()
757}
758
759/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
760/// Default ON: the token a burst emits at its own boundary is drawn from the request's
761/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
762/// every boundary) without touching greedy, which is byte-unaffected either way.
763pub fn spec_sampled_boundary_on() -> bool {
764    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
765    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
766}
767
768/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
769/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
770/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
771/// restores the pre-lane posture (each burst restarts the window from its own prompt
772/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
773/// must keep refusing penalized sampled prefix-cache restores, because the restored
774/// session's continuation burst is handed no prompt slice at all.
775pub fn spec_pen_session_on() -> bool {
776    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
777    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
778}
779
780/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
781/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
782/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
783/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
784/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
785/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
786pub fn spec_restore_republish_on() -> bool {
787    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
788    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
789}
790
791/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
792/// the argmax the pre-lane code would have emitted from the same row. This is how the
793/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
794fn spec_boundary_trace() -> bool {
795    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
796    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
797}
798
799/// llama-parity floor for the penalty window when the request does not ask for a bigger
800/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
801/// non-identity penalty, so this floor only matters to explicit small windows and to the
802/// CLI env path.
803const PEN_WINDOW_FLOOR: usize = 64;
804
805/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
806/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
807/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
808/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
809/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
810/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
811/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
812/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
813/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
814/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
815/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
816/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
817const PEN_WINDOW_MAX: usize = 8192;
818
819/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
820/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
821/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
822/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
823/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
824/// client actually asked us to penalize, where the pre-lane code had NOTHING.
825fn pen_window_seed(
826    session_committed: &[u32],
827    burst_prompt: &[u32],
828    penalty_last_n: usize,
829) -> Vec<u32> {
830    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
831    let take_prompt = burst_prompt.len().min(win);
832    let take_sess = (win - take_prompt).min(session_committed.len());
833    let mut hist = Vec::with_capacity(take_sess + take_prompt);
834    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
835    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
836    hist
837}
838
839/// Draw a BOUNDARY token from the target distribution the request asked for
840/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
841/// every burst boundary".
842///
843/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
844/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
845/// row after the last committed token on a continuation burst; the prefix-cache entry's
846/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
847/// regimes, so a sampled stream took a greedy token once per burst — measured, not
848/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
849/// customer asked for a sampled token, so this draws one.
850///
851/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
852/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
853/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
854/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
855/// composition means `sample_check`'s distributional oracle covers this draw too, and the
856/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
857///
858/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
859/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
860/// stream the accept walk uses — never a second, independently seeded stream (which would be
861/// a new distributional bug: two streams from one seed correlate wherever their counters
862/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
863/// to the cold session's own first draw from the same logits row, which is what preserves the
864/// sampled-hit lane's per-seed hit==cold byte identity.
865#[allow(clippy::too_many_arguments)]
866pub fn sample_boundary_token_dev(
867    e: &Engine,
868    logits: &CudaSlice<f32>,
869    n_vocab: usize,
870    sp: &SpecSampling,
871    pen_hist: &[u32],
872    sctr: &mut u32,
873    site: &str,
874) -> Result<u32, Box<dyn std::error::Error>> {
875    debug_assert!(
876        sp.temp > 0.0,
877        "boundary sampling is the sampled regime only"
878    );
879    // Own copy: penalize_logits mutates in place and the caller's row is live state
880    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
881    let mut col = e.zeros(n_vocab)?;
882    e.copy_into(&mut col, 0, logits, n_vocab)?;
883    let pen_on = sp.penalty_last_n > 0
884        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
885    if pen_on && !pen_hist.is_empty() {
886        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
887        let w0 = pen_hist
888            .len()
889            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
890        let hist = &pen_hist[w0..];
891        let hd = e.htod_u32_v(hist)?;
892        e.penalize_logits(
893            &mut col,
894            &hd,
895            hist.len(),
896            sp.penalty_repeat,
897            sp.penalty_freq,
898            sp.penalty_present,
899            n_vocab,
900        )?;
901    }
902    let rows0 = e.htod_i32(&[0])?;
903    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
904    e.filter_stats(
905        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
906        sp.top_p, sp.min_p,
907    )?;
908    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
909    let mut perturb = e.zeros(n_vocab)?;
910    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
911    *sctr = sctr.wrapping_add(1);
912    let td = e.argmax_token_device(&perturb, n_vocab)?;
913    let tok = e.dtoh_u32_one(&td)?;
914    if spec_boundary_trace() {
915        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
916        let raw = e.argmax_token_device(logits, n_vocab)?;
917        let greedy = e.dtoh_u32_one(&raw)?;
918        eprintln!(
919            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
920             deviates={} temp={} sctr={}",
921            (tok != greedy) as u8,
922            sp.temp,
923            sctr.wrapping_sub(1),
924        );
925    }
926    Ok(tok)
927}
928
929/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
930/// host `Vec<f32>`).
931#[allow(clippy::too_many_arguments)]
932pub fn sample_boundary_token(
933    e: &Engine,
934    logits: &[f32],
935    sp: &SpecSampling,
936    pen_hist: &[u32],
937    sctr: &mut u32,
938    site: &str,
939) -> Result<u32, Box<dyn std::error::Error>> {
940    let n_vocab = logits.len();
941    let d = e.htod(logits)?;
942    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
943}
944
945struct SpecPipeTraceClock {
946    pair: usize,
947    started: std::time::Instant,
948}
949
950#[derive(Clone)]
951struct SpecPipeTraceCtx {
952    clock: std::sync::Arc<SpecPipeTraceClock>,
953    round: usize,
954    lane: usize,
955}
956
957struct SpecPipeTraceMarker {
958    trace: SpecPipeTraceCtx,
959    phase: &'static str,
960    edge: &'static str,
961    slot: Option<usize>,
962}
963
964unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
965    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
966    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
967    let slot = marker
968        .slot
969        .map(|v| v.to_string())
970        .unwrap_or_else(|| "-".into());
971    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
972    use std::io::Write as _;
973    let stderr = std::io::stderr();
974    let mut stderr = stderr.lock();
975    let _ = writeln!(
976        stderr,
977        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
978         slot={slot} t_ms={t_ms:.3}",
979        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
980    );
981}
982
983fn enqueue_spec_pipe_trace_marker(
984    stream: &cudarc::driver::CudaStream,
985    trace: Option<&SpecPipeTraceCtx>,
986    phase: &'static str,
987    edge: &'static str,
988    slot: Option<usize>,
989) -> Result<(), Box<dyn std::error::Error>> {
990    let Some(trace) = trace else {
991        return Ok(());
992    };
993    let marker = Box::new(SpecPipeTraceMarker {
994        trace: trace.clone(),
995        phase,
996        edge,
997        slot,
998    });
999    let raw = Box::into_raw(marker);
1000    let result = unsafe {
1001        cudarc::driver::result::stream::launch_host_function(
1002            stream.cu_stream(),
1003            spec_pipe_trace_marker,
1004            raw.cast(),
1005        )
1006    };
1007    if let Err(err) = result {
1008        unsafe {
1009            drop(Box::from_raw(raw));
1010        }
1011        return Err(err.into());
1012    }
1013    Ok(())
1014}
1015
1016#[derive(Default)]
1017struct SpecPipeProgress {
1018    setup_done: [bool; 2],
1019    draft_done: [usize; 2],
1020    stage0_done: [usize; 2],
1021    verify_done: [usize; 2],
1022    accept_done: [usize; 2],
1023    finished: [bool; 2],
1024    aborted: bool,
1025}
1026
1027/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1028/// keeps its existing call stack and round locals; this object only orders phase entry. The
1029/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1030/// cannot be interleaved by the two host threads.
1031struct SpecPipeSync {
1032    progress: std::sync::Mutex<SpecPipeProgress>,
1033    changed: std::sync::Condvar,
1034    primary: std::sync::Mutex<()>,
1035    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1036}
1037
1038impl SpecPipeSync {
1039    fn new() -> Self {
1040        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1041        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1042            std::sync::Arc::new(SpecPipeTraceClock {
1043                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1044                started: std::time::Instant::now(),
1045            })
1046        });
1047        Self {
1048            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1049            changed: std::sync::Condvar::new(),
1050            primary: std::sync::Mutex::new(()),
1051            trace,
1052        }
1053    }
1054}
1055
1056#[derive(Clone)]
1057struct SpecPipeLane {
1058    sync: std::sync::Arc<SpecPipeSync>,
1059    lane: usize,
1060}
1061
1062impl SpecPipeLane {
1063    fn peer(&self) -> usize {
1064        1 - self.lane
1065    }
1066
1067    fn aborted() -> Box<dyn std::error::Error> {
1068        "paired speculative peer aborted".into()
1069    }
1070
1071    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1072        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1073            clock: clock.clone(),
1074            round,
1075            lane: self.lane,
1076        })
1077    }
1078
1079    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1080        let mut p = self.sync.progress.lock().unwrap();
1081        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1082            p = self.sync.changed.wait(p).unwrap();
1083        }
1084        if p.aborted {
1085            Err(Self::aborted())
1086        } else {
1087            Ok(())
1088        }
1089    }
1090
1091    fn setup_end(&self) {
1092        let mut p = self.sync.progress.lock().unwrap();
1093        p.setup_done[self.lane] = true;
1094        self.sync.changed.notify_all();
1095    }
1096
1097    fn draft_begin(
1098        &self,
1099        round: usize,
1100    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1101        let peer = self.peer();
1102        let mut p = self.sync.progress.lock().unwrap();
1103        loop {
1104            if p.aborted {
1105                return Err(Self::aborted());
1106            }
1107            let setup_ready =
1108                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1109            let prior_ready = p.accept_done[self.lane] >= round
1110                && (p.accept_done[peer] >= round || p.finished[peer]);
1111            let turn_ready = if self.lane == 0 {
1112                true
1113            } else {
1114                p.draft_done[0] > round || p.finished[0]
1115            };
1116            if setup_ready && prior_ready && turn_ready {
1117                break;
1118            }
1119            p = self.sync.changed.wait(p).unwrap();
1120        }
1121        drop(p);
1122        Ok(self.sync.primary.lock().unwrap())
1123    }
1124
1125    fn draft_end(&self, round: usize) {
1126        let mut p = self.sync.progress.lock().unwrap();
1127        p.draft_done[self.lane] = round + 1;
1128        self.sync.changed.notify_all();
1129    }
1130
1131    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1132    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1133    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1134        let peer = self.peer();
1135        let mut p = self.sync.progress.lock().unwrap();
1136        loop {
1137            if p.aborted {
1138                return Err(Self::aborted());
1139            }
1140            let ready = if self.lane == 0 {
1141                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1142            } else {
1143                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1144            };
1145            if ready {
1146                return Ok(self.lane == 0 || p.finished[peer]);
1147            }
1148            p = self.sync.changed.wait(p).unwrap();
1149        }
1150    }
1151
1152    fn stage0_end(&self, round: usize) {
1153        let mut p = self.sync.progress.lock().unwrap();
1154        p.stage0_done[self.lane] = round + 1;
1155        self.sync.changed.notify_all();
1156    }
1157
1158    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1159    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1160    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1161        let mut p = self.sync.progress.lock().unwrap();
1162        while !p.aborted
1163            && !(p.stage0_done[self.lane] > round
1164                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1165        {
1166            p = self.sync.changed.wait(p).unwrap();
1167        }
1168        if p.aborted {
1169            Err(Self::aborted())
1170        } else {
1171            Ok(())
1172        }
1173    }
1174
1175    fn verify_end(&self, round: usize) {
1176        let mut p = self.sync.progress.lock().unwrap();
1177        p.verify_done[self.lane] = round + 1;
1178        self.sync.changed.notify_all();
1179    }
1180
1181    fn accept_begin(
1182        &self,
1183        round: usize,
1184    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1185        let mut p = self.sync.progress.lock().unwrap();
1186        loop {
1187            if p.aborted {
1188                return Err(Self::aborted());
1189            }
1190            let ready = if self.lane == 0 {
1191                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1192            } else {
1193                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1194            };
1195            if ready {
1196                break;
1197            }
1198            p = self.sync.changed.wait(p).unwrap();
1199        }
1200        drop(p);
1201        Ok(self.sync.primary.lock().unwrap())
1202    }
1203
1204    fn accept_end(&self, round: usize) {
1205        let mut p = self.sync.progress.lock().unwrap();
1206        p.accept_done[self.lane] = round + 1;
1207        self.sync.changed.notify_all();
1208    }
1209
1210    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1211        self.sync.primary.lock().unwrap()
1212    }
1213
1214    fn finish(&self, failed: bool) {
1215        let mut p = self.sync.progress.lock().unwrap();
1216        p.finished[self.lane] = true;
1217        p.aborted |= failed;
1218        self.sync.changed.notify_all();
1219    }
1220}
1221
1222struct SpecPipeFinish<'a> {
1223    lane: &'a SpecPipeLane,
1224    closed: bool,
1225}
1226
1227impl<'a> SpecPipeFinish<'a> {
1228    fn new(lane: &'a SpecPipeLane) -> Self {
1229        Self {
1230            lane,
1231            closed: false,
1232        }
1233    }
1234
1235    fn close(&mut self, failed: bool) {
1236        self.lane.finish(failed);
1237        self.closed = true;
1238    }
1239}
1240
1241impl Drop for SpecPipeFinish<'_> {
1242    fn drop(&mut self) {
1243        if !self.closed {
1244            self.lane.finish(true);
1245        }
1246    }
1247}
1248
1249/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1250/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1251/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1252/// binds that context before touching the session, joins before returning, and never aliases the
1253/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1254/// session type Send.
1255struct SpecPipeSessionPtr(*mut SpecSession);
1256
1257unsafe impl Send for SpecPipeSessionPtr {}
1258
1259impl SpecPipeSessionPtr {
1260    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1261        unsafe { &mut *self.0 }
1262    }
1263}
1264
1265/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1266/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1267/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1268/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1269/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1270/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1271/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1272/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1273/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1274///
1275/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1276/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1277/// load-bearing:
1278///
1279/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1280///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1281///   This is all the key used to carry.
1282/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1283///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1284///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1285///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1286///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1287///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1288///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1289///
1290/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1291/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1292/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1293/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1294/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1295#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1296pub(crate) struct SampledGraphKey {
1297    seed: u64,
1298    temp_bits: u32,
1299    k: usize,
1300    top_k: i32,
1301    top_p_bits: u32,
1302    min_p_bits: u32,
1303    pen_on: bool,
1304}
1305
1306impl SampledGraphKey {
1307    pub(crate) fn new(
1308        seed: u64,
1309        temp: f32,
1310        k: usize,
1311        top_k: i32,
1312        top_p: f32,
1313        min_p: f32,
1314        pen_on: bool,
1315    ) -> Self {
1316        SampledGraphKey {
1317            seed,
1318            temp_bits: temp.to_bits(),
1319            k,
1320            top_k,
1321            top_p_bits: top_p.to_bits(),
1322            min_p_bits: min_p.to_bits(),
1323            pen_on,
1324        }
1325    }
1326
1327    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1328    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1329    /// the key can never drift apart (they were three separate expressions before this lane, and
1330    /// the launch site simply forgot to ask).
1331    pub(crate) fn pure_temp(&self) -> bool {
1332        self.top_k == 0
1333            && f32::from_bits(self.top_p_bits) >= 1.0
1334            && f32::from_bits(self.min_p_bits) <= 0.0
1335            && !self.pen_on
1336    }
1337}
1338
1339pub(crate) struct DraftGraphCtx {
1340    g_tok: CudaSlice<u32>,
1341    g_pos: CudaSlice<i32>,
1342    g_seed: CudaSlice<f32>,
1343    g_p: CudaSlice<f32>,
1344    g_ctr: CudaSlice<u32>,
1345    g_q: CudaSlice<f32>,
1346    g_perturb: CudaSlice<f32>,
1347    q_slots: Vec<CudaSlice<f32>>,
1348    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1349    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1350    /// per-position contents the host re-uploads before each replay (the graph-promote
1351    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1352    g_dmask: CudaSlice<u32>,
1353    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1354    graph_masked: bool,
1355    graph: Option<cudarc::driver::CudaGraph>,
1356    graph_s: Option<cudarc::driver::CudaGraph>,
1357    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1358    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1359    failed: DraftGraphFallback,
1360    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1361    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1362    s_key: Option<SampledGraphKey>,
1363    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1364    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1365    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1366    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1367    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1368    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1369    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1370    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1371    keeper: Vec<Box<dyn std::any::Any + Send>>,
1372    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1373}
1374
1375/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1376/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1377///
1378/// Three contracts:
1379/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1380///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1381///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1382///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1383///   fallback from paying a doomed capture attempt every burst).
1384/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1385///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1386///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1387///   actually set (quiet on the common clean-resume path).
1388/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1389///   capture attempt whose own failure would re-flip loudly.
1390#[derive(Default)]
1391pub(crate) struct DraftGraphFallback {
1392    greedy: bool,
1393    sampled: bool,
1394}
1395impl DraftGraphFallback {
1396    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1397        if self.greedy {
1398            return None;
1399        }
1400        self.greedy = true;
1401        Some(format!(
1402            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1403        ))
1404    }
1405    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1406        if self.sampled {
1407            return None;
1408        }
1409        self.sampled = true;
1410        Some(format!(
1411            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1412        ))
1413    }
1414    fn greedy_failed(&self) -> bool {
1415        self.greedy
1416    }
1417    fn sampled_failed(&self) -> bool {
1418        self.sampled
1419    }
1420    fn clear_greedy(&mut self) {
1421        self.greedy = false;
1422    }
1423    fn clear_sampled(&mut self) {
1424        self.sampled = false;
1425    }
1426    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1427    /// was set (so clean resumes stay quiet).
1428    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1429        if !self.greedy && !self.sampled {
1430            return None;
1431        }
1432        let which = match (self.greedy, self.sampled) {
1433            (true, true) => "greedy+sampled",
1434            (true, false) => "greedy",
1435            _ => "sampled",
1436        };
1437        self.greedy = false;
1438        self.sampled = false;
1439        Some(format!(
1440            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1441        ))
1442    }
1443}
1444
1445impl DraftGraphCtx {
1446    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1447        Ok(DraftGraphCtx {
1448            g_tok: e.alloc_u32_zeroed(1)?,
1449            g_pos: e.htod_i32(&[0])?,
1450            g_seed: e.zeros(n_embd)?,
1451            g_p: e.zeros(1)?,
1452            g_ctr: e.alloc_u32_zeroed(1)?,
1453            g_q: e.zeros(qlen)?,
1454            g_perturb: e.zeros(qlen)?,
1455            q_slots: Vec::new(),
1456            g_dmask: e.alloc_u32_zeroed(1)?,
1457            graph_masked: false,
1458            graph: None,
1459            graph_s: None,
1460            failed: DraftGraphFallback::default(),
1461            s_key: None,
1462            keeper: Vec::new(),
1463            keeper_s: Vec::new(),
1464        })
1465    }
1466}
1467
1468pub(crate) struct MtpScratch {
1469    kv: KvLayer,
1470    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1471    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1472    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1473    /// smaller host-indexed SWA ring instead.
1474    cap: usize,
1475}
1476
1477fn mtp_scratch_layout(
1478    cfg: &memra_gguf::config::ModelConfig,
1479    geom: Option<&crate::hybrid::DraftGeom>,
1480) -> (usize, usize, usize, usize) {
1481    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1482    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1483    let head_dim_k = cfg.head_dim_k as usize;
1484    let head_dim_v = cfg.head_dim_v as usize;
1485    assert!(
1486        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1487        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1488    );
1489    let kv_dim_k = head_dim_k * n_head_kv;
1490    let kv_dim_v = head_dim_v * n_head_kv;
1491    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1492    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1493    let (kbb, vbb) = crate::kv_blk_bytes();
1494    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1495    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1496    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1497}
1498
1499impl MtpScratch {
1500    fn new(
1501        e: &Engine,
1502        cfg: &memra_gguf::config::ModelConfig,
1503        cap: usize,
1504        geom: Option<&crate::hybrid::DraftGeom>,
1505    ) -> Result<Self, Box<dyn std::error::Error>> {
1506        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1507        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1508        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1509        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1510        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1511        let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1512            let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1513            Some(crate::cache::KvRing::new(
1514                crate::cache::swa_ring_rows(window, cap),
1515                window,
1516            ))
1517        } else {
1518            None
1519        };
1520        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1521        Ok(MtpScratch {
1522            kv: KvLayer {
1523                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1524                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1525                kv_dim_k,
1526                kv_dim_v,
1527                k_tok_bytes,
1528                v_tok_bytes,
1529                len: 0,
1530                ring,
1531                len_d: e.htod_i32(&[0])?,
1532            },
1533            cap,
1534        })
1535    }
1536    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1537    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1538    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1539    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1540        if self
1541            .kv
1542            .ring
1543            .as_ref()
1544            .is_some_and(|ring| !ring.can_rewind_to(n))
1545        {
1546            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1547        }
1548        self.kv.len = n;
1549        e.set_i32_one(&mut self.kv.len_d, n as i32)
1550    }
1551
1552    fn can_rewind_to(&self, n: usize) -> bool {
1553        self.kv
1554            .ring
1555            .as_ref()
1556            .is_none_or(|ring| ring.can_rewind_to(n))
1557    }
1558}
1559
1560/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1561/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1562/// full weight reads per round — recomputing columns the verify had already produced
1563/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1564/// to "after the first j verify columns" WITHOUT re-running the trunk:
1565/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1566///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1567///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1568///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1569///   pure-copy ring rebuild.
1570/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1571///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1572///   target: j <= t-1).
1573/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1574/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1575struct GdnStash {
1576    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1577    q_l2: CudaSlice<f32>,
1578    k_l2: CudaSlice<f32>,
1579    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1580    g_log: CudaSlice<f32>,
1581    beta: CudaSlice<f32>, // [t, num_v]
1582}
1583struct VerifyCkpt {
1584    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1585    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1586}
1587/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1588pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1589
1590impl VerifyCkpt {
1591    fn new(n_layer: usize) -> Self {
1592        VerifyCkpt {
1593            gdn: (0..n_layer).map(|_| None).collect(),
1594            cols: (0..n_layer).map(|_| None).collect(),
1595        }
1596    }
1597}
1598
1599/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1600/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1601/// a logical round number.
1602struct VerifyBoundaryTicket {
1603    rt: &'static crate::pp::PpNRt,
1604    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1605    slot: usize,
1606    pos0: usize,
1607    t: usize,
1608    payload: usize,
1609    n_st: usize,
1610    pipelined: bool,
1611    pp_anatomy: bool,
1612    pp_started: std::time::Instant,
1613    reverse_ms: f64,
1614    stage0_ms: f64,
1615    tx_ms: f64,
1616    trace: Option<SpecPipeTraceCtx>,
1617}
1618
1619/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1620/// increment-2 controller can also be armed by the server's fresh-process research door.
1621#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1622pub enum OptiForkGateMode {
1623    Disabled,
1624    Hit,
1625    Miss,
1626    Alternate,
1627    Abort,
1628    Controller,
1629}
1630
1631static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1632static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1633    std::sync::atomic::AtomicU32::new(0);
1634static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1635static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1636static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1637static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1638static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1639static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1640static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1641static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1642static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1643static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1644    std::sync::atomic::AtomicU64::new(0);
1645static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
1646    std::sync::atomic::AtomicU64::new(0);
1647static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1648
1649impl OptiForkGateMode {
1650    fn code(self) -> u8 {
1651        match self {
1652            Self::Disabled => 0,
1653            Self::Hit => 1,
1654            Self::Miss => 2,
1655            Self::Alternate => 3,
1656            Self::Abort => 4,
1657            Self::Controller => 5,
1658        }
1659    }
1660
1661    fn configured() -> Self {
1662        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
1663            1 => Self::Hit,
1664            2 => Self::Miss,
1665            3 => Self::Alternate,
1666            4 => Self::Abort,
1667            5 => Self::Controller,
1668            _ => Self::Disabled,
1669        }
1670    }
1671
1672    fn action(self, generation: u64) -> OptiForkAction {
1673        match self {
1674            Self::Hit => OptiForkAction::Hit,
1675            Self::Miss => OptiForkAction::Miss,
1676            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
1677            Self::Alternate => OptiForkAction::Miss,
1678            Self::Abort => OptiForkAction::Abort,
1679            Self::Disabled | Self::Controller => {
1680                unreachable!("non-forced mode cannot choose a forced fork action")
1681            }
1682        }
1683    }
1684
1685    fn is_forced(self) -> bool {
1686        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
1687    }
1688}
1689
1690/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
1691pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
1692    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
1693}
1694
1695/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
1696/// two-token draft-probability product. Serving can call this only through its explicit
1697/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
1698pub fn set_optipipe_controller_threshold(threshold: f32) {
1699    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
1700    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
1701    set_optipipe_gate_mode(OptiForkGateMode::Controller);
1702}
1703
1704#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1705pub struct OptiForkGateStats {
1706    pub attempts: u64,
1707    pub hits: u64,
1708    pub misses: u64,
1709    pub abort_drains: u64,
1710    pub refusals: u64,
1711    pub gate_checks: u64,
1712    pub gate_admits: u64,
1713    pub gate_rejects: u64,
1714    pub reconciles: u64,
1715    pub wasted_draft_tokens: u64,
1716    pub shadow_draft_tokens: u64,
1717    pub breaker_trips: u64,
1718}
1719
1720#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1721pub struct OptiForkStateIdentity {
1722    pub trunk_kv_bytes: usize,
1723    pub recurrent_bytes: usize,
1724    pub scratch_kv_bytes: usize,
1725    pub hidden_bytes: usize,
1726}
1727
1728pub fn reset_optipipe_gate_stats() {
1729    for counter in [
1730        &OPTI_FORK_ATTEMPTS,
1731        &OPTI_FORK_HITS,
1732        &OPTI_FORK_MISSES,
1733        &OPTI_FORK_ABORT_DRAINS,
1734        &OPTI_FORK_REFUSALS,
1735        &OPTI_GATE_CHECKS,
1736        &OPTI_GATE_ADMITS,
1737        &OPTI_GATE_REJECTS,
1738        &OPTI_RECONCILES,
1739        &OPTI_WASTED_DRAFT_TOKENS,
1740        &OPTI_SHADOW_DRAFT_TOKENS,
1741        &OPTI_BREAKER_TRIPS,
1742    ] {
1743        counter.store(0, std::sync::atomic::Ordering::Relaxed);
1744    }
1745}
1746
1747pub fn optipipe_gate_stats() -> OptiForkGateStats {
1748    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
1749    OptiForkGateStats {
1750        attempts: load(&OPTI_FORK_ATTEMPTS),
1751        hits: load(&OPTI_FORK_HITS),
1752        misses: load(&OPTI_FORK_MISSES),
1753        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
1754        refusals: load(&OPTI_FORK_REFUSALS),
1755        gate_checks: load(&OPTI_GATE_CHECKS),
1756        gate_admits: load(&OPTI_GATE_ADMITS),
1757        gate_rejects: load(&OPTI_GATE_REJECTS),
1758        reconciles: load(&OPTI_RECONCILES),
1759        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
1760        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
1761        breaker_trips: load(&OPTI_BREAKER_TRIPS),
1762    }
1763}
1764
1765#[derive(Clone, Copy, Debug)]
1766struct OptiControllerPolicy {
1767    threshold: f32,
1768    consecutive_misses: u8,
1769    breaker_tripped: bool,
1770}
1771
1772impl OptiControllerPolicy {
1773    fn configured() -> Self {
1774        Self {
1775            threshold: f32::from_bits(
1776                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
1777            ),
1778            consecutive_misses: 0,
1779            breaker_tripped: false,
1780        }
1781    }
1782
1783    fn admit(&self, q_proxy: f32) -> bool {
1784        q_proxy.is_finite()
1785            && (0.0..=1.0).contains(&q_proxy)
1786            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
1787    }
1788
1789    /// Returns true exactly when this resolution newly trips the three-miss breaker.
1790    fn resolve(&mut self, hit: bool) -> bool {
1791        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
1792        // every optimistic opportunity, so the safety breaker is measured separately and must
1793        // not silently turn this arm into "three attempts then serial".
1794        if self.threshold == 0.0 {
1795            self.consecutive_misses = 0;
1796            return false;
1797        }
1798        if hit {
1799            self.consecutive_misses = 0;
1800            return false;
1801        }
1802        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
1803        if !self.breaker_tripped && self.consecutive_misses >= 3 {
1804            self.breaker_tripped = true;
1805            return true;
1806        }
1807        false
1808    }
1809}
1810
1811#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1812enum OptiForkAction {
1813    Hit,
1814    Miss,
1815    Abort,
1816}
1817
1818#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1819struct OptiForkGeneration {
1820    id: u64,
1821    slot: usize,
1822}
1823
1824#[derive(Default)]
1825struct OptiForkGenerationTracker {
1826    next: u64,
1827    live: [Option<u64>; 2],
1828}
1829
1830impl OptiForkGenerationTracker {
1831    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
1832        let generation = OptiForkGeneration {
1833            id: self.next,
1834            slot: (self.next & 1) as usize,
1835        };
1836        if let Some(live) = self.live[generation.slot] {
1837            return Err(format!(
1838                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
1839                generation.slot,
1840            )
1841            .into());
1842        }
1843        self.next += 1;
1844        self.live[generation.slot] = Some(generation.id);
1845        Ok(generation)
1846    }
1847
1848    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
1849        match self.live[generation.slot] {
1850            Some(id) if id == generation.id => {
1851                self.live[generation.slot] = None;
1852                Ok(())
1853            }
1854            other => Err(format!(
1855                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
1856                generation.id, generation.slot,
1857            )
1858            .into()),
1859        }
1860    }
1861}
1862
1863struct OptiForkSeedGeneration {
1864    h_seed: CudaSlice<f32>,
1865    fill_prev: CudaSlice<f32>,
1866    scratch_len: usize,
1867}
1868
1869/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
1870/// generic cache helper accepts one device and therefore cannot copy GDN state split across
1871/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
1872/// device ownership.
1873fn opti_snapshot_stage_owned(
1874    e: &Engine,
1875    cache: &Cache,
1876    rt: &'static crate::pp::PpNRt,
1877    fence: &[usize],
1878) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
1879    let n = cache.kv.len();
1880    let mut snapshot = crate::cache::CacheSnapshot {
1881        kv_len: vec![None; n],
1882        conv: (0..n).map(|_| None).collect(),
1883        ssm: (0..n).map(|_| None).collect(),
1884        pos: cache.pos,
1885    };
1886    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
1887    Ok(snapshot)
1888}
1889
1890fn opti_snapshot_stage_owned_into(
1891    e: &Engine,
1892    cache: &Cache,
1893    rt: &'static crate::pp::PpNRt,
1894    fence: &[usize],
1895    snapshot: &mut crate::cache::CacheSnapshot,
1896) -> Result<(), Box<dyn std::error::Error>> {
1897    if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
1898        return Err("optipipe stage-owned snapshot shape mismatch".into());
1899    }
1900    for stage in 0..rt.n_stages() {
1901        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
1902    }
1903    snapshot.pos = cache.pos;
1904    Ok(())
1905}
1906
1907/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
1908/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
1909/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
1910/// either point would capture one side of the fork at the wrong generation.
1911fn opti_snapshot_one_stage_owned_into(
1912    e: &Engine,
1913    cache: &Cache,
1914    rt: &'static crate::pp::PpNRt,
1915    fence: &[usize],
1916    stage: usize,
1917    snapshot: &mut crate::cache::CacheSnapshot,
1918) -> Result<(), Box<dyn std::error::Error>> {
1919    if fence.len() != rt.n_stages() + 1
1920        || snapshot.kv_len.len() != cache.kv.len()
1921        || stage >= rt.n_stages()
1922    {
1923        return Err("optipipe single-stage snapshot shape mismatch".into());
1924    }
1925    let _scope = rt.enter(stage);
1926    let owner = rt.engine(stage, e);
1927    for il in fence[stage]..fence[stage + 1] {
1928        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
1929        match &cache.recur[il] {
1930            Some(recur) => {
1931                match snapshot.conv[il].as_mut() {
1932                    Some(dst) => {
1933                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
1934                    }
1935                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
1936                }
1937                match snapshot.ssm[il].as_mut() {
1938                    Some(dst) => {
1939                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
1940                    }
1941                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
1942                }
1943            }
1944            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
1945                return Err(
1946                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
1947                );
1948            }
1949            None => {}
1950        }
1951    }
1952    snapshot.pos = cache.pos;
1953    Ok(())
1954}
1955
1956/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
1957/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
1958/// resolve, so the reconcile tables and conditional restores are stage-local.
1959struct OptiForkState {
1960    mode: OptiForkGateMode,
1961    controller: Option<OptiControllerPolicy>,
1962    generations: OptiForkGenerationTracker,
1963    active_snapshot_slot: usize,
1964    alternate_snapshot: crate::cache::CacheSnapshot,
1965    seeds: [OptiForkSeedGeneration; 2],
1966    rt: &'static crate::pp::PpNRt,
1967    fence: [usize; 3],
1968    split: usize,
1969    len_ptrs: CudaSlice<u64>,
1970    saved_lens: CudaSlice<i32>,
1971    forced_acc: CudaSlice<u32>,
1972    valid: CudaSlice<u32>,
1973    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1974    logical_payload_bytes: [usize; 2],
1975}
1976
1977struct OptiForkTicket {
1978    generation: OptiForkGeneration,
1979    boundary: Option<VerifyBoundaryTicket>,
1980    drain: std::sync::Arc<cudarc::driver::CudaStream>,
1981    settled: bool,
1982}
1983
1984struct OptiControllerTicket {
1985    generation: OptiForkGeneration,
1986    boundary: Option<VerifyBoundaryTicket>,
1987    ckpt: Option<VerifyCkpt>,
1988    verify_tokens: [u32; 2],
1989    draft_prob: f32,
1990    eager_seed: Option<CudaSlice<f32>>,
1991    q_proxy: f32,
1992    scratch_len: usize,
1993    issued_at: std::time::Instant,
1994    drain: std::sync::Arc<cudarc::driver::CudaStream>,
1995    settled: bool,
1996}
1997
1998struct OptiControllerPrepared {
1999    verify_tokens: [u32; 2],
2000    draft_prob: f32,
2001    eager_seed: Option<CudaSlice<f32>>,
2002    q_proxy: f32,
2003    scratch_len: usize,
2004}
2005
2006impl OptiControllerTicket {
2007    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2008        self.boundary
2009            .take()
2010            .expect("controller boundary ticket already consumed")
2011    }
2012
2013    fn take_ckpt(&mut self) -> VerifyCkpt {
2014        self.ckpt
2015            .take()
2016            .expect("controller verify checkpoint already consumed")
2017    }
2018
2019    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
2020        self.eager_seed.take()
2021    }
2022
2023    fn settle(&mut self) {
2024        self.settled = true;
2025    }
2026}
2027
2028impl Drop for OptiControllerTicket {
2029    fn drop(&mut self) {
2030        if !self.settled {
2031            let _ = self.drain.synchronize();
2032            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2033        }
2034    }
2035}
2036
2037impl OptiForkTicket {
2038    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2039        self.boundary
2040            .take()
2041            .expect("fork ticket boundary already consumed")
2042    }
2043
2044    fn settle(&mut self) {
2045        self.settled = true;
2046    }
2047}
2048
2049impl Drop for OptiForkTicket {
2050    fn drop(&mut self) {
2051        if !self.settled {
2052            let _ = self.drain.synchronize();
2053            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2054        }
2055    }
2056}
2057
2058impl OptiForkState {
2059    #[allow(clippy::too_many_arguments)]
2060    fn new(
2061        e: &Engine,
2062        cache: &Cache,
2063        mode: OptiForkGateMode,
2064        alternate_snapshot: crate::cache::CacheSnapshot,
2065        h_seed: &CudaSlice<f32>,
2066        fill_prev: &CudaSlice<f32>,
2067        rt: &'static crate::pp::PpNRt,
2068        split: usize,
2069        n_layer: usize,
2070    ) -> Result<Self, Box<dyn std::error::Error>> {
2071        let fence = [0, split, n_layer];
2072        let mut logical_payload_bytes = [0usize; 2];
2073        for stage in 0..2 {
2074            for il in fence[stage]..fence[stage + 1] {
2075                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
2076                    .as_ref()
2077                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2078                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
2079                    .as_ref()
2080                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2081            }
2082        }
2083        let seeds = [
2084            OptiForkSeedGeneration {
2085                h_seed: e.clone_dtod(h_seed)?,
2086                fill_prev: e.clone_dtod(fill_prev)?,
2087                scratch_len: 0,
2088            },
2089            OptiForkSeedGeneration {
2090                h_seed: e.clone_dtod(h_seed)?,
2091                fill_prev: e.clone_dtod(fill_prev)?,
2092                scratch_len: 0,
2093            },
2094        ];
2095        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
2096            let _stage = rt.enter(0);
2097            let e0 = rt.engine(0, e);
2098            (
2099                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
2100                e0.htod_i32(&vec![0; split])?,
2101                e0.alloc_u32_zeroed(2)?,
2102                e0.alloc_u32_zeroed(1)?,
2103                e0.stream(),
2104            )
2105        };
2106        logical_payload_bytes[0] += seeds
2107            .iter()
2108            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
2109            .sum::<usize>();
2110        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
2111            + saved_lens.len() * std::mem::size_of::<i32>()
2112            + forced_acc.len() * std::mem::size_of::<u32>()
2113            + valid.len() * std::mem::size_of::<u32>();
2114        Ok(Self {
2115            mode,
2116            controller: (mode == OptiForkGateMode::Controller)
2117                .then(OptiControllerPolicy::configured),
2118            generations: OptiForkGenerationTracker::default(),
2119            active_snapshot_slot: 0,
2120            alternate_snapshot,
2121            seeds,
2122            rt,
2123            fence,
2124            split,
2125            len_ptrs,
2126            saved_lens,
2127            forced_acc,
2128            valid,
2129            stage0_stream,
2130            logical_payload_bytes,
2131        })
2132    }
2133
2134    fn reserve(
2135        &mut self,
2136        current_snapshot: &mut crate::cache::CacheSnapshot,
2137    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2138        let generation = self.generations.reserve()?;
2139        if generation.slot != self.active_snapshot_slot {
2140            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2141            self.active_snapshot_slot = generation.slot;
2142        }
2143        Ok(generation)
2144    }
2145
2146    fn capture_seed(
2147        &mut self,
2148        e: &Engine,
2149        generation: OptiForkGeneration,
2150        h_seed: &CudaSlice<f32>,
2151        fill_prev: &CudaSlice<f32>,
2152        scratch_len: usize,
2153    ) -> Result<(), Box<dyn std::error::Error>> {
2154        let seed = &mut self.seeds[generation.slot];
2155        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
2156        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
2157        seed.scratch_len = scratch_len;
2158        Ok(())
2159    }
2160
2161    fn ticket(
2162        &self,
2163        generation: OptiForkGeneration,
2164        boundary: VerifyBoundaryTicket,
2165    ) -> OptiForkTicket {
2166        OptiForkTicket {
2167            generation,
2168            boundary: Some(boundary),
2169            drain: self.stage0_stream.clone(),
2170            settled: false,
2171        }
2172    }
2173
2174    #[allow(clippy::too_many_arguments)]
2175    fn controller_ticket(
2176        &self,
2177        generation: OptiForkGeneration,
2178        boundary: VerifyBoundaryTicket,
2179        ckpt: VerifyCkpt,
2180        verify_tokens: [u32; 2],
2181        draft_prob: f32,
2182        eager_seed: Option<CudaSlice<f32>>,
2183        q_proxy: f32,
2184        scratch_len: usize,
2185    ) -> OptiControllerTicket {
2186        OptiControllerTicket {
2187            generation,
2188            boundary: Some(boundary),
2189            ckpt: Some(ckpt),
2190            verify_tokens,
2191            draft_prob,
2192            eager_seed,
2193            q_proxy,
2194            scratch_len,
2195            issued_at: std::time::Instant::now(),
2196            drain: self.stage0_stream.clone(),
2197            settled: false,
2198        }
2199    }
2200
2201    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2202        self.generations.reserve()
2203    }
2204
2205    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
2206        &mut self.alternate_snapshot
2207    }
2208
2209    fn promote_successor_snapshot(
2210        &mut self,
2211        current_snapshot: &mut crate::cache::CacheSnapshot,
2212        generation: OptiForkGeneration,
2213    ) {
2214        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2215        self.active_snapshot_slot = generation.slot;
2216    }
2217
2218    fn queue_actual_reconcile(
2219        &mut self,
2220        e: &Engine,
2221        snapshot: &crate::cache::CacheSnapshot,
2222        acc: &CudaSlice<u32>,
2223        optimistic_pending: u32,
2224        base: usize,
2225    ) -> Result<(), Box<dyn std::error::Error>> {
2226        let saved: Vec<i32> = (0..self.split)
2227            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2228            .collect();
2229        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
2230        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
2231        // the validity/reconcile kernels must never peer-read acc before it is written. The
2232        // increment-1 harness uses primary stage 0, where stream order already provides this.
2233        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
2234            self.rt.fence_stages_behind(&e.stream())?;
2235        }
2236        let _stage = self.rt.enter(0);
2237        let e0 = self.rt.engine(0, e);
2238        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2239        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
2240        e0.spec_fork_reconcile_kv(
2241            &self.len_ptrs,
2242            &self.saved_lens,
2243            acc,
2244            &self.valid,
2245            base,
2246            self.split,
2247        )
2248    }
2249
2250    fn finish_actual_reconcile(
2251        &mut self,
2252        e: &Engine,
2253        cache: &mut Cache,
2254        snapshot: &crate::cache::CacheSnapshot,
2255        n_acc: usize,
2256        base: usize,
2257        hit: bool,
2258    ) -> Result<(), Box<dyn std::error::Error>> {
2259        if hit {
2260            return Ok(());
2261        }
2262        let len_delta = base + n_acc;
2263        for il in 0..self.split {
2264            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2265                kv.len = saved + len_delta;
2266            }
2267        }
2268        {
2269            let _stage = self.rt.enter(1);
2270            let e1 = self.rt.engine(1, e);
2271            for il in self.split..self.fence[2] {
2272                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2273                    kv.len = saved + len_delta;
2274                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2275                }
2276            }
2277        }
2278        self.rt.publish_to(0, &e.stream())?;
2279        Ok(())
2280    }
2281
2282    fn cancel_controller_ticket(
2283        &mut self,
2284        e: &Engine,
2285        cache: &mut Cache,
2286        scratch: &mut MtpScratch,
2287        snapshot: &crate::cache::CacheSnapshot,
2288        ticket: &mut OptiControllerTicket,
2289    ) -> Result<(), Box<dyn std::error::Error>> {
2290        {
2291            let _stage = self.rt.enter(0);
2292            let e0 = self.rt.engine(0, e);
2293            for il in 0..self.split {
2294                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2295                    kv.len = saved;
2296                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
2297                }
2298            }
2299        }
2300        scratch.set_len(e, snapshot.pos)?;
2301        ticket.settle();
2302        self.generations.retire(ticket.generation)?;
2303        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2304        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
2305        eprintln!(
2306            "[opti-controller] tail-drain generation={} slot={}",
2307            ticket.generation.id, ticket.generation.slot,
2308        );
2309        Ok(())
2310    }
2311
2312    #[allow(clippy::too_many_arguments)]
2313    fn reconcile(
2314        &mut self,
2315        e: &Engine,
2316        cache: &mut Cache,
2317        scratch: &mut MtpScratch,
2318        snapshot: &crate::cache::CacheSnapshot,
2319        h_seed: &mut CudaSlice<f32>,
2320        fill_prev: &mut CudaSlice<f32>,
2321        generation: OptiForkGeneration,
2322        action: OptiForkAction,
2323        optimistic_pending: u32,
2324    ) -> Result<(), Box<dyn std::error::Error>> {
2325        debug_assert!(action != OptiForkAction::Abort);
2326        let miss_started = std::time::Instant::now();
2327        let keep = action == OptiForkAction::Hit;
2328        let saved: Vec<i32> = (0..self.split)
2329            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2330            .collect();
2331        let seed = &self.seeds[generation.slot];
2332        {
2333            let _stage = self.rt.enter(0);
2334            let e0 = self.rt.engine(0, e);
2335            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2336            let forced = if keep {
2337                [1u32, optimistic_pending]
2338            } else {
2339                [0u32, optimistic_pending]
2340            };
2341            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
2342            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
2343            e0.spec_fork_reconcile_kv(
2344                &self.len_ptrs,
2345                &self.saved_lens,
2346                &self.forced_acc,
2347                &self.valid,
2348                0,
2349                self.split,
2350            )?;
2351            for il in 0..self.split {
2352                if let Some(recur) = cache.recur[il].as_mut() {
2353                    let conv = snapshot.conv[il]
2354                        .as_ref()
2355                        .ok_or("optipipe stage0 snapshot missing conv state")?;
2356                    let ssm = snapshot.ssm[il]
2357                        .as_ref()
2358                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
2359                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2360                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2361                }
2362            }
2363            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2364            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2365        }
2366
2367        if keep {
2368            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2369            return Ok(());
2370        }
2371
2372        for il in 0..self.split {
2373            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2374                kv.len = saved;
2375            }
2376        }
2377        scratch.set_len(e, seed.scratch_len)?;
2378        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2379        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2380        let caller = e.stream();
2381        self.rt.publish_to(0, &caller)?;
2382        caller.synchronize()?;
2383        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2384        eprintln!(
2385            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2386            generation.id, generation.slot,
2387        );
2388        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2389        Ok(())
2390    }
2391
2392    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2393        self.generations.retire(generation)
2394    }
2395}
2396
2397impl HybridModel {
2398    fn opti_graph_draft_step(
2399        &self,
2400        e: &Engine,
2401        mtp: &MtpHead,
2402        dctx: &mut DraftGraphCtx,
2403        scratch: &mut MtpScratch,
2404        d_vocab: usize,
2405    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2406        dctx.graph
2407            .as_ref()
2408            .ok_or("optipipe controller requires the greedy draft graph")?
2409            .launch()?;
2410        scratch.kv.len += 1;
2411        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2412        if (idx as usize) >= d_vocab {
2413            return Err(
2414                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2415            );
2416        }
2417        let probability = e.dtoh(&dctx.g_p)?[0];
2418        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2419            return Err(format!("optipipe draft probability is invalid: {probability}").into());
2420        }
2421        let token = match &mtp.d2t {
2422            Some(map) => map[idx as usize],
2423            None => idx,
2424        };
2425        if token != idx {
2426            e.set_u32_one(&mut dctx.g_tok, token)?;
2427        }
2428        Ok((token, probability))
2429    }
2430
2431    #[allow(clippy::too_many_arguments)]
2432    fn opti_controller_draft_step(
2433        &self,
2434        e: &Engine,
2435        mtp: &MtpHead,
2436        dctx: &mut DraftGraphCtx,
2437        scratch: &mut MtpScratch,
2438        d_vocab: usize,
2439        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2440        eager_pos: usize,
2441        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2442    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2443        if dctx.graph.is_some() {
2444            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2445        }
2446        let (input_token, input_seed) = eager_state
2447            .take()
2448            .ok_or("optipipe eager continuation seed is unavailable")?;
2449        let (logits, next_seed) = self.mtp_head_forward_dev(
2450            e,
2451            mtp,
2452            input_token,
2453            &input_seed,
2454            scratch,
2455            eager_pos,
2456            embd_dev,
2457            None,
2458        )?;
2459        let token_d = e.argmax_token_device(&logits, d_vocab)?;
2460        let idx = e.dtoh_u32_one(&token_d)?;
2461        if (idx as usize) >= d_vocab {
2462            return Err(format!(
2463                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2464            )
2465            .into());
2466        }
2467        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2468        let probability = e.dtoh(&probability_d)?[0];
2469        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2470            return Err(
2471                format!("optipipe eager draft probability is invalid: {probability}").into(),
2472            );
2473        }
2474        let token = match &mtp.d2t {
2475            Some(map) => map[idx as usize],
2476            None => idx,
2477        };
2478        *eager_state = Some((token, next_seed));
2479        Ok((token, probability))
2480    }
2481
2482    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2483    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2484    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2485    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2486    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2487    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2488    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2489    /// transfer + host argmax per draft token from the K-token draft chain.
2490    #[allow(clippy::too_many_arguments)]
2491    fn mtp_head_forward_dev(
2492        &self,
2493        e: &Engine,
2494        mtp: &MtpHead,
2495        e_tok: u32,
2496        h_seed: &CudaSlice<f32>,
2497        scratch: &mut MtpScratch,
2498        mtp_pos: usize,
2499        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2500        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2501        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2502        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2503        mask: Option<(&CudaSlice<u32>, usize)>,
2504    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2505        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
2506        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
2507        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
2508        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
2509        static ANAT_NS: [AtomicU64; 5] = [
2510            AtomicU64::new(0),
2511            AtomicU64::new(0),
2512            AtomicU64::new(0),
2513            AtomicU64::new(0),
2514            AtomicU64::new(0),
2515        ];
2516        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
2517        let anat = {
2518            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2519            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
2520        };
2521        if anat {
2522            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
2523        }
2524        let t_all = std::time::Instant::now();
2525        let mut t_ph = std::time::Instant::now();
2526        let mut anat_mark = |i: usize,
2527                             e: &Engine,
2528                             t: &mut std::time::Instant|
2529         -> Result<(), Box<dyn std::error::Error>> {
2530            if anat {
2531                e.stream().synchronize()?;
2532                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
2533                *t = std::time::Instant::now();
2534            }
2535            Ok(())
2536        };
2537        let cfg = &self.cfg;
2538        let n_embd = cfg.n_embd as usize;
2539        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2540        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2541        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2542        let eps = cfg.rms_eps;
2543        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2544
2545        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2546        // expands this one row on CPU and transfers n_embd f32 values instead.
2547        let e_emb = match embd_dev {
2548            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2549            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2550        };
2551
2552        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2553        let mut e_norm = e.zeros(n_embd)?;
2554        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2555        let mut h_norm = e.zeros(n_embd)?;
2556        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2557
2558        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2559        let mut concat = e.zeros(2 * n_embd)?;
2560        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2561        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2562
2563        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2564        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2565
2566        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2567        let mut a_norm = e.zeros(di)?;
2568        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2569        anat_mark(0, e, &mut t_ph)?;
2570
2571        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2572        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2573        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2574        // advances only the device counter).
2575        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2576            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2577            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2578            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2579            // whose host-side mirror the caller does).
2580            (Mixer::Full(fa), Some(g)) => {
2581                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2582            }
2583            (Mixer::Full(fa), None) => {
2584                let out =
2585                    self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2586                scratch.kv.len += 1;
2587                out
2588            }
2589            (Mixer::Linear(_), _) => {
2590                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2591            }
2592            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2593        };
2594        anat_mark(1, e, &mut t_ph)?;
2595
2596        // op 7: x1 = inpSA + attn_out
2597        let mut x1 = e.zeros(di)?;
2598        e.add(&inp_sa, &attn_out, &mut x1, di)?;
2599
2600        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
2601        let mut z = e.zeros(di)?;
2602        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2603
2604        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2605        let ffn_out = match &mtp.ffn {
2606            crate::hybrid::Ffn::Dense {
2607                ffn_gate,
2608                ffn_up,
2609                ffn_down,
2610            } => {
2611                let n_ff = ffn_gate.out_features();
2612                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2613                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2614                    (
2615                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2616                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2617                    )
2618                } else {
2619                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2620                };
2621                let mut act = e.zeros(n_ff)?;
2622                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2623                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2624                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2625                // passes None, which is `ffn_act`'s dispatch verbatim.
2626                Self::ffn_act_lim(
2627                    e,
2628                    &self.cfg,
2629                    &gate,
2630                    &up,
2631                    1.0,
2632                    1.0,
2633                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2634                    &mut act,
2635                    n_ff,
2636                )?;
2637                e.matmul(ffn_down, &act, 1)?
2638            }
2639            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
2640            // so they never alias trunk layer 0's cache keys.
2641            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
2642        };
2643        anat_mark(2, e, &mut t_ph)?;
2644
2645        // op 10: h_nextn = x1 + ffn_out (at di)
2646        let mut h_inner = e.zeros(di)?;
2647        e.add(&x1, &ffn_out, &mut h_inner, di)?;
2648
2649        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
2650        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
2651        let h_nextn = match mtp.geom.as_ref() {
2652            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
2653            None => h_inner,
2654        };
2655
2656        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
2657        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
2658        let mut final_h = e.zeros(n_embd)?;
2659        e.rms_norm(
2660            &h_nextn,
2661            final_norm.float_data(),
2662            &mut final_h,
2663            n_embd,
2664            1,
2665            eps,
2666        )?;
2667
2668        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
2669        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
2670        let mut logits = e.matmul(head, &final_h, 1)?;
2671        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
2672        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
2673        if let Some((mask_d, mw)) = mask {
2674            let d_vocab = head.out_features();
2675            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
2676        }
2677        anat_mark(3, e, &mut t_ph)?;
2678        if anat {
2679            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
2680            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
2681            if n % 128 == 0 {
2682                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
2683                eprintln!(
2684                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
2685                    us(0),
2686                    us(1),
2687                    us(2),
2688                    us(3),
2689                    us(4)
2690                );
2691            }
2692        }
2693        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
2694        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
2695        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
2696    }
2697
2698    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
2699    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
2700    /// the dc path, and all three are properties of this arch's MTP block:
2701    ///
2702    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
2703    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
2704    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
2705    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
2706    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
2707    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
2708    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
2709    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
2710    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
2711    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
2712    ///    resolved `Step35MtpGeom`, never from `cfg`.
2713    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
2714    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
2715    ///    fused-into-wq `q_gate_split` form the dc arm handles.
2716    ///
2717    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
2718    /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
2719    /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
2720    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
2721    ///
2722    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
2723    /// caller must not mirror.
2724    fn mtp_step35_attn(
2725        &self,
2726        e: &Engine,
2727        fa: &FullAttnLayer,
2728        g: &crate::hybrid::Step35MtpGeom,
2729        h: &CudaSlice<f32>,
2730        pos_d: &CudaSlice<i32>,
2731        scratch: &mut MtpScratch,
2732    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2733        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
2734        let eps = self.cfg.rms_eps;
2735        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
2736        let n_embd = self.cfg.n_embd as usize;
2737        let gw = fa
2738            .attn_gate
2739            .as_ref()
2740            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
2741
2742        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
2743            && e.uses_q8_1_fast(&fa.wk)
2744            && e.uses_q8_1_fast(&fa.wv)
2745            && e.uses_q8_1_fast(gw)
2746        {
2747            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
2748            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
2749                Some(t3) => t3,
2750                None => (
2751                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
2752                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
2753                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
2754                ),
2755            };
2756            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
2757        } else {
2758            (
2759                e.matmul(&fa.wq, h, 1)?,
2760                e.matmul(&fa.wk, h, 1)?,
2761                e.matmul(&fa.wv, h, 1)?,
2762                e.matmul(gw, h, 1)?,
2763            )
2764        };
2765
2766        let mut q = e.uninit(nh * hd)?;
2767        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
2768        let mut k = e.uninit(nkv * hd)?;
2769        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
2770        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
2771        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
2772        // the resolved flag, not the constant, so an all-full sibling stays correct.
2773        let ff = if g.swa {
2774            None
2775        } else {
2776            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2777        };
2778        #[cfg(debug_assertions)]
2779        if let Some(ff) = ff {
2780            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
2781        }
2782        e.rope_neox2(
2783            &mut q,
2784            &mut k,
2785            pos_d,
2786            hd,
2787            g.n_rot,
2788            nh,
2789            nkv,
2790            1,
2791            g.rope_base,
2792            1.0,
2793            ff,
2794        )?;
2795
2796        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
2797        // length on the host anyway, and the windowed view below needs it there to compute the
2798        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
2799        // dc-family consumer of this scratch still agree.
2800        let kv = &mut scratch.kv;
2801        assert!(
2802            kv.len < scratch.cap,
2803            "step35 MTP scratch overflow ({} >= {})",
2804            kv.len,
2805            scratch.cap
2806        );
2807        let next_len = kv.len + 1;
2808        let (off, t_kv) = if g.swa && next_len > g.window {
2809            (next_len - g.window, g.window)
2810        } else {
2811            (0, next_len)
2812        };
2813        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
2814        e.append_kv_quantized(
2815            &k,
2816            &v0,
2817            &mut kv.k,
2818            &mut kv.v,
2819            write_row,
2820            kv.kv_dim_k,
2821            kv.kv_dim_v,
2822            kv.k_tok_bytes,
2823            kv.v_tok_bytes,
2824            false,
2825        )?;
2826        kv.len = next_len;
2827        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2828        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
2829        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
2830        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
2831        // therefore live, not theoretical.
2832        let physical = kv.physical_rows(off, off + t_kv)?;
2833        let k_view = e.view_u8_range(
2834            &kv.k,
2835            physical.start * kv.k_tok_bytes,
2836            physical.end * kv.k_tok_bytes,
2837        );
2838        let v_view = e.view_u8_range(
2839            &kv.v,
2840            physical.start * kv.v_tok_bytes,
2841            physical.end * kv.v_tok_bytes,
2842        );
2843        let mut attn = e.uninit(nh * hd)?;
2844        e.fa_decode_kvmod(
2845            &q,
2846            &k_view,
2847            &v_view,
2848            &mut attn,
2849            hd,
2850            nh,
2851            nkv,
2852            t_kv,
2853            scale,
2854            kv.k_tok_bytes,
2855            kv.v_tok_bytes,
2856            false,
2857        )?;
2858
2859        let mut ag = e.uninit(nh * hd)?;
2860        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
2861        Ok(e.matmul(&fa.wo, &ag, 1)?)
2862    }
2863
2864    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
2865    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
2866    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
2867    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
2868    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
2869    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
2870    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
2871    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
2872    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
2873    fn mtp_full_attn_dc(
2874        &self,
2875        e: &Engine,
2876        fa: &FullAttnLayer,
2877        h: &CudaSlice<f32>,
2878        pos_d: &CudaSlice<i32>,
2879        scratch: &mut MtpScratch,
2880        geom: Option<&crate::hybrid::DraftGeom>,
2881    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2882        let cfg = &self.cfg;
2883        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
2884        let geometry = cfg.full_attention_geometry_at(mtp_il);
2885        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
2886        let n_head_kv = geom
2887            .map(|g| g.n_head_kv)
2888            .unwrap_or(geometry.n_head_kv as usize);
2889        let head_dim = geometry.head_dim_k as usize;
2890        let eps = cfg.rms_eps;
2891        let scale = geometry.attention_scale();
2892        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
2893        let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
2894
2895        let (qf, mut k, v) =
2896            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
2897                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
2898                (
2899                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
2900                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
2901                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
2902                )
2903            } else {
2904                (
2905                    e.matmul(&fa.wq, h, 1)?,
2906                    e.matmul(&fa.wk, h, 1)?,
2907                    e.matmul(&fa.wv, h, 1)?,
2908                )
2909            };
2910        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
2911        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
2912        let (mut q, gate) = if gated {
2913            let mut q = e.zeros(n_head * head_dim)?;
2914            let mut gate = e.zeros(n_head * head_dim)?;
2915            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
2916            (q, Some(gate))
2917        } else {
2918            (qf, None)
2919        };
2920
2921        let mut qn = e.zeros(n_head * head_dim)?;
2922        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
2923        q = qn;
2924        let mut kn = e.zeros(n_head_kv * head_dim)?;
2925        e.rms_norm(
2926            &k,
2927            fa.k_norm.float_data(),
2928            &mut kn,
2929            head_dim,
2930            n_head_kv,
2931            eps,
2932        )?;
2933        k = kn;
2934        let rope_dims = geometry.n_rot as usize;
2935        e.rope_neox(
2936            &mut q,
2937            pos_d,
2938            head_dim,
2939            rope_dims,
2940            n_head,
2941            1,
2942            geometry.rope_base,
2943            1.0,
2944        )?;
2945        e.rope_neox(
2946            &mut k,
2947            pos_d,
2948            head_dim,
2949            rope_dims,
2950            n_head_kv,
2951            1,
2952            geometry.rope_base,
2953            1.0,
2954        )?;
2955
2956        let kv = &mut scratch.kv;
2957        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
2958        e.append_kv_quantized_dc(
2959            &k,
2960            &v,
2961            &mut kv.k,
2962            &mut kv.v,
2963            &kv.len_d,
2964            kv.kv_dim_k,
2965            kv.kv_dim_v,
2966            kv.k_tok_bytes,
2967            kv.v_tok_bytes,
2968            false,
2969        )?;
2970        e.inc_seqlen(&mut kv.len_d)?;
2971        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
2972        // key range from the device counter.
2973        let k_view = e.view_u8(&kv.k, kv.k.len());
2974        let v_view = e.view_u8(&kv.v, kv.v.len());
2975        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
2976        let mut attn = e.zeros(n_head * head_dim)?;
2977        e.fa_decode_dc(
2978            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
2979            scale, ktb, vtb, false,
2980        )?;
2981
2982        let attn_g = match &gate {
2983            Some(gate) => {
2984                let mut gsig = e.zeros(n_head * head_dim)?;
2985                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
2986                let mut ag = e.zeros(n_head * head_dim)?;
2987                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
2988                ag
2989            }
2990            None => attn,
2991        };
2992        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
2993    }
2994
2995    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
2996    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
2997    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
2998    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
2999    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
3000    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
3001    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
3002    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
3003    #[allow(clippy::too_many_arguments)]
3004    fn mtp_kv_fill(
3005        &self,
3006        e: &Engine,
3007        mtp: &MtpHead,
3008        tokens: &[u32],
3009        h: &CudaSlice<f32>,
3010        pos0: usize,
3011        scratch: &mut MtpScratch,
3012        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3013    ) -> Result<(), Box<dyn std::error::Error>> {
3014        let cfg = &self.cfg;
3015        let n_embd = cfg.n_embd as usize;
3016        let eps = cfg.rms_eps;
3017        let t = tokens.len();
3018        assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
3019        assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
3020        let Mixer::Full(fa) = &mtp.mixer else {
3021            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3022        };
3023        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
3024        let pos_d = e.htod_i32(&pos_vec)?;
3025
3026        // ops A/1/2: embed + the two input norms, T-wide.
3027        let e_emb = match embd_dev {
3028            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3029            None => e.htod(&self.embd.gather(n_embd, tokens))?,
3030        };
3031        let mut e_norm = e.zeros(t * n_embd)?;
3032        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
3033        let mut h_norm = e.zeros(t * n_embd)?;
3034        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
3035
3036        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
3037        let mut concat = e.zeros(t * 2 * n_embd)?;
3038        for i in 0..t {
3039            e.copy_view_into(
3040                &mut concat,
3041                i * 2 * n_embd,
3042                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
3043                n_embd,
3044            )?;
3045            e.copy_view_into(
3046                &mut concat,
3047                i * 2 * n_embd + n_embd,
3048                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
3049                n_embd,
3050            )?;
3051        }
3052
3053        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
3054        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3055        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
3056        let mut a_norm = e.zeros(t * di)?;
3057        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
3058
3059        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
3060        // the fill only has to leave correct K/V rows behind for later chains to attend over.
3061        let n_head_kv = mtp
3062            .geom
3063            .as_ref()
3064            .map(|g| g.n_head_kv)
3065            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
3066            .unwrap_or_else(|| {
3067                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3068                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
3069            });
3070        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3071        let geometry = cfg.full_attention_geometry_at(mtp_il);
3072        let head_dim = geometry.head_dim_k as usize;
3073        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
3074        let v = e.matmul(&fa.wv, &a_norm, t)?;
3075        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
3076        e.rms_norm(
3077            &k,
3078            fa.k_norm.float_data(),
3079            &mut kn,
3080            head_dim,
3081            n_head_kv * t,
3082            eps,
3083        )?;
3084        k = kn;
3085        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
3086        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
3087        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
3088        // writes K rows the attention arm then re-derives at a different theta: correct-looking
3089        // output with dead acceptance, invisible to the exactness gates.
3090        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
3091            Some(s) => (
3092                s.n_rot,
3093                s.rope_base,
3094                if s.swa {
3095                    None
3096                } else {
3097                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3098                },
3099            ),
3100            None => (geometry.n_rot as usize, geometry.rope_base, None),
3101        };
3102        #[cfg(debug_assertions)]
3103        if let Some(ff) = ff {
3104            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
3105        }
3106        match ff {
3107            Some(f) => e.rope_neox_ff(
3108                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
3109            )?,
3110            None => e.rope_neox(
3111                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3112            )?,
3113        }
3114
3115        let kv = &mut scratch.kv;
3116        // Match the trunk prime contract: a chunk may need the aligned window immediately before
3117        // its first row, so preserve that prefix when the physical tail rebases at wrap.
3118        let retain_from = kv
3119            .ring
3120            .as_ref()
3121            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
3122            .unwrap_or(0);
3123        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
3124        for i in 0..t {
3125            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
3126            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
3127            e.append_kv_quantized_view(
3128                &k_row,
3129                &v_row,
3130                &mut kv.k,
3131                &mut kv.v,
3132                write_row + i,
3133                kv.kv_dim_k,
3134                kv.kv_dim_v,
3135                kv.k_tok_bytes,
3136                kv.v_tok_bytes,
3137                false,
3138            )?;
3139        }
3140        kv.len = pos0 + t;
3141        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3142        Ok(())
3143    }
3144
3145    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
3146    /// every varying input device-resident —
3147    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
3148    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
3149    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
3150    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
3151    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
3152    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
3153    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
3154    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
3155    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
3156    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
3157    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
3158    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
3159    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
3160    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
3161    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
3162    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
3163    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
3164    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
3165    #[allow(clippy::too_many_arguments)]
3166    fn mtp_head_forward_cap(
3167        &self,
3168        e: &Engine,
3169        mtp: &MtpHead,
3170        tok_d: &mut CudaSlice<u32>,
3171        pos_d: &mut CudaSlice<i32>,
3172        h_seed_d: &mut CudaSlice<f32>,
3173        p_d: &mut CudaSlice<f32>,
3174        scratch: &mut MtpScratch,
3175        with_prob: bool,
3176        with_head: bool,
3177        embd_gpu: &CudaSlice<u8>,
3178        embd_qt: i32,
3179        embd_rb: usize,
3180        d_vocab: usize,
3181        sampled_cap: Option<(
3182            &mut CudaSlice<u32>,
3183            &mut CudaSlice<f32>,
3184            &mut CudaSlice<f32>,
3185            u64,
3186            f32,
3187        )>,
3188        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
3189        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
3190        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
3191        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
3192        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
3193        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
3194        mask_cap: Option<(&CudaSlice<u32>, usize)>,
3195    ) -> Result<(), Box<dyn std::error::Error>> {
3196        let cfg = &self.cfg;
3197        let n_embd = cfg.n_embd as usize;
3198        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
3199        // whose device-counter key bound always starts at row 0 — it cannot express this block's
3200        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
3201        // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
3202        // refuses step35 heads explicitly (SWA refusal), so the eager chain
3203        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
3204        // panic) is what the two capture sites and the round-stream capture already handle by
3205        // degrading to eager / stream-off.
3206        if mtp.step35.is_some() {
3207            return Err(
3208                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
3209                        block's SWA view offset; same root cause as the dc decode refusal) — the \
3210                        eager draft chain serves this arch"
3211                    .into(),
3212            );
3213        }
3214        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
3215        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3216        let eps = cfg.rms_eps;
3217        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
3218        let mut e_norm = e.zeros(n_embd)?;
3219        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3220        let mut h_norm = e.zeros(n_embd)?;
3221        e.rms_norm(
3222            &*h_seed_d,
3223            mtp.hnorm.float_data(),
3224            &mut h_norm,
3225            n_embd,
3226            1,
3227            eps,
3228        )?;
3229        let mut concat = e.zeros(2 * n_embd)?;
3230        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3231        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3232        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3233        let mut a_norm = e.zeros(di)?;
3234        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3235        let attn_out = match &mtp.mixer {
3236            Mixer::Full(fa) => {
3237                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
3238            }
3239            Mixer::Linear(_) => {
3240                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3241            }
3242            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3243        };
3244        let mut x1 = e.zeros(di)?;
3245        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3246        let mut z = e.zeros(di)?;
3247        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3248        let ffn_out = match &mtp.ffn {
3249            crate::hybrid::Ffn::Dense {
3250                ffn_gate,
3251                ffn_up,
3252                ffn_down,
3253            } => {
3254                let n_ff = ffn_gate.out_features();
3255                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3256                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3257                    (
3258                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3259                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3260                    )
3261                } else {
3262                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3263                };
3264                let mut act = e.zeros(n_ff)?;
3265                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
3266                e.matmul(ffn_down, &act, 1)?
3267            }
3268            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
3269            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
3270            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
3271            // error arm degrades the caller to eager/stream-off.
3272            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
3273                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
3274            }
3275            crate::hybrid::Ffn::Moe(_) => {
3276                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
3277            }
3278        };
3279        let mut h_inner = e.zeros(di)?;
3280        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3281        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
3282        let h_nextn = match mtp.geom.as_ref() {
3283            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3284            None => h_inner,
3285        };
3286        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
3287        let final_h = if with_head || spec_hpost() {
3288            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3289            let mut fh = e.zeros(n_embd)?;
3290            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
3291            Some(fh)
3292        } else {
3293            None
3294        };
3295        if with_head {
3296            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3297            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
3298            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
3299            // before the argmax — proposals become legal by construction. Contents-only
3300            // per-replay upload keeps the capture valid.
3301            if let Some((mask_d, mw)) = mask_cap {
3302                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3303            }
3304            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
3305                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
3306                // own buffer is pool-recycled after the capture body returns, so it can't be the
3307                // retention target), bump the device event counter, gumbel-perturb reading it,
3308                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
3309                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
3310                e.sctr_inc(ctr_d)?;
3311                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
3312                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
3313                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
3314                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
3315                if with_prob {
3316                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3317                }
3318            } else {
3319                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
3320                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
3321                // p-min under a draft mask reads the MASKED row: confidence relative to the
3322                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
3323                // is the right semantics for "does the drafter know what comes next here" and
3324                // the same row the pick came from. Draft-quality only — verify arbitrates.
3325                if with_prob {
3326                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3327                }
3328            }
3329        }
3330        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
3331        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
3332        if let Some((out, slot, d2t)) = stream_pack {
3333            e.pack_tok_p(tok_d, p_d, out, slot)?;
3334            if let Some(map) = d2t {
3335                e.tok_map_u32(tok_d, map)?;
3336            }
3337        }
3338        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
3339        if spec_hpost() {
3340            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
3341        } else {
3342            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
3343        }
3344        // advance the draft rope position in-graph.
3345        e.inc_seqlen(pos_d)?;
3346        Ok(())
3347    }
3348
3349    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
3350    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
3351    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
3352    /// Advances `cache.pos` by T.
3353    pub fn decode_step_t(
3354        &self,
3355        e: &Engine,
3356        tokens: &[u32],
3357        pos0: usize,
3358        cache: &mut Cache,
3359    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3360        if self.is_gemma4_e4b() {
3361            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
3362        }
3363        if self.cfg.gemma4.is_some() {
3364            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
3365        }
3366        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
3367    }
3368
3369    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
3370    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
3371    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
3372    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
3373    pub fn decode_step_t_h(
3374        &self,
3375        e: &Engine,
3376        tokens: &[u32],
3377        pos0: usize,
3378        cache: &mut Cache,
3379    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3380        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
3381    }
3382
3383    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
3384    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
3385    pub fn decode_step_t_h_emb(
3386        &self,
3387        e: &Engine,
3388        tokens: &[u32],
3389        pos0: usize,
3390        cache: &mut Cache,
3391        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3392    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3393        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
3394        Ok((e.dtoh(&logits_d)?, h_seed))
3395    }
3396
3397    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
3398    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
3399    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
3400    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
3401    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
3402    pub fn decode_step_t_h_emb_dev(
3403        &self,
3404        e: &Engine,
3405        tokens: &[u32],
3406        pos0: usize,
3407        cache: &mut Cache,
3408        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3409    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3410        let n_embd = self.cfg.n_embd as usize;
3411        let t = tokens.len();
3412        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3413        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3414        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3415        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3416        Ok((logits, hs))
3417    }
3418
3419    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3420    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3421    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3422    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3423    /// retains/copies — they never change what any kernel computes).
3424    fn decode_step_t_core(
3425        &self,
3426        e: &Engine,
3427        tokens: &[u32],
3428        pos0: usize,
3429        cache: &mut Cache,
3430        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3431        mut ckpt: Option<&mut VerifyCkpt>,
3432    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3433        self.decode_step_t_core_stream(e, tokens, pos0, cache, embd_dev, ckpt.take(), None, None)
3434    }
3435
3436    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3437    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3438    fn decode_step_t_core_pipelined(
3439        &self,
3440        e: &Engine,
3441        tokens: &[u32],
3442        pos0: usize,
3443        cache: &mut Cache,
3444        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3445        mut ckpt: Option<&mut VerifyCkpt>,
3446        pipe: &SpecPipeLane,
3447        round: usize,
3448    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3449        let fence = crate::pp::pp_cuts(self.layers.len())
3450            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3451        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3452            return Err("two-session speculative pipeline requires the PP verify split".into());
3453        }
3454        let interval_fence = pipe.stage0_begin(round)?;
3455        let ticket = self.verify_stage0_issue(
3456            e,
3457            tokens,
3458            pos0,
3459            cache,
3460            embd_dev,
3461            ckpt.as_deref_mut(),
3462            None,
3463            &fence,
3464            Some(interval_fence),
3465            pipe.trace(round),
3466        )?;
3467        pipe.stage0_end(round);
3468        pipe.stage1_begin(round)?;
3469        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3470        pipe.verify_end(round);
3471        Ok(result)
3472    }
3473
3474    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3475    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3476    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3477    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3478    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3479    #[allow(clippy::too_many_arguments)]
3480    fn decode_step_t_core_stream(
3481        &self,
3482        e: &Engine,
3483        tokens: &[u32],
3484        pos0: usize,
3485        cache: &mut Cache,
3486        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3487        mut ckpt: Option<&mut VerifyCkpt>,
3488        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3489        pp_pipe: Option<bool>,
3490    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3491        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3492        // exactly as the eager and batched steps do. This is the single funnel every verify
3493        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3494        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3495        // is untouched.
3496        //
3497        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3498        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3499        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3500        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3501        // or a placement whose PpNRt fails to build — so a config that would still walk the
3502        // whole trunk on one stream refuses instead of regressing 28x.
3503        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3504            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3505                return self.decode_step_t_core_ppn(
3506                    e,
3507                    tokens,
3508                    pos0,
3509                    cache,
3510                    embd_dev,
3511                    ckpt.take(),
3512                    stream,
3513                    &fence,
3514                    pp_pipe,
3515                );
3516            }
3517        }
3518        crate::pp::refuse_unsplit_if_remote(
3519            "decode_step_t (spec verify)",
3520            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3521             split (decode_step_t_core_ppn); or run spec on one device",
3522        )?;
3523        let cfg = &self.cfg;
3524        let n_embd = cfg.n_embd as usize;
3525        let eps = cfg.rms_eps;
3526        let t = tokens.len();
3527        let pos_d = match stream {
3528            Some((_, ctr)) => {
3529                let mut p = e.alloc_uninit::<i32>(t)?;
3530                e.pos_iota(ctr, &mut p, t)?;
3531                p
3532            }
3533            None => {
3534                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3535                e.htod_i32(&pos_vec)?
3536            }
3537        };
3538
3539        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3540        let x = match (stream, embd_dev) {
3541            (Some((vtok, _)), Some((g, qt, rb))) => {
3542                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3543            }
3544            (None, Some((g, qt, rb))) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3545            _ => e.htod(&self.embd.gather(n_embd, tokens))?,
3546        };
3547
3548        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3549        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3550        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3551        let x = self.verify_layers(
3552            e,
3553            x,
3554            0,
3555            self.layers.len(),
3556            &pos_d,
3557            pos0,
3558            t,
3559            cache,
3560            ckpt.take(),
3561            stream,
3562        )?;
3563
3564        let mut hn = vbuf(e, t * n_embd)?;
3565        let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3566        let logits = if serving_head {
3567            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3568            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3569            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3570            // serve one batched numeric class at every live width, including B=1. Keep the
3571            // verify head in that same class; other generic families retain the decode-exact
3572            // head that their run-spec contract pins.
3573            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3574            e.matmul(&self.output, &hn, t)?
3575        } else {
3576            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3577            e.matmul_decode_exact(&self.output, &hn, t)?
3578        };
3579        // stream: the device pos counter owns position; host mirror reconciles at drain.
3580        if stream.is_none() {
3581            cache.pos += t;
3582        }
3583        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3584        Ok((logits, if spec_hpost() { hn } else { x }))
3585    }
3586
3587    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3588    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3589    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3590    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3591    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3592    /// the payload).
3593    ///
3594    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3595    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3596    /// receipts):
3597    ///
3598    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
3599    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
3600    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
3601    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
3602    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
3603    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
3604    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
3605    ///
3606    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
3607    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
3608    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
3609    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
3610    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
3611    ///
3612    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
3613    ///    sharded loader leaves the table with stage 0 by construction).
3614    ///
3615    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
3616    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
3617    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
3618    ///    model, every round.
3619    ///
3620    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
3621    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
3622    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
3623    /// through the primary context by UVA — the same read the batched serving epilogue's
3624    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
3625    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
3626    ///
3627    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
3628    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
3629    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
3630    ///
3631    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
3632    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
3633    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
3634    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
3635    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
3636    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
3637    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
3638    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
3639    #[allow(clippy::too_many_arguments)]
3640    fn decode_step_t_core_ppn(
3641        &self,
3642        e: &Engine,
3643        tokens: &[u32],
3644        pos0: usize,
3645        cache: &mut Cache,
3646        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3647        mut ckpt: Option<&mut VerifyCkpt>,
3648        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3649        fence: &[usize],
3650        pp_pipe: Option<bool>,
3651    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3652        let ticket = self.verify_stage0_issue(
3653            e,
3654            tokens,
3655            pos0,
3656            cache,
3657            embd_dev,
3658            ckpt.as_deref_mut(),
3659            stream,
3660            fence,
3661            pp_pipe,
3662            None,
3663        )?;
3664        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
3665    }
3666
3667    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
3668    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
3669    #[allow(clippy::too_many_arguments)]
3670    fn verify_stage0_issue(
3671        &self,
3672        e: &Engine,
3673        tokens: &[u32],
3674        pos0: usize,
3675        cache: &mut Cache,
3676        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3677        mut ckpt: Option<&mut VerifyCkpt>,
3678        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3679        fence: &[usize],
3680        pp_pipe: Option<bool>,
3681        trace: Option<SpecPipeTraceCtx>,
3682    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
3683        assert!(
3684            !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
3685            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
3686             (the gemma4 arms have their own decode_step_t twins)"
3687        );
3688        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
3689            return Err(
3690                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
3691                 boundary itself is host-staged, but device-resident verify still peer-reads \
3692                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
3693                 serving on this host class; spec requires local per-stage inputs first."
3694                    .into(),
3695            );
3696        }
3697        let rt = crate::pp::PpNRt::get(e)?;
3698        let n_st = fence.len() - 1;
3699        assert_eq!(
3700            rt.n_stages(),
3701            n_st,
3702            "PpNRt stage count {} != fence stages {n_st}",
3703            rt.n_stages()
3704        );
3705        let n_embd = self.cfg.n_embd as usize;
3706        let t = tokens.len();
3707        let payload = t * n_embd;
3708        if pp_pipe.is_some() {
3709            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
3710        }
3711        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
3712        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
3713        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
3714        // the report below names exactly two stages and must never imply it measured middle ones.
3715        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
3716        let pp_started = std::time::Instant::now();
3717        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
3718        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
3719        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
3720        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
3721        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
3722        // stage stream and the wait would self-order into a no-op.
3723        let caller_stream = e.stream();
3724        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
3725        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
3726        // the primary stream still holds queued reads of them — with event tracking elided,
3727        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
3728        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
3729        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
3730        // stage stream behind the caller before enqueueing new stage work.
3731        let reverse_started = std::time::Instant::now();
3732        if pp_pipe != Some(false) {
3733            rt.fence_stages_behind(&caller_stream)?;
3734        }
3735        if pp_pipe == Some(true) {
3736            // Both session verifies must alternate boundary slots even when the ordinary
3737            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
3738            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
3739            rt.prepare_overlap_slots(0, payload)?;
3740        }
3741        if pp_anatomy {
3742            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
3743            // prices any primary-stream rollback/refresh tail inherited from the prior round.
3744            for s in 0..n_st {
3745                let _st = rt.enter(s);
3746                rt.engine(s, e).stream().synchronize()?;
3747            }
3748            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
3749        }
3750
3751        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
3752        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
3753        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3754            match stream {
3755                Some((_, ctr)) => {
3756                    let mut p = es.alloc_uninit::<i32>(t)?;
3757                    es.pos_iota(ctr, &mut p, t)?;
3758                    Ok(p)
3759                }
3760                None => {
3761                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3762                    es.htod_i32(&pos_vec)
3763                }
3764            }
3765        };
3766
3767        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
3768        let slot = {
3769            let _st0 = rt.enter(0);
3770            let e0 = rt.engine(0, e);
3771            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
3772            let stage0_started = std::time::Instant::now();
3773            let pos_d = stage_pos(e0)?;
3774            let x = match (stream, embd_dev) {
3775                (Some((vtok, _)), Some((g, qt, rb))) => {
3776                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3777                }
3778                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3779                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
3780            };
3781            let x = self.verify_layers(
3782                e0,
3783                x,
3784                fence[0],
3785                fence[1],
3786                &pos_d,
3787                pos0,
3788                t,
3789                cache,
3790                ckpt.as_deref_mut(),
3791                stream,
3792            )?;
3793            if pp_anatomy {
3794                e0.stream().synchronize()?;
3795                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
3796            }
3797            let tx_started = std::time::Instant::now();
3798            let slot = if pp_pipe.is_some() {
3799                rt.tx_pipelined(0, &x, payload)?
3800            } else {
3801                rt.tx(0, &x, payload)?
3802            };
3803            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
3804            if pp_anatomy {
3805                e0.stream().synchronize()?;
3806                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
3807            }
3808            slot
3809            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
3810        };
3811
3812        Ok(VerifyBoundaryTicket {
3813            rt,
3814            caller_stream,
3815            slot,
3816            pos0,
3817            t,
3818            payload,
3819            n_st,
3820            pipelined: pp_pipe.is_some(),
3821            pp_anatomy,
3822            pp_started,
3823            reverse_ms,
3824            stage0_ms,
3825            tx_ms,
3826            trace,
3827        })
3828    }
3829
3830    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
3831    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
3832    #[allow(clippy::too_many_arguments)]
3833    fn verify_stage1_finish(
3834        &self,
3835        e: &Engine,
3836        ticket: VerifyBoundaryTicket,
3837        cache: &mut Cache,
3838        mut ckpt: Option<&mut VerifyCkpt>,
3839        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3840        fence: &[usize],
3841        publish_to_caller: bool,
3842    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3843        let VerifyBoundaryTicket {
3844            rt,
3845            caller_stream,
3846            slot,
3847            pos0,
3848            t,
3849            payload,
3850            n_st,
3851            pipelined,
3852            pp_anatomy,
3853            pp_started,
3854            reverse_ms,
3855            stage0_ms,
3856            tx_ms,
3857            trace,
3858        } = ticket;
3859        let n_embd = self.cfg.n_embd as usize;
3860        let eps = self.cfg.rms_eps;
3861        let mut slot = slot;
3862        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
3863        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
3864            match stream {
3865                Some((_, ctr)) => {
3866                    let mut p = es.alloc_uninit::<i32>(t)?;
3867                    es.pos_iota(ctr, &mut p, t)?;
3868                    Ok(p)
3869                }
3870                None => {
3871                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3872                    es.htod_i32(&pos_vec)
3873                }
3874            }
3875        };
3876
3877        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
3878        for s in 1..n_st - 1 {
3879            let _st = rt.enter(s);
3880            let es = rt.engine(s, e);
3881            let pos_d = stage_pos(es)?;
3882            let x = rt.rx(s - 1, slot, payload)?;
3883            let x = self.verify_layers(
3884                es,
3885                x,
3886                fence[s],
3887                fence[s + 1],
3888                &pos_d,
3889                pos0,
3890                t,
3891                cache,
3892                ckpt.as_deref_mut(),
3893                stream,
3894            )?;
3895            slot = if pipelined {
3896                rt.tx_pipelined(s, &x, payload)?
3897            } else {
3898                rt.tx(s, &x, payload)?
3899            };
3900        }
3901
3902        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
3903        let _stl = rt.enter(n_st - 1);
3904        let el = rt.engine(n_st - 1, e);
3905        let pos_d = stage_pos(el)?;
3906        let rx_started = std::time::Instant::now();
3907        let x = rt.rx(n_st - 2, slot, payload)?;
3908        if pp_anatomy {
3909            el.stream().synchronize()?;
3910            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
3911        }
3912        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
3913        let stage1_started = std::time::Instant::now();
3914        let x = self.verify_layers(
3915            el,
3916            x,
3917            fence[n_st - 1],
3918            fence[n_st],
3919            &pos_d,
3920            pos0,
3921            t,
3922            cache,
3923            ckpt.as_deref_mut(),
3924            stream,
3925        )?;
3926
3927        let mut hn = vbuf(el, payload)?;
3928        let logits = if self.cfg.step35.is_some() {
3929            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
3930            // Verify must not switch numeric class merely because the same session speculates.
3931            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3932            el.matmul(&self.output, &hn, t)?
3933        } else {
3934            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3935            el.matmul_decode_exact(&self.output, &hn, t)?
3936        };
3937        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
3938        if pp_anatomy {
3939            el.stream().synchronize()?;
3940            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
3941        }
3942        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
3943        // stream. Order the caller's stream behind that work before the buffers escape this
3944        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
3945        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
3946        // the following arm's KV in the same process).
3947        if publish_to_caller {
3948            rt.publish_to(n_st - 1, &caller_stream)?;
3949        }
3950        if pp_anatomy {
3951            if publish_to_caller {
3952                caller_stream.synchronize()?;
3953            }
3954            eprintln!(
3955                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
3956                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
3957                pp_started.elapsed().as_secs_f64() * 1e3,
3958            );
3959        }
3960        // stream: the device pos counter owns position; host mirror reconciles at drain.
3961        if stream.is_none() {
3962            cache.pos += t;
3963        }
3964        Ok((logits, if spec_hpost() { hn } else { x }))
3965    }
3966
3967    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
3968    ///
3969    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
3970    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
3971    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
3972    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
3973    /// bytes when a request moves from batched plain serving into speculative verify. Run the
3974    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
3975    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
3976    /// every norm/projection/FFN uses exactly the live serving dispatch.
3977    #[allow(clippy::too_many_arguments)]
3978    fn step35_verify_batch_layers(
3979        &self,
3980        e: &Engine,
3981        mut x: CudaSlice<f32>,
3982        lo: usize,
3983        hi: usize,
3984        pos0: usize,
3985        t: usize,
3986        cache: &mut Cache,
3987    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3988        let n_embd = self.cfg.n_embd as usize;
3989        self.cfg
3990            .step35
3991            .as_ref()
3992            .ok_or("step35 verify batch requires step35 cfg")?;
3993        let mut ph_last = std::time::Instant::now();
3994        for il in lo..hi {
3995            let mut next = e.uninit(t * n_embd)?;
3996            for r in 0..t {
3997                let mut row = e.uninit(n_embd)?;
3998                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
3999                // The caller owns this verify's position. During controller overlap, cache.pos
4000                // still describes generation N while this stage-0 walk belongs to N+1.
4001                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4002                let mut one = [&mut *cache];
4003                let out = self.step35_decode_batch_layers(
4004                    e,
4005                    row,
4006                    &mut one,
4007                    &row_pos,
4008                    il,
4009                    il + 1,
4010                    &mut ph_last,
4011                )?;
4012                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4013            }
4014            self.dflash_tap(e, cache, il, &next, t)?;
4015            x = next;
4016        }
4017        Ok(x)
4018    }
4019
4020    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
4021    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
4022    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
4023    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
4024    /// prefix-keep, not all-or-nothing).
4025    pub(crate) fn dspark_verify_t_am(
4026        &self,
4027        e: &Engine,
4028        tokens: &[u32],
4029        pos0: usize,
4030        cache: &mut Cache,
4031    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4032        let (logits, _hn) =
4033            self.decode_step_t_core_stream(e, tokens, pos0, cache, None, None, None, None)?;
4034        let t = tokens.len();
4035        let v = self.output.out_features();
4036        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4037        for r in 0..t {
4038            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4039        }
4040        Ok(e.dtoh_u32(&am_d)?)
4041    }
4042
4043    /// DSpark verify with the MTP column-stash armed: identical forward to
4044    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
4045    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
4046    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
4047    pub(crate) fn dspark_verify_t_am_ckpt(
4048        &self,
4049        e: &Engine,
4050        tokens: &[u32],
4051        pos0: usize,
4052        cache: &mut Cache,
4053    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4054        let mut ck = VerifyCkpt::new(self.layers.len());
4055        let (logits, _hn) = self.decode_step_t_core_stream(
4056            e,
4057            tokens,
4058            pos0,
4059            cache,
4060            None,
4061            Some(&mut ck),
4062            None,
4063            None,
4064        )?;
4065        let t = tokens.len();
4066        let v = self.output.out_features();
4067        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4068        for r in 0..t {
4069            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4070        }
4071        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
4072    }
4073
4074    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
4075    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
4076    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
4077    pub(crate) fn dspark_commit_prefix(
4078        &self,
4079        e: &Engine,
4080        cache: &mut Cache,
4081        snap: &crate::cache::CacheSnapshot,
4082        ckpt: &DsparkVerifyCkpt,
4083        keep: usize,
4084    ) -> Result<(), Box<dyn std::error::Error>> {
4085        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
4086    }
4087
4088    /// Qwen35-family verify trunk in the live serving numeric class.
4089    ///
4090    /// Serving intentionally keeps this architecture in the generic batched program even at
4091    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
4092    ///
4093    /// Two arms, one numeric class:
4094    /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
4095    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
4096    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
4097    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
4098    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
4099    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
4100    ///   program its isolated serving step would). One weight read per layer per round
4101    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
4102    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
4103    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
4104    ///   serving layer body, preserving single-session autoregressive cache order (the
4105    ///   correctness reference; also the rollback seam for the t-parallel arm).
4106    ///
4107    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
4108    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
4109    #[allow(clippy::too_many_arguments)]
4110    fn qwen35_verify_batch_layers(
4111        &self,
4112        e: &Engine,
4113        x: CudaSlice<f32>,
4114        lo: usize,
4115        hi: usize,
4116        pos0: usize,
4117        t: usize,
4118        cache: &mut Cache,
4119        ckpt: Option<&mut VerifyCkpt>,
4120        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4121    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4122        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
4123        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
4124        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
4125        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
4126        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
4127        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
4128        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
4129            || !matches!(
4130                self.cfg.arch,
4131                memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
4132            )
4133            || t > 16;
4134        if rowwise {
4135            if stream.is_some() {
4136                // rowwise replays per row with host cache.pos — irreconcilable with a
4137                // device position counter. Burst callers must keep t <= 16 and the
4138                // ROWWISE env unset; refusing beats silently mispositioned rows.
4139                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
4140                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
4141                    .into());
4142            }
4143            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
4144        } else {
4145            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream)
4146        }
4147    }
4148
4149    /// The per-row correctness reference: replay each verify row through the authoritative
4150    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
4151    #[allow(clippy::too_many_arguments)]
4152    fn qwen35_verify_rowwise(
4153        &self,
4154        e: &Engine,
4155        mut x: CudaSlice<f32>,
4156        lo: usize,
4157        hi: usize,
4158        pos0: usize,
4159        t: usize,
4160        cache: &mut Cache,
4161        mut ckpt: Option<&mut VerifyCkpt>,
4162    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4163        let n_embd = self.cfg.n_embd as usize;
4164        let saved_pos = cache.pos;
4165        let mut ph_last = std::time::Instant::now();
4166        for il in lo..hi {
4167            let mut next = e.uninit(t * n_embd)?;
4168            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4169                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4170                    Some(Vec::with_capacity(t - 1))
4171                } else {
4172                    None
4173                };
4174            for r in 0..t {
4175                cache.pos = pos0 + r;
4176                let mut row = e.uninit(n_embd)?;
4177                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4178                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4179                let mut one = [&mut *cache];
4180                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
4181                let out = match self.decode_batch_layers(
4182                    e,
4183                    row,
4184                    &mut one,
4185                    &ctx,
4186                    &row_pos,
4187                    &mut ph_last,
4188                ) {
4189                    Ok(out) => out,
4190                    Err(error) => {
4191                        cache.pos = saved_pos;
4192                        return Err(error);
4193                    }
4194                };
4195                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4196                if r + 1 < t {
4197                    if let Some(states) = col_states.as_mut() {
4198                        let recur = cache.recur[il]
4199                            .as_ref()
4200                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
4201                        states.push((
4202                            e.clone_dtod(&recur.conv_state)?,
4203                            e.clone_dtod(&recur.ssm_state)?,
4204                        ));
4205                    }
4206                }
4207            }
4208            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4209                checkpoint.cols[il] = Some(states);
4210            }
4211            x = next;
4212        }
4213        cache.pos = saved_pos;
4214        Ok(x)
4215    }
4216
4217    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
4218    ///
4219    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
4220    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
4221    /// pins the serving batch tier already carries:
4222    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
4223    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
4224    ///     alone;
4225    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
4226    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
4227    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
4228    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
4229    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
4230    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
4231    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
4232    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
4233    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
4234    /// program its isolated B=1 serving step would.
4235    ///
4236    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
4237    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
4238    #[allow(clippy::too_many_arguments)]
4239    fn qwen35_verify_tparallel(
4240        &self,
4241        e: &Engine,
4242        mut x: CudaSlice<f32>,
4243        lo: usize,
4244        hi: usize,
4245        pos0: usize,
4246        t: usize,
4247        cache: &mut Cache,
4248        mut ckpt: Option<&mut VerifyCkpt>,
4249        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4250    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4251        use cudarc::driver::DevicePtr;
4252        let cfg = &self.cfg;
4253        let n_embd = cfg.n_embd as usize;
4254        let eps = cfg.rms_eps;
4255        let head_dim_global = cfg.head_dim_k as usize;
4256        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
4257        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
4258        let pos_d = match stream {
4259            Some((_, ctr)) => {
4260                let mut p = e.alloc_uninit::<i32>(t)?;
4261                e.pos_iota(ctr, &mut p, t)?;
4262                p
4263            }
4264            None => {
4265                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
4266                e.htod_i32(&pos_host)?
4267            }
4268        };
4269        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
4270        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
4271        let pos_rows: Vec<CudaSlice<i32>> = match stream {
4272            Some((_, ctr)) => (0..t)
4273                .map(|r| {
4274                    let mut b = e.alloc_uninit::<i32>(1)?;
4275                    e.i32_copy_add(ctr, &mut b, r as i32)?;
4276                    Ok(b)
4277                })
4278                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
4279            None => (0..t)
4280                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
4281                .collect::<Result<_, _>>()?,
4282        };
4283        let seqs_append =
4284            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
4285        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
4286
4287        for il in lo..hi {
4288            let layer = &self.layers[il];
4289            // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
4290            let anorm = layer.attn_norm.float_data();
4291            let mut xn = e.uninit(t * n_embd)?;
4292            e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
4293            let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
4294
4295            let mixed: CudaSlice<f32> = match &layer.mixer {
4296                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4297                Mixer::Full(fa) => {
4298                    let geometry = cfg.full_attention_geometry_at(il as u32);
4299                    let n_head = geometry.n_head as usize;
4300                    let n_head_kv = geometry.n_head_kv as usize;
4301                    let head_dim = geometry.head_dim_k as usize;
4302                    let rope_dims = geometry.n_rot as usize;
4303                    let rope_base = geometry.rope_base;
4304                    let scale = geometry.attention_scale();
4305                    // Batched projections: one weight read serves all T rows.
4306                    let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
4307                    let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
4308                    let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
4309                    let gated =
4310                        geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4311                    let (mut q, gate) = if gated {
4312                        let mut qs = e.uninit(t * n_head * head_dim)?;
4313                        let mut gs = e.uninit(t * n_head * head_dim)?;
4314                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
4315                        (qs, Some(gs))
4316                    } else {
4317                        (qf, None)
4318                    };
4319                    let mut qn = e.uninit(t * n_head * head_dim)?;
4320                    e.rms_norm(
4321                        &q,
4322                        fa.q_norm.float_data(),
4323                        &mut qn,
4324                        head_dim,
4325                        t * n_head,
4326                        eps,
4327                    )?;
4328                    q = qn;
4329                    let mut kn = e.uninit(t * n_head_kv * head_dim)?;
4330                    e.rms_norm(
4331                        &k,
4332                        fa.k_norm.float_data(),
4333                        &mut kn,
4334                        head_dim,
4335                        t * n_head_kv,
4336                        eps,
4337                    )?;
4338                    k = kn;
4339                    e.rope_neox(
4340                        &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
4341                    )?;
4342                    e.rope_neox(
4343                        &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4344                    )?;
4345
4346                    // Per-row append + attend: row r sees rows 0..r in KV (causal within the
4347                    // draft), each through the b_n=1 serving kernels at its own t_kv.
4348                    let q_dim = n_head * head_dim;
4349                    let kv_dim = n_head_kv * head_dim;
4350                    let mut attn = e.uninit(t * q_dim)?;
4351                    let (kdk, kdv, ktb, vtb, kv_view) = {
4352                        let kvl = cache.kv[il].as_ref().unwrap();
4353                        let s = &e.gpu.stream();
4354                        let (pk, _g) = kvl.k.device_ptr(s);
4355                        let (pv, _g2) = kvl.v.device_ptr(s);
4356                        (
4357                            kvl.kv_dim_k,
4358                            kvl.kv_dim_v,
4359                            kvl.k_tok_bytes,
4360                            kvl.v_tok_bytes,
4361                            e.htod_u64(&[pk as u64, pv as u64])?,
4362                        )
4363                    };
4364                    if let Some((_, ctr)) = stream {
4365                        // STREAM ARM (2b): one batched dc append + the multi-row dc attention
4366                        // — the generic stream arm's exact shape (rows kernels are pinned
4367                        // byte-identical to the per-row programs by kernel-check). Host len
4368                        // stays a stale lower bound; the burst drain reconciles it.
4369                        let kvl = cache.kv[il].as_mut().unwrap();
4370                        e.append_kv_quantized_rows_dc(
4371                            &k,
4372                            &v,
4373                            &mut kvl.k,
4374                            &mut kvl.v,
4375                            ctr,
4376                            t,
4377                            kdk,
4378                            kdv,
4379                            ktb,
4380                            vtb,
4381                            Engine::kv_fp8_on(),
4382                        )?;
4383                        let upper = (kvl.len + t + 64).min(cache.max_ctx);
4384                        let k_view = e.view_u8(&kvl.k, upper * ktb);
4385                        let v_view = e.view_u8(&kvl.v, upper * vtb);
4386                        e.fa_decode_rows_dc(
4387                            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr,
4388                            upper, t, scale, ktb, vtb, 0, false,
4389                        )?;
4390                    } else {
4391                        for r in 0..t {
4392                            // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
4393                            // whose row 0 is this row (arithmetic-free materialization copies,
4394                            // same as decode's per-seq fallback arm).
4395                            let mut k_row = e.uninit(kv_dim)?;
4396                            e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
4397                            let mut v_row = e.uninit(kv_dim)?;
4398                            e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
4399                            let pos_row = &pos_rows[r];
4400                            let kvl = cache.kv[il].as_mut().unwrap();
4401                            if seqs_append {
4402                                e.append_kv_quantized_seqs(
4403                                    &k_row,
4404                                    &v_row,
4405                                    &kv_view.slice(0..2),
4406                                    pos_row,
4407                                    1,
4408                                    kdk,
4409                                    kdv,
4410                                    ktb,
4411                                    vtb,
4412                                )?;
4413                                kvl.len += 1;
4414                            } else {
4415                                e.append_kv_quantized_view(
4416                                    &k_row.slice(0..kv_dim),
4417                                    &v_row.slice(0..kv_dim),
4418                                    &mut kvl.k,
4419                                    &mut kvl.v,
4420                                    kvl.len,
4421                                    kvl.kv_dim_k,
4422                                    kvl.kv_dim_v,
4423                                    kvl.k_tok_bytes,
4424                                    kvl.v_tok_bytes,
4425                                    Engine::kv_fp8_on(),
4426                                )?;
4427                                kvl.len += 1;
4428                            }
4429                            let t_kv = kvl.len;
4430                            let mut q_row = e.uninit(q_dim)?;
4431                            e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
4432                            let mut a_row = e.uninit(q_dim)?;
4433                            if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
4434                                let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
4435                                e.fa_decode_batch_seqs_v4(
4436                                    &q_row,
4437                                    &kv_view.slice(0..2),
4438                                    pos_row,
4439                                    &mut a_row,
4440                                    head_dim,
4441                                    n_head,
4442                                    n_head_kv,
4443                                    1,
4444                                    t_kv,
4445                                    scale,
4446                                    sp0_r,
4447                                    ktb,
4448                                    vtb,
4449                                )?;
4450                            } else {
4451                                let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
4452                                let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
4453                                let mut a_view = a_row.slice_mut(0..q_dim);
4454                                e.fa_decode_kvmod_view(
4455                                    &q_row.slice(0..q_dim),
4456                                    &k_view,
4457                                    &v_view,
4458                                    &mut a_view,
4459                                    head_dim,
4460                                    n_head,
4461                                    n_head_kv,
4462                                    t_kv,
4463                                    scale,
4464                                    kvl.k_tok_bytes,
4465                                    kvl.v_tok_bytes,
4466                                    Engine::kv_fp8_on(),
4467                                )?;
4468                            }
4469                            e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
4470                        }
4471                    }
4472
4473                    // Output gate (element-wise) + o-proj at m=T.
4474                    let attn_g = match &gate {
4475                        Some(g) => {
4476                            let n = t * q_dim;
4477                            let mut gsig = e.uninit(n)?;
4478                            e.sigmoid(g, &mut gsig, n)?;
4479                            let mut ag = e.uninit(n)?;
4480                            e.mul(&attn, &gsig, &mut ag, n)?;
4481                            ag
4482                        }
4483                        None => attn,
4484                    };
4485                    e.matmul(&fa.wo, &attn_g, t)?
4486                }
4487                // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
4488                // per-row serving-kernel chain cannot run (host state swaps keyed on host
4489                // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
4490                // rebuild — the per-row chain only produces per-column clones). GDN rides
4491                // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
4492                // and its one-scan recurrence is pinned bit-identical to T chained T=1
4493                // steps (its header + kernel-check). Position-independent, so no counter
4494                // plumbing is needed. Guards mirror the generic call site exactly.
4495                Mixer::Linear(la) if stream.is_some() => {
4496                    if !(t >= 3 || (t == 2 && spec_m2()))
4497                        || !self.mixer_in_q8_1_fast(e, &layer.mixer)
4498                        || !e.uses_q8_1_fast(&la.ssm_out)
4499                    {
4500                        return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
4501                                    (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
4502                            .into());
4503                    }
4504                    let want = ckpt.is_some();
4505                    let (out, stash) = self.linear_attn_verify_t(
4506                        e,
4507                        la,
4508                        &xn,
4509                        Some((&hq, &hd)),
4510                        t,
4511                        cache,
4512                        il,
4513                        want,
4514                    )?;
4515                    if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4516                        ck.gdn[il] = Some(st);
4517                    }
4518                    out
4519                }
4520                Mixer::Linear(la) => {
4521                    let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
4522                    let d_state = ssm.state_size as usize;
4523                    let num_k = ssm.group_count as usize;
4524                    let num_v = ssm.time_step_rank as usize;
4525                    let d_conv = ssm.conv_kernel as usize;
4526                    let key_dim = d_state * num_k;
4527                    let value_dim = d_state * num_v;
4528                    let conv_dim = key_dim * 2 + value_dim;
4529                    let gdn_scale = 1.0 / (d_state as f32).sqrt();
4530
4531                    // ---- batched projections: one weight read for all T rows ----
4532                    let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
4533                    let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
4534                    let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
4535                    let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
4536                    let beta_w = la.ssm_beta.out_features();
4537                    let alpha_w = la.ssm_alpha.out_features();
4538                    let qkv_w = la.wqkv.out_features();
4539
4540                    // ---- per-row state chain through the b_n=1 serving kernels ----
4541                    // 6-entry alternating pointer table expresses the ping-pong without a
4542                    // rebuild per row: even rows scan s0 -> s1, odd rows s1 -> s0. Host
4543                    // handles swap per row so ckpt clones the canonical state (and the
4544                    // post-verify canonical handle matches the last write), exactly as the
4545                    // rowwise arm leaves them.
4546                    let table = {
4547                        let rl = cache.recur[il].as_ref().unwrap();
4548                        let s = &e.gpu.stream();
4549                        let (pc, _g0) = rl.conv_state.device_ptr(s);
4550                        let (p0, _g1) = rl.ssm_state.device_ptr(s);
4551                        let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
4552                        e.htod_u64(&[
4553                            pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
4554                        ])?
4555                    };
4556                    let mut o_all = e.uninit(t * value_dim)?;
4557                    let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4558                        if ckpt.is_some() && t >= 2 {
4559                            Some(Vec::with_capacity(t - 1))
4560                        } else {
4561                            None
4562                        };
4563                    // Per-row scratch reused across rows (uninit is cheap but not free at
4564                    // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
4565                    // [T, ...] buffers — zero arithmetic-free copies in this loop.
4566                    let mut conv_out = e.uninit(conv_dim)?;
4567                    let mut q_l2 = e.uninit(value_dim)?;
4568                    let mut k_l2 = e.uninit(value_dim)?;
4569                    let mut v_gd = e.uninit(value_dim)?;
4570                    let mut beta_b = e.uninit(num_v)?;
4571                    let mut g_log = e.uninit(num_v)?;
4572                    for r in 0..t {
4573                        let base = if r % 2 == 0 { 0 } else { 3 };
4574                        let conv_view = table.slice(base..base + 1);
4575                        let in_view = table.slice(base + 1..base + 2);
4576                        let out_view = table.slice(base + 2..base + 3);
4577                        e.ssm_conv1d_fused_decode_b_view(
4578                            &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
4579                            &conv_view,
4580                            la.ssm_conv1d.float_data(),
4581                            &mut conv_out,
4582                            conv_dim,
4583                            d_conv,
4584                            1,
4585                        )?;
4586                        e.gdn_prep_decode_b_view(
4587                            &conv_out,
4588                            &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
4589                            &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
4590                            la.ssm_dt.float_data(),
4591                            la.ssm_a.float_data(),
4592                            &mut q_l2,
4593                            &mut k_l2,
4594                            &mut v_gd,
4595                            &mut beta_b,
4596                            &mut g_log,
4597                            d_state,
4598                            num_v,
4599                            num_k,
4600                            key_dim,
4601                            eps,
4602                            conv_dim,
4603                            1,
4604                        )?;
4605                        let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
4606                        e.gdn_scan_s128_batched_view(
4607                            &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row,
4608                            num_v, 1, gdn_scale,
4609                        )?;
4610                        {
4611                            let rl = cache.recur[il].as_mut().unwrap();
4612                            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
4613                        }
4614                        if r + 1 < t {
4615                            if let Some(states) = col_states.as_mut() {
4616                                let recur = cache.recur[il]
4617                                    .as_ref()
4618                                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
4619                                states.push((
4620                                    e.clone_dtod(&recur.conv_state)?,
4621                                    e.clone_dtod(&recur.ssm_state)?,
4622                                ));
4623                            }
4624                        }
4625                    }
4626                    if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4627                        checkpoint.cols[il] = Some(states);
4628                    }
4629
4630                    // ---- batched gated norm + out-projection at m=T ----
4631                    if e.uses_q8_1_fast(&la.ssm_out) {
4632                        let (gq, gd) = e.gated_rmsnorm_q8_1(
4633                            &o_all,
4634                            la.ssm_norm.float_data(),
4635                            &z,
4636                            d_state,
4637                            t * num_v,
4638                            eps,
4639                        )?;
4640                        let g0 = e.zeros(0)?;
4641                        e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
4642                    } else {
4643                        let mut gn = e.uninit(t * value_dim)?;
4644                        e.gated_rmsnorm(
4645                            &o_all,
4646                            la.ssm_norm.float_data(),
4647                            &z,
4648                            &mut gn,
4649                            d_state,
4650                            t * num_v,
4651                            eps,
4652                        )?;
4653                        e.matmul(&la.ssm_out, &gn, t)?
4654                    }
4655                }
4656            };
4657
4658            // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
4659            let pnorm = layer.post_attn_norm.float_data();
4660            let mut x1 = e.uninit(t * n_embd)?;
4661            let mut zn = e.uninit(t * n_embd)?;
4662            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
4663            let ffn_out = match &layer.ffn {
4664                crate::hybrid::Ffn::Dense {
4665                    ffn_gate,
4666                    ffn_up,
4667                    ffn_down,
4668                } => {
4669                    assert!(
4670                        self.cfg.m3.is_none(),
4671                        "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
4672                    );
4673                    let n_ff = ffn_gate.out_features();
4674                    let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
4675                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
4676                    let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
4677                    let mut act = e.uninit(t * n_ff)?;
4678                    e.silu_mul(&g, &u, &mut act, t * n_ff)?;
4679                    let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
4680                    e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
4681                }
4682                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
4683            };
4684            let mut x2 = e.uninit(t * n_embd)?;
4685            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
4686            // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
4687            self.dflash_tap(e, cache, il, &x2, t)?;
4688            x = x2;
4689        }
4690        Ok(x)
4691    }
4692
4693    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
4694    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
4695    /// carried in from outside the range) and exits with the range's final residual materialized
4696    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
4697    /// instead of one.
4698    ///
4699    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
4700    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
4701    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
4702    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
4703    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
4704    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
4705    /// code — there is no "split version" of the verify math.
4706    ///
4707    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
4708    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
4709    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
4710    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
4711    #[allow(clippy::too_many_arguments)]
4712    fn verify_layers(
4713        &self,
4714        e: &Engine,
4715        mut x: CudaSlice<f32>,
4716        lo: usize,
4717        hi: usize,
4718        pos_d: &CudaSlice<i32>,
4719        pos0: usize,
4720        t: usize,
4721        cache: &mut Cache,
4722        mut ckpt: Option<&mut VerifyCkpt>,
4723        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4724    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4725        if self.cfg.step35.is_some() {
4726            if stream.is_some() {
4727                return Err(
4728                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
4729                            cannot express the SWA offset KV view)"
4730                        .into(),
4731                );
4732            }
4733            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
4734        }
4735        if self.qwen35_serving_class() {
4736            return self.qwen35_verify_batch_layers(
4737                e,
4738                x,
4739                lo,
4740                hi,
4741                pos0,
4742                t,
4743                cache,
4744                ckpt.take(),
4745                stream,
4746            );
4747        }
4748        let n_embd = self.cfg.n_embd as usize;
4749        let eps = self.cfg.rms_eps;
4750        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
4751        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
4752        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
4753        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
4754        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
4755        // residual the next layer needs) as its `res` output. Falls back to the separate add
4756        // when the next layer is off the fused-q8 path.
4757        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
4758        for il in lo..hi {
4759            let layer = &self.layers[il];
4760            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
4761            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
4762            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
4763            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
4764            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
4765            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
4766            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
4767            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
4768            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
4769            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
4770            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
4771            // projections only; Linear mixer: the batched arm — the per-column fallback needs
4772            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
4773            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
4774            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
4775            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
4776            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
4777            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
4778            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
4779            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
4780            let lin_q8_only = match &layer.mixer {
4781                Mixer::Linear(la) => {
4782                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
4783                }
4784                Mixer::Full(_) if self.cfg.step35.is_some() => false,
4785                _ => true,
4786            };
4787            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
4788            // a non-fused layer still performs the residual add.
4789            let taken = pending.take();
4790            let (h, h_q8) = if norm_fused && lin_q8_only {
4791                let pair = match taken {
4792                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
4793                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
4794                    Some((x1p, f1p)) => {
4795                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
4796                        let p = e.add_rms_norm_q8_1(
4797                            &x1p,
4798                            &f1p,
4799                            layer.attn_norm.float_data(),
4800                            &mut x2,
4801                            n_embd,
4802                            t,
4803                            eps,
4804                        )?;
4805                        x = x2;
4806                        p
4807                    }
4808                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
4809                };
4810                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
4811            } else {
4812                if let Some((x1p, f1p)) = taken {
4813                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
4814                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
4815                    x = x2;
4816                }
4817                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
4818                if norm_fused {
4819                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4820                } else {
4821                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
4822                }
4823                (h, None)
4824            };
4825            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
4826
4827            let mixed = match &layer.mixer {
4828                Mixer::Full(fa) => self.full_attn_verify(
4829                    e,
4830                    fa,
4831                    &h,
4832                    h_q8_ref,
4833                    pos_d,
4834                    t,
4835                    cache,
4836                    il,
4837                    stream.map(|(_, c)| c),
4838                )?,
4839                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4840                Mixer::Linear(la) => {
4841                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
4842                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
4843                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
4844                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
4845                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
4846                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
4847                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
4848                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
4849                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
4850                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
4851                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
4852                    if (t >= 3 || (t == 2 && spec_m2()))
4853                        && mixer_fast
4854                        && e.uses_q8_1_fast(&la.ssm_out)
4855                    {
4856                        let want = ckpt.is_some();
4857                        let (out, stash) =
4858                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
4859                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
4860                            ck.gdn[il] = Some(st);
4861                        }
4862                        out
4863                    } else {
4864                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
4865                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4866                            if ckpt.is_some() && t >= 2 {
4867                                Some(Vec::with_capacity(t - 1))
4868                            } else {
4869                                None
4870                            };
4871                        for col in 0..t {
4872                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
4873                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
4874                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
4875                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
4876                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
4877                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
4878                            // (pure dtod — cannot change any computed value). Last column skipped:
4879                            // rebuild targets are j <= t-1 columns.
4880                            if let Some(cs) = col_states.as_mut() {
4881                                if col + 1 < t {
4882                                    let rl = cache.recur[il].as_ref().unwrap();
4883                                    cs.push((
4884                                        e.clone_dtod(&rl.conv_state)?,
4885                                        e.clone_dtod(&rl.ssm_state)?,
4886                                    ));
4887                                }
4888                            }
4889                        }
4890                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
4891                            // ReplaySSM-assessment instrumentation (2026-07-30): the
4892                            // per-column clones are the only true state snapshots left in
4893                            // the verify (the batched path stashes INPUTS and replays).
4894                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
4895                                static ONCE: std::sync::Once = std::sync::Once::new();
4896                                let bytes: usize =
4897                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
4898                                ONCE.call_once(|| eprintln!(
4899                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
4900                                    cs.len(), bytes as f64 / 1e6));
4901                            }
4902                            ck.cols[il] = Some(cs);
4903                        }
4904                        out
4905                    }
4906                }
4907            };
4908
4909            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
4910            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
4911            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
4912            let ffn_fuse = match &layer.ffn {
4913                crate::hybrid::Ffn::Dense {
4914                    ffn_gate, ffn_up, ..
4915                } => {
4916                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
4917                        && e.uses_q8_1_fast(ffn_gate)
4918                        && e.uses_q8_1_fast(ffn_up)
4919                }
4920                crate::hybrid::Ffn::Moe(_) => false,
4921            };
4922            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
4923            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
4924            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
4925            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
4926            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
4927            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
4928            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
4929            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
4930            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
4931            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
4932            // mirror decode's dispatch or spec self-consistency fails.
4933            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
4934            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
4935            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
4936            let mut z = e.zeros(0)?; // replaced below on the unfused arms
4937            let z_q8 = if fuse_q8 {
4938                Some(e.add_rms_norm_q8_1(
4939                    &x,
4940                    &mixed,
4941                    layer.post_attn_norm.float_data(),
4942                    &mut x1,
4943                    n_embd,
4944                    t,
4945                    eps,
4946                )?)
4947            } else {
4948                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
4949                if ffn_fuse {
4950                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
4951                    e.rms_norm_decode(
4952                        &x1,
4953                        layer.post_attn_norm.float_data(),
4954                        &mut zf,
4955                        n_embd,
4956                        t,
4957                        eps,
4958                    )?;
4959                } else {
4960                    e.add_rms_norm(
4961                        &x,
4962                        &mixed,
4963                        layer.post_attn_norm.float_data(),
4964                        &mut x1,
4965                        &mut zf,
4966                        n_embd,
4967                        t,
4968                        eps,
4969                    )?;
4970                }
4971                z = zf;
4972                None
4973            };
4974            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
4975            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
4976            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
4977            let ffn_out = match &layer.ffn {
4978                crate::hybrid::Ffn::Dense {
4979                    ffn_gate,
4980                    ffn_up,
4981                    ffn_down,
4982                } => {
4983                    let n_ff = ffn_gate.out_features();
4984                    if let Some((zq, zd)) = z_q8.as_ref() {
4985                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
4986                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
4987                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
4988                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
4989                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
4990                        // structure at nrows=t.
4991                        let pair =
4992                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
4993                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
4994                                None => None,
4995                            };
4996                        let (gate, gs, up, us) = match pair {
4997                            Some(x4) => x4,
4998                            None => (
4999                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
5000                                1.0, // scale already applied inside _pre
5001                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
5002                                1.0,
5003                            ),
5004                        };
5005                        if e.uses_q8_1_fast(ffn_down) {
5006                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
5007                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
5008                        } else {
5009                            let mut act = vbuf(e, t * n_ff)?;
5010                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
5011                            e.matmul_decode_exact(ffn_down, &act, t)?
5012                        }
5013                    } else {
5014                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
5015                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
5016                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
5017                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
5018                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
5019                        let (gate, up) =
5020                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
5021                                Some(pair) => pair,
5022                                None => (
5023                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
5024                                    e.matmul_decode_exact(ffn_up, &z, t)?,
5025                                ),
5026                            };
5027                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5028                        Self::ffn_act_lim(
5029                            e,
5030                            &self.cfg,
5031                            &gate,
5032                            &up,
5033                            1.0,
5034                            1.0,
5035                            dense_lim,
5036                            &mut act,
5037                            t * n_ff,
5038                        )?;
5039                        e.matmul_decode_exact(ffn_down, &act, t)?
5040                    }
5041                }
5042                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5043            };
5044            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
5045            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
5046            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
5047            pending = Some((x1, ffn_out));
5048        }
5049        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
5050        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
5051        if let Some((x1p, f1p)) = pending.take() {
5052            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5053            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5054            x = x2;
5055        }
5056        Ok(x)
5057    }
5058    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
5059    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
5060    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
5061    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
5062    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
5063    /// ssm state exactly like T sequential decode steps.
5064    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
5065    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
5066    #[allow(clippy::too_many_arguments)]
5067    fn linear_attn_verify_t(
5068        &self,
5069        e: &Engine,
5070        la: &LinearAttnLayer,
5071        h: &CudaSlice<f32>,
5072        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5073        t: usize,
5074        cache: &mut Cache,
5075        il: usize,
5076        want_stash: bool,
5077    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
5078        let cfg = &self.cfg;
5079        let ssm = cfg.ssm.as_ref().unwrap();
5080        let d_state = ssm.state_size as usize;
5081        let num_k = ssm.group_count as usize;
5082        let num_v = ssm.time_step_rank as usize;
5083        let d_conv = ssm.conv_kernel as usize;
5084        let key_dim = d_state * num_k;
5085        let conv_dim = key_dim * 2 + d_state * num_v;
5086        let eps = cfg.rms_eps;
5087        let scale = 1.0 / (d_state as f32).sqrt();
5088
5089        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
5090        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
5091        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
5092        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
5093        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
5094        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
5095        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
5096        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
5097        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
5098        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
5099        // Bit-identical per (tensor,token,row) — see spec_fused_t().
5100        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
5101        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
5102        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
5103        // and feeds every projection; the caller guaranteed all four input projections are
5104        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
5105        let h_q8_t = if h_q8.is_none()
5106            && spec_fused_t()
5107            && (2..=4).contains(&t)
5108            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
5109                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
5110        {
5111            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
5112        } else {
5113            None
5114        };
5115        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
5116        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
5117            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
5118        let (qkv_mixed, z) = {
5119            let mut fused = None;
5120            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
5121                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5122                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
5123            } else if let Some((hq, hd)) = hq8_any {
5124                if spec_fused_t() && (2..=4).contains(&t) {
5125                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
5126                }
5127            }
5128            match (fused, hq8_any) {
5129                (Some(pair), _) => pair,
5130                (None, Some((hq, hd))) if h_q8.is_some() => (
5131                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
5132                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
5133                ),
5134                (None, _) => (
5135                    e.matmul_decode_exact(&la.wqkv, h, t)?,
5136                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
5137                ),
5138            }
5139        };
5140        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
5141        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
5142        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
5143        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
5144        let (beta_raw, alpha) = if t == 1 {
5145            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5146            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
5147                Some(((mut b, bs), (mut a, as_))) => {
5148                    if bs != 1.0 {
5149                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
5150                    }
5151                    if as_ != 1.0 {
5152                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
5153                    }
5154                    (b, a)
5155                }
5156                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
5157                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
5158                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
5159                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
5160                    Some((b, a)) => (b, a),
5161                    None => (
5162                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
5163                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
5164                    ),
5165                },
5166            }
5167        } else {
5168            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
5169            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
5170            let mut nvfp4_fused = None;
5171            let mut q8_fused = None;
5172            if let Some((hq, hd)) = hq8_any {
5173                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
5174                    nvfp4_fused =
5175                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5176                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
5177                        static ONCE: std::sync::Once = std::sync::Once::new();
5178                        ONCE.call_once(|| {
5179                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
5180                        });
5181                    }
5182                }
5183                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
5184                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5185                }
5186            }
5187            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
5188                if bs != 1.0 {
5189                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
5190                }
5191                if as_ != 1.0 {
5192                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
5193                }
5194                (b, a)
5195            } else if let Some(pair) = q8_fused {
5196                pair
5197            } else {
5198                match hq8_any {
5199                    Some((hq, hd)) if h_q8.is_some() => (
5200                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
5201                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
5202                    ),
5203                    _ => (
5204                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
5205                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
5206                    ),
5207                }
5208            }
5209        };
5210
5211        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
5212        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
5213        let rl = cache.recur[il].as_mut().unwrap();
5214        let mut conv_out = e.uninit(conv_dim * t)?;
5215        e.ssm_conv1d_tm_state(
5216            &qkv_mixed,
5217            &mut rl.conv_state,
5218            la.ssm_conv1d.float_data(),
5219            &mut conv_out,
5220            conv_dim,
5221            t,
5222            d_conv,
5223        )?;
5224
5225        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
5226        let mut q_g = e.uninit(d_state * num_v * t)?;
5227        let mut k_g = e.uninit(d_state * num_v * t)?;
5228        let mut v_g = e.uninit(d_state * num_v * t)?;
5229        e.qkv_to_gdn_repack(
5230            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
5231        )?;
5232        let mut q_l2 = e.uninit(d_state * num_v * t)?;
5233        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
5234        let mut k_l2 = e.uninit(d_state * num_v * t)?;
5235        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
5236        let mut beta = e.uninit(t * num_v)?;
5237        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
5238        let mut g_log = e.uninit(t * num_v)?;
5239        e.gdn_glog(
5240            &alpha,
5241            la.ssm_dt.float_data(),
5242            la.ssm_a.float_data(),
5243            &mut g_log,
5244            num_v,
5245            t,
5246        )?;
5247
5248        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
5249        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
5250        let mut o = e.uninit(d_state * num_v * t)?;
5251        {
5252            let crate::cache::RecurLayer {
5253                ssm_state,
5254                ssm_state_alt,
5255                ..
5256            } = rl;
5257            e.gdn_scan_s128(
5258                &q_l2,
5259                &k_l2,
5260                &v_g,
5261                &g_log,
5262                &beta,
5263                ssm_state,
5264                ssm_state_alt,
5265                &mut o,
5266                num_v,
5267                t,
5268                scale,
5269            )?;
5270        }
5271        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5272
5273        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
5274        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
5275        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
5276        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
5277        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
5278        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
5279        let out = if e.uses_q8_1_fast(&la.ssm_out) {
5280            let (gq, gd) =
5281                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
5282            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
5283        } else {
5284            let mut gn = e.uninit(d_state * num_v * t)?;
5285            e.gated_rmsnorm(
5286                &o,
5287                la.ssm_norm.float_data(),
5288                &z,
5289                &mut gn,
5290                d_state,
5291                num_v * t,
5292                eps,
5293            )?;
5294            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
5295            // would fall to dp4a with a different FP reduction order — same class of bug as
5296            // the input projs).
5297            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
5298        };
5299        let stash = if want_stash {
5300            Some(GdnStash {
5301                qkv_mixed,
5302                q_l2,
5303                k_l2,
5304                v_g,
5305                g_log,
5306                beta,
5307            })
5308        } else {
5309            None
5310        };
5311        Ok((out, stash))
5312    }
5313
5314    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
5315    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
5316    /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
5317    ///   are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
5318    ///   verify-probe gates), so keeping them == replaying them.
5319    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
5320    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
5321    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
5322    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
5323    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
5324    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
5325    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
5326    fn commit_verified_prefix(
5327        &self,
5328        e: &Engine,
5329        cache: &mut Cache,
5330        snap: &crate::cache::CacheSnapshot,
5331        ckpt: &VerifyCkpt,
5332        j: usize,
5333        kv_lens_done: bool,
5334        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
5335    ) -> Result<(), Box<dyn std::error::Error>> {
5336        let cfg = &self.cfg;
5337        let ssm = cfg.ssm.as_ref().unwrap();
5338        let d_state = ssm.state_size as usize;
5339        let num_k = ssm.group_count as usize;
5340        let num_v = ssm.time_step_rank as usize;
5341        let d_conv = ssm.conv_kernel as usize;
5342        let conv_dim = d_state * num_k * 2 + d_state * num_v;
5343        let scale = 1.0 / (d_state as f32).sqrt();
5344        for il in 0..self.layers.len() {
5345            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
5346                kvl.len = saved + j;
5347                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
5348                if !kv_lens_done {
5349                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5350                }
5351            }
5352            if let Some(rl) = cache.recur[il].as_mut() {
5353                if let Some(st) = &ckpt.gdn[il] {
5354                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
5355                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
5356                    if let Some((acc, base, t_v)) = dev_j {
5357                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
5358                        e.ssm_conv_ring_rebuild_dc(
5359                            &st.qkv_mixed,
5360                            ring_old,
5361                            &mut rl.conv_state,
5362                            conv_dim,
5363                            acc,
5364                            base,
5365                            t_v,
5366                            d_conv,
5367                        )?;
5368                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
5369                        e.gdn_scan_s128_dc(
5370                            &st.q_l2,
5371                            &st.k_l2,
5372                            &st.v_g,
5373                            &st.g_log,
5374                            &st.beta,
5375                            state_in,
5376                            &mut rl.ssm_state,
5377                            &mut o,
5378                            num_v,
5379                            acc,
5380                            base,
5381                            t_v,
5382                            scale,
5383                        )?;
5384                    } else {
5385                        e.ssm_conv_ring_rebuild(
5386                            &st.qkv_mixed,
5387                            ring_old,
5388                            &mut rl.conv_state,
5389                            conv_dim,
5390                            j,
5391                            d_conv,
5392                        )?;
5393                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
5394                        e.gdn_scan_s128(
5395                            &st.q_l2,
5396                            &st.k_l2,
5397                            &st.v_g,
5398                            &st.g_log,
5399                            &st.beta,
5400                            state_in,
5401                            &mut rl.ssm_state,
5402                            &mut o,
5403                            num_v,
5404                            j,
5405                            scale,
5406                        )?;
5407                    }
5408                } else if let Some(cols) = &ckpt.cols[il] {
5409                    let (c, s) = &cols[j - 1];
5410                    e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
5411                    e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
5412                } else {
5413                    return Err(
5414                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
5415                    );
5416                }
5417            }
5418        }
5419        cache.pos = snap.pos + j;
5420        Ok(())
5421    }
5422
5423    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
5424    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
5425    fn commit_verified_prefix_stream(
5426        &self,
5427        e: &Engine,
5428        cache: &mut Cache,
5429        snap: &crate::cache::CacheSnapshot,
5430        ckpt: &VerifyCkpt,
5431        acc: &CudaSlice<u32>,
5432        base: usize,
5433        t_v: usize,
5434    ) -> Result<(), Box<dyn std::error::Error>> {
5435        let cfg = &self.cfg;
5436        let ssm = cfg.ssm.as_ref().unwrap();
5437        let d_state = ssm.state_size as usize;
5438        let num_k = ssm.group_count as usize;
5439        let num_v = ssm.time_step_rank as usize;
5440        let d_conv = ssm.conv_kernel as usize;
5441        let conv_dim = d_state * num_k * 2 + d_state * num_v;
5442        let scale = 1.0 / (d_state as f32).sqrt();
5443        for il in 0..self.layers.len() {
5444            if let Some(rl) = cache.recur[il].as_mut() {
5445                let st = ckpt.gdn[il]
5446                    .as_ref()
5447                    .ok_or("stream restore: batched-linear stash missing")?;
5448                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
5449                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
5450                e.ssm_conv_ring_rebuild_dc(
5451                    &st.qkv_mixed,
5452                    ring_old,
5453                    &mut rl.conv_state,
5454                    conv_dim,
5455                    acc,
5456                    base,
5457                    t_v,
5458                    d_conv,
5459                )?;
5460                let mut o = e.uninit(d_state * num_v * t_v)?;
5461                e.gdn_scan_s128_dc(
5462                    &st.q_l2,
5463                    &st.k_l2,
5464                    &st.v_g,
5465                    &st.g_log,
5466                    &st.beta,
5467                    state_in,
5468                    &mut rl.ssm_state,
5469                    &mut o,
5470                    num_v,
5471                    acc,
5472                    base,
5473                    t_v,
5474                    scale,
5475                )?;
5476            }
5477        }
5478        Ok(())
5479    }
5480
5481    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
5482    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
5483    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
5484    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
5485    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
5486    pub fn decode_step_t_aux2(
5487        &self,
5488        e: &Engine,
5489        tokens: &[u32],
5490        pos0: usize,
5491        cache: &mut Cache,
5492        aux_layers: &[usize],
5493        pred_col: Option<usize>,
5494    ) -> Result<
5495        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
5496        Box<dyn std::error::Error>,
5497    > {
5498        let cfg = &self.cfg;
5499        let n_embd = cfg.n_embd as usize;
5500        let eps = cfg.rms_eps;
5501        let t = tokens.len();
5502        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5503        let pos_d = e.htod_i32(&pos_vec)?;
5504        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
5505        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
5506        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
5507        let want_pred = pred_col.is_some();
5508
5509        for (il, layer) in self.layers.iter().enumerate() {
5510            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
5511            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
5512            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
5513            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
5514            if norm_fused {
5515                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5516            } else {
5517                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5518            }
5519            let mixed = match &layer.mixer {
5520                Mixer::Full(fa) => {
5521                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
5522                }
5523                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
5524                Mixer::Linear(la) => {
5525                    let mut out = e.zeros(t * n_embd)?;
5526                    for col in 0..t {
5527                        let mut h_col = e.zeros(n_embd)?;
5528                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
5529                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
5530                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
5531                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
5532                    }
5533                    out
5534                }
5535            };
5536            let ffn_fuse = match &layer.ffn {
5537                crate::hybrid::Ffn::Dense {
5538                    ffn_gate, ffn_up, ..
5539                } => {
5540                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
5541                        && e.uses_q8_1_fast(ffn_gate)
5542                        && e.uses_q8_1_fast(ffn_up)
5543                }
5544                crate::hybrid::Ffn::Moe(_) => false,
5545            };
5546            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
5547            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
5548            if ffn_fuse {
5549                e.add(&x, &mixed, &mut x1, t * n_embd)?;
5550                e.rms_norm_decode(
5551                    &x1,
5552                    layer.post_attn_norm.float_data(),
5553                    &mut z,
5554                    n_embd,
5555                    t,
5556                    eps,
5557                )?;
5558            } else {
5559                e.add_rms_norm(
5560                    &x,
5561                    &mixed,
5562                    layer.post_attn_norm.float_data(),
5563                    &mut x1,
5564                    &mut z,
5565                    n_embd,
5566                    t,
5567                    eps,
5568                )?;
5569            }
5570            let ffn_out = match &layer.ffn {
5571                crate::hybrid::Ffn::Dense {
5572                    ffn_gate,
5573                    ffn_up,
5574                    ffn_down,
5575                } => {
5576                    let n_ff = ffn_gate.out_features();
5577                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
5578                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
5579                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5580                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
5581                    Self::ffn_act_lim(
5582                        e,
5583                        &self.cfg,
5584                        &gate,
5585                        &up,
5586                        1.0,
5587                        1.0,
5588                        self.cfg.clamp_shexp_at(il as u32),
5589                        &mut act,
5590                        t * n_ff,
5591                    )?;
5592                    e.matmul_decode_exact(ffn_down, &act, t)?
5593                }
5594                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5595            };
5596            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5597            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5598            if aux_layers.contains(&il) {
5599                let mut a = e.zeros(n_embd)?;
5600                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5601                aux_last.push(a);
5602                if let Some(pc) = pred_col {
5603                    let mut ap = e.zeros(n_embd)?;
5604                    e.copy_view_into(
5605                        &mut ap,
5606                        0,
5607                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
5608                        n_embd,
5609                    )?;
5610                    aux_pred.push(ap);
5611                }
5612            }
5613            x = x2;
5614        }
5615        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
5616        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5617        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
5618        let host = e.dtoh(&logits)?;
5619        cache.pos += t;
5620        Ok((
5621            host,
5622            aux_last,
5623            if want_pred { Some(aux_pred) } else { None },
5624        ))
5625    }
5626
5627    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
5628    /// `step35_decode_attn`.
5629    ///
5630    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
5631    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
5632    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
5633    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
5634    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
5635    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
5636    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
5637    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
5638    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
5639    /// position of each query row. A batched twin would have to reproduce all of that AND the
5640    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
5641    /// take one `base_len`, not a per-row offset).
5642    ///
5643    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
5644    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
5645    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
5646    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
5647    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
5648    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
5649    /// step35 twin is a perf lane's job and must be gated against this arm.
5650    ///
5651    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
5652    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
5653    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
5654    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
5655    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
5656    #[allow(clippy::too_many_arguments)]
5657    fn step35_verify(
5658        &self,
5659        e: &Engine,
5660        fa: &FullAttnLayer,
5661        h: &CudaSlice<f32>,
5662        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5663        t: usize,
5664        cache: &mut Cache,
5665        il: usize,
5666    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5667        let n_embd = self.cfg.n_embd as usize;
5668        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
5669        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
5670        // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
5671        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
5672        // cannot regress it into silently reading an empty buffer.
5673        assert_eq!(
5674            h.len(),
5675            t * n_embd,
5676            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
5677             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
5678            h_q8.is_some()
5679        );
5680        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
5681        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
5682        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
5683        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
5684        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
5685        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
5686        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
5687        for r in 0..t {
5688            // Absolute position of this query row. `cache.pos` is the committed length at round
5689            // start and every row before r has already been appended by this loop, so the r-th
5690            // verify token sits at cache.pos + r — the same position eager decode would give it.
5691            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
5692            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
5693            e.copy_view_into(
5694                &mut h_row,
5695                0,
5696                &h.slice(r * n_embd..(r + 1) * n_embd),
5697                n_embd,
5698            )?;
5699            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
5700            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
5701            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
5702            debug_assert_eq!(
5703                o.len(),
5704                n_embd,
5705                "step35_decode_attn returns post-wo [n_embd]"
5706            );
5707            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
5708        }
5709        Ok(out)
5710    }
5711
5712    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
5713    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
5714    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
5715    #[allow(clippy::too_many_arguments)]
5716    fn full_attn_verify(
5717        &self,
5718        e: &Engine,
5719        fa: &FullAttnLayer,
5720        h: &CudaSlice<f32>,
5721        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5722        pos_d: &CudaSlice<i32>,
5723        t: usize,
5724        cache: &mut Cache,
5725        il: usize,
5726        stream_ctr: Option<&CudaSlice<i32>>,
5727    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5728        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
5729        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
5730        // its own arm. A verify that silently computes different attention than decode defeats the
5731        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
5732        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
5733        // shape and not laziness.
5734        if self.cfg.step35.is_some() {
5735            if stream_ctr.is_some() {
5736                return Err(
5737                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5738                            cannot express the SWA offset KV view; same root cause as the dc \
5739                            decode refusal) — run spec without the stream arm"
5740                        .into(),
5741                );
5742            }
5743            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
5744        }
5745        let cfg = &self.cfg;
5746        let geometry = cfg.full_attention_geometry_at(il as u32);
5747        let n_head = geometry.n_head as usize;
5748        let n_head_kv = geometry.n_head_kv as usize;
5749        let head_dim = geometry.head_dim_k as usize;
5750        let eps = cfg.rms_eps;
5751        let scale = geometry.attention_scale();
5752        let n_embd = cfg.n_embd as usize;
5753
5754        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
5755        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
5756        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
5757        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
5758        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
5759        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
5760        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
5761        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
5762        let (qf, mut k, v) = {
5763            let mut fused = None;
5764            let qkv_fast =
5765                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
5766            if t == 1 && qkv_fast {
5767                let (hq_o, hd_o);
5768                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5769                    Some(p) => p,
5770                    None => {
5771                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
5772                        (&hq_o, &hd_o)
5773                    }
5774                };
5775                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
5776            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
5777                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
5778                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
5779                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
5780                let (hq_o, hd_o);
5781                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
5782                    Some(p) => p,
5783                    None => {
5784                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
5785                        (&hq_o, &hd_o)
5786                    }
5787                };
5788                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
5789            }
5790            match (fused, h_q8) {
5791                (Some(triple), _) => triple,
5792                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
5793                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
5794                (None, Some((hq, hd))) if qkv_fast => (
5795                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
5796                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
5797                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
5798                ),
5799                (None, _) => (
5800                    e.matmul_decode_exact(&fa.wq, h, t)?,
5801                    e.matmul_decode_exact(&fa.wk, h, t)?,
5802                    e.matmul_decode_exact(&fa.wv, h, t)?,
5803                ),
5804            }
5805        };
5806        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
5807        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5808        let (mut q, gate) = if gated {
5809            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5810            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
5811            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
5812            (q, Some(gate))
5813        } else {
5814            (qf, None)
5815        };
5816
5817        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
5818        e.rms_norm(
5819            &q,
5820            fa.q_norm.float_data(),
5821            &mut qn,
5822            head_dim,
5823            n_head * t,
5824            eps,
5825        )?;
5826        q = qn;
5827        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
5828        e.rms_norm(
5829            &k,
5830            fa.k_norm.float_data(),
5831            &mut kn,
5832            head_dim,
5833            n_head_kv * t,
5834            eps,
5835        )?;
5836        k = kn;
5837        let rope_dims = geometry.n_rot as usize;
5838        e.rope_neox(
5839            &mut q,
5840            pos_d,
5841            head_dim,
5842            rope_dims,
5843            n_head,
5844            t,
5845            geometry.rope_base,
5846            1.0,
5847        )?;
5848        e.rope_neox(
5849            &mut k,
5850            pos_d,
5851            head_dim,
5852            rope_dims,
5853            n_head_kv,
5854            t,
5855            geometry.rope_base,
5856            1.0,
5857        )?;
5858
5859        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
5860        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
5861        let kvl = cache.kv[il].as_mut().unwrap();
5862        let (kv_dim_k, kv_dim_v, ktb, vtb) =
5863            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
5864        if let Some(ctr) = stream_ctr {
5865            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
5866            // math on a (block, token) grid, documented byte-identical); host len is a stale
5867            // LOWER BOUND under pre-issue (drain reconciles it).
5868            e.append_kv_quantized_rows_dc(
5869                &k,
5870                &v,
5871                &mut kvl.k,
5872                &mut kvl.v,
5873                ctr,
5874                t,
5875                kv_dim_k,
5876                kv_dim_v,
5877                ktb,
5878                vtb,
5879                crate::Engine::kv_fp8_on(),
5880            )?;
5881        } else {
5882            for i in 0..t {
5883                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
5884                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
5885                e.append_kv_quantized_view(
5886                    &k_row,
5887                    &v_row,
5888                    &mut kvl.k,
5889                    &mut kvl.v,
5890                    kvl.len + i,
5891                    kv_dim_k,
5892                    kv_dim_v,
5893                    ktb,
5894                    vtb,
5895                    crate::Engine::kv_fp8_on(),
5896                )?;
5897            }
5898            kvl.len += t;
5899        }
5900
5901        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
5902        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
5903        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
5904        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
5905        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
5906        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
5907        // keys. The verify appends all T tokens first but bounds the key range per row.
5908        //
5909        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
5910        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
5911        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
5912        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
5913        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
5914        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
5915        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
5916        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
5917        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
5918        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
5919        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
5920        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
5921        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
5922        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
5923        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
5924        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
5925        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
5926        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
5927        if let Some(ctr) = stream_ctr {
5928            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
5929            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
5930            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
5931            let upper = kvl.len + t + 64;
5932            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
5933            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
5934            e.fa_decode_rows_dc(
5935                &q,
5936                &k_view,
5937                &v_view,
5938                &mut attn,
5939                head_dim,
5940                n_head,
5941                n_head_kv,
5942                ctr,
5943                upper.min(cache.max_ctx),
5944                t,
5945                scale,
5946                ktb,
5947                vtb,
5948                0,
5949                false,
5950            )?;
5951        } else if spec_lean() && t == 1 {
5952            let t_kv = base_len + 1;
5953            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
5954            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
5955            e.fa_decode_kvmod(
5956                &q,
5957                &k_view,
5958                &v_view,
5959                &mut attn,
5960                head_dim,
5961                n_head,
5962                n_head_kv,
5963                t_kv,
5964                scale,
5965                ktb,
5966                vtb,
5967                crate::Engine::kv_fp8_on(),
5968            )?;
5969        } else if e.fa_rows_eligible(base_len, head_dim) {
5970            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
5971            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
5972            e.fa_decode_rows(
5973                &q,
5974                &k_view,
5975                &v_view,
5976                &mut attn,
5977                head_dim,
5978                n_head,
5979                n_head_kv,
5980                base_len,
5981                t,
5982                scale,
5983                ktb,
5984                vtb,
5985                None,
5986                false,
5987                crate::Engine::kv_fp8_on(),
5988                None,
5989            )?;
5990        } else {
5991            for r in 0..t {
5992                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
5993                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
5994                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
5995                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
5996                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
5997                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
5998                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
5999                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
6000                e.fa_decode_kvmod(
6001                    &q_row,
6002                    &k_view_r,
6003                    &v_view_r,
6004                    &mut attn_row,
6005                    head_dim,
6006                    n_head,
6007                    n_head_kv,
6008                    t_kv_r,
6009                    scale,
6010                    ktb,
6011                    vtb,
6012                    crate::Engine::kv_fp8_on(),
6013                )?;
6014                e.copy_into(
6015                    &mut attn,
6016                    r * n_head * head_dim,
6017                    &attn_row,
6018                    n_head * head_dim,
6019                )?;
6020            }
6021        }
6022
6023        let attn_g = match &gate {
6024            Some(gate) => {
6025                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
6026                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
6027                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
6028                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
6029                ag
6030            }
6031            None => attn,
6032        };
6033        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
6034        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
6035        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
6036    }
6037
6038    /// Context-linear bytes for a plain serving session's trunk cache.
6039    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
6040        crate::cache::cache_bytes_per_token(&self.cfg)
6041    }
6042
6043    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
6044    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
6045        (
6046            self.plain_session_kv_bytes_per_token(),
6047            crate::cache::cache_ring_bytes_per_token(&self.cfg),
6048            crate::cache::cache_ring_row_cap(&self.cfg),
6049        )
6050    }
6051
6052    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
6053    /// scratch. With no MTP head this equals the plain coefficient.
6054    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
6055        let scratch = self
6056            .mtp
6057            .as_ref()
6058            .map(|mtp| {
6059                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6060                k + v
6061            })
6062            .unwrap_or(0);
6063        self.plain_session_kv_bytes_per_token()
6064            .saturating_add(scratch)
6065    }
6066
6067    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
6068    /// capped by the same SWA ring rows as the trunk.
6069    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
6070        let total = self.spec_session_kv_bytes_per_token();
6071        let (_, mut ring, rows) = self.plain_session_kv_shape();
6072        if rows > 0 {
6073            ring = ring.saturating_add(
6074                self.mtp
6075                    .as_ref()
6076                    .map(|mtp| {
6077                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6078                        k + v
6079                    })
6080                    .unwrap_or(0),
6081            );
6082        }
6083        (total, ring, rows)
6084    }
6085
6086    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
6087    /// the NextN head to draft K tokens then verifies them in one batched target forward.
6088    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
6089    /// acceptance rate. `k` = draft length per round.
6090    ///
6091    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
6092    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
6093    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
6094    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
6095    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
6096    /// captured graph references is event-free; the spec loop is strictly single-stream.
6097    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
6098    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
6099    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
6100    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
6101    /// generate_spec_inner2.
6102    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
6103    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
6104    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
6105    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
6106    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
6107    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
6108    pub fn new_session(
6109        &self,
6110        e: &Engine,
6111        max_ctx: usize,
6112    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
6113        Ok(SpecSession {
6114            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
6115            // is the SERVING spec-session path, and with the ppN door open across two cards a
6116            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
6117            // round — the wrong-card class already fixed on the two batched serving paths
6118            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
6119            // branch, same allocations), so single-device behavior is byte-unchanged.
6120            cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
6121            scratch: MtpScratch::new(
6122                e,
6123                &self.cfg,
6124                max_ctx,
6125                self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6126            )?,
6127            committed: Vec::new(),
6128            last_h: None,
6129            next_pred: None,
6130            sctr: 0,
6131            uctr: 0,
6132            draft_ctx: None,
6133            pending_tok: None,
6134            turn_ckpt: None,
6135            telem: SpecTelemetryCounters::default(),
6136            capture_at: None,
6137            boundary_capture: None,
6138        })
6139    }
6140
6141    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
6142    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
6143    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
6144    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
6145    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
6146    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
6147    /// worker always receives a fully-warm continuation session (committed = whole
6148    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
6149    /// boundary logits on the empty-suffix shape).
6150    ///
6151    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
6152    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
6153    /// request, and plain feeds a carried suffix via eager `decode_step` below
6154    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
6155    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
6156    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
6157    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
6158    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
6159    /// burst prime.
6160    ///
6161    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
6162    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
6163    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
6164    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
6165    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
6166    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
6167    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
6168    /// cold session draws from the identical row at counter 0 and then runs its rounds from
6169    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
6170    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
6171    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
6172    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
6173    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
6174    ///
6175    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
6176    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
6177    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
6178    /// and are never routed here.
6179    ///
6180    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
6181    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
6182    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
6183    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
6184    /// entry stays published for the next request.
6185    #[allow(clippy::too_many_arguments)]
6186    pub fn spec_session_from_restored(
6187        &self,
6188        e: &Engine,
6189        mut cache: Cache,
6190        prefix: Vec<u32>,
6191        suffix: &[u32],
6192        draft_k: &CudaSlice<u8>,
6193        draft_v: &CudaSlice<u8>,
6194        draft_k_tok_bytes: usize,
6195        draft_v_tok_bytes: usize,
6196        draft_len: usize,
6197        last_h: &[f32],
6198        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
6199        // when a suffix follows — the feed's own logits are the boundary then.
6200        boundary_logits: &[f32],
6201        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
6202        // ONE place instead of being half-applied by the worker.
6203        sampling: Option<SpecSampling>,
6204        require_anchor: bool,
6205        max_ctx: usize,
6206    ) -> Result<SpecSession, (Option<Cache>, String)> {
6207        let pos = prefix.len();
6208        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
6209            Err((Some(cache), msg))
6210        };
6211        if self.mtp.is_none() {
6212            return fail(cache, "no MTP head attached (nothing to draft with)".into());
6213        }
6214        if pos == 0 {
6215            return fail(cache, "empty committed prefix".into());
6216        }
6217        if cache.pos != pos {
6218            let msg = format!(
6219                "restored cache pos {} != restored prefix len {pos}",
6220                cache.pos
6221            );
6222            return fail(cache, msg);
6223        }
6224        if draft_len != pos {
6225            return fail(
6226                cache,
6227                format!("draft plane len {draft_len} != restored prefix len {pos}"),
6228            );
6229        }
6230        if pos + suffix.len() >= max_ctx {
6231            return fail(
6232                cache,
6233                format!(
6234                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
6235                    pos + suffix.len(),
6236                ),
6237            );
6238        }
6239        let mut scratch = match MtpScratch::new(
6240            e,
6241            &self.cfg,
6242            max_ctx,
6243            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6244        ) {
6245            Ok(s) => s,
6246            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
6247        };
6248        if scratch.kv.ring.is_some() {
6249            return fail(
6250                cache,
6251                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
6252            );
6253        }
6254        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
6255            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
6256        {
6257            return fail(
6258                cache,
6259                format!(
6260                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
6261                     {}/{} bytes/token (stale entry across a format change)",
6262                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
6263                ),
6264            );
6265        }
6266        if pos > scratch.cap {
6267            return fail(
6268                cache,
6269                format!(
6270                    "draft plane rows {pos} exceed scratch capacity {}",
6271                    scratch.cap
6272                ),
6273            );
6274        }
6275        let kb = pos * draft_k_tok_bytes;
6276        let vb = pos * draft_v_tok_bytes;
6277        if draft_k.len() < kb || draft_v.len() < vb {
6278            return fail(
6279                cache,
6280                format!(
6281                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
6282                    draft_k.len(),
6283                    draft_v.len(),
6284                ),
6285            );
6286        }
6287        if kb > 0 {
6288            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
6289                return fail(cache, format!("draft K restore copy failed: {err}"));
6290            }
6291        }
6292        if vb > 0 {
6293            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
6294                return fail(cache, format!("draft V restore copy failed: {err}"));
6295            }
6296        }
6297        if let Err(err) = scratch.set_len(e, pos) {
6298            return fail(cache, format!("draft scratch len set failed: {err}"));
6299        }
6300        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
6301            // anchor upload failure is acceptance-only when a suffix feed follows (fill
6302            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
6303            // burst entry asserts committed + last_h + next_pred) — the caller says which.
6304            e.htod(last_h).ok()
6305        } else {
6306            None
6307        };
6308        if require_anchor && last_h_dev.is_none() {
6309            return fail(
6310                cache,
6311                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
6312            );
6313        }
6314        let mut committed = prefix;
6315        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
6316        // what the empty-suffix continuation assert in the burst entry requires.
6317        let next_pred;
6318        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
6319        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
6320        // drawing its own first token from the same row.
6321        let mut sctr = 0u32;
6322        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
6323        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
6324        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
6325        // after the suffix joins `committed` below.
6326        let mut boundary_capture: Option<SpecBoundaryCapture> = None;
6327        if !suffix.is_empty() {
6328            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
6329            // From here on the trunk cache mutates: failures return Err((None, _)) and
6330            // the worker serves the request cold-plain instead of reusing the carrier.
6331            let dirty =
6332                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
6333            let n_embd = self.cfg.n_embd as usize;
6334            let t = suffix.len();
6335            let mut h_rows = match e.uninit(t * n_embd) {
6336                Ok(b) => b,
6337                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
6338            };
6339            let mut feed_logits = Vec::new();
6340            let batched = t >= crate::hybrid_forward::PRIME_MIN_T
6341                && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
6342                && !e.frozen_cpu_experts_prefer_tokenwise_prime();
6343            if batched {
6344                // prefill_tick's prime arm: one request-level prime_cache call.
6345                match self.prime_cache(e, suffix, &mut cache, 0) {
6346                    Ok((l, _h_seed, hiddens)) => {
6347                        if let Err(err) = e.copy_into(&mut h_rows, 0, &hiddens, t * n_embd) {
6348                            return dirty(format!("suffix hidden copy: {err}"));
6349                        }
6350                        feed_logits = l;
6351                    }
6352                    Err(err) => return dirty(format!("suffix prime failed: {err}")),
6353                }
6354            } else {
6355                // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
6356                for (i, &tok) in suffix.iter().enumerate() {
6357                    match self.decode_step_h(e, tok, &mut cache) {
6358                        Ok((l, h)) => {
6359                            if let Err(err) = e.copy_into(&mut h_rows, i * n_embd, &h, n_embd) {
6360                                return dirty(format!("suffix hidden copy: {err}"));
6361                            }
6362                            feed_logits = l;
6363                        }
6364                        Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
6365                    }
6366                }
6367            }
6368            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
6369            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
6370            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
6371            // with T). Fill failures are acceptance-only — truncate to the restored rows
6372            // and continue; the burst's own set_len keeps the invariant.
6373            let mtp = self.mtp.as_ref().expect("mtp checked above");
6374            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6375            let embd_gpu = if spec_host_embd() {
6376                None
6377            } else {
6378                Some(
6379                    self.embd_gpu
6380                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6381                )
6382            };
6383            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6384            let fill_chunk = 4096usize;
6385            let mut filled = true;
6386            let mut start = 0usize;
6387            'fill: while start < t {
6388                let end = (start + fill_chunk).min(t);
6389                let tc = end - start;
6390                let Ok(mut phs) = e.zeros(tc * n_embd) else {
6391                    filled = false;
6392                    break 'fill;
6393                };
6394                let (src_lo, dst_off, n_copy) = if start == 0 {
6395                    (0, n_embd, (tc - 1) * n_embd)
6396                } else {
6397                    ((start - 1) * n_embd, 0, tc * n_embd)
6398                };
6399                if start == 0 {
6400                    if let Some(lh) = last_h_dev.as_ref() {
6401                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
6402                            filled = false;
6403                            break 'fill;
6404                        }
6405                    }
6406                }
6407                if n_copy > 0
6408                    && e.copy_view_into(
6409                        &mut phs,
6410                        dst_off,
6411                        &h_rows.slice(src_lo..src_lo + n_copy),
6412                        n_copy,
6413                    )
6414                    .is_err()
6415                {
6416                    filled = false;
6417                    break 'fill;
6418                }
6419                if self
6420                    .mtp_kv_fill(
6421                        e,
6422                        mtp,
6423                        &suffix[start..end],
6424                        &phs,
6425                        pos + start,
6426                        &mut scratch,
6427                        embd_dev,
6428                    )
6429                    .is_err()
6430                {
6431                    filled = false;
6432                    break 'fill;
6433                }
6434                start = end;
6435            }
6436            if !filled {
6437                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
6438                // so keep only the restored rows resident and let verify arbitrate.
6439                if let Err(err) = scratch.set_len(e, pos) {
6440                    return dirty(format!("scratch truncation after failed fill: {err}"));
6441                }
6442            }
6443            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
6444            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
6445            // finding (d)). Pre-lane, publication was armed only for COLD sessions
6446            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
6447            // non-continuation burst — but a converted hit's first burst IS a continuation,
6448            // so a growing conversation learned exactly ONE boundary and turn 3 could never
6449            // hit a longer prefix than turn 2 did.
6450            //
6451            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
6452            // line — the trunk is primed over the whole prompt, nothing is generated, and the
6453            // draft plane rows [0..prompt) are filled just above. That is a complete
6454            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
6455            // publishes; the worker's existing publication sweep picks it up because it is
6456            // keyed on `boundary_capture.is_some()` and is sampler- and resume-independent.
6457            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
6458            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
6459            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
6460            // publication is an optimization, never a correctness dependency.
6461            if spec_restore_republish_on() {
6462                debug_assert_eq!(
6463                    cache.pos,
6464                    pos + t,
6465                    "extended-entry capture must sit at the restored session's prompt end",
6466                );
6467                if let Ok(snap) = cache.snapshot(e) {
6468                    boundary_capture = Some(SpecBoundaryCapture {
6469                        snap,
6470                        pos: pos + t,
6471                        logits: feed_logits.clone(),
6472                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
6473                    });
6474                }
6475            }
6476            // continuation seed: the feed's boundary logits ARE the plain path's boundary
6477            // logits (same program), so greedy's argmax here is plain's first emitted token,
6478            // and the sampled draw is the cold sampled session's own first token.
6479            next_pred = Some(if sampled {
6480                let sp = sampling.expect("sampled implies a sampler");
6481                // `committed` is still the restored prefix here; the suffix joins it below —
6482                // so this is the last-N window over the WHOLE prompt, exactly the cold
6483                // session's own window at its first token.
6484                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
6485                match sample_boundary_token(
6486                    e,
6487                    &feed_logits,
6488                    &sp,
6489                    &hist,
6490                    &mut sctr,
6491                    "restore-suffix-feed",
6492                ) {
6493                    Ok(t) => t,
6494                    // the trunk is already fed: hand nothing back, the worker serves the
6495                    // request cold-plain. Never fall back to an argmax — that would put a
6496                    // greedy token in a sampled stream to save a slow path.
6497                    Err(err) => {
6498                        return dirty(format!("boundary token draw failed: {err}"));
6499                    }
6500                }
6501            } else {
6502                argmax(&feed_logits) as u32
6503            });
6504            let mut lh = match e.uninit(n_embd) {
6505                Ok(b) => b,
6506                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
6507            };
6508            if let Err(err) = e.copy_view_into(
6509                &mut lh,
6510                0,
6511                &h_rows.slice((t - 1) * n_embd..t * n_embd),
6512                n_embd,
6513            ) {
6514                return dirty(format!("boundary hidden copy: {err}"));
6515            }
6516            last_h_dev = Some(lh);
6517            committed.extend_from_slice(suffix);
6518        } else {
6519            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
6520            // ENTRY's boundary logits are the boundary row, and this is the token the cold
6521            // session emits from that same row. Owned here rather than in the worker so the
6522            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
6523            if boundary_logits.is_empty() {
6524                return fail(
6525                    cache,
6526                    "full-cover restore without the entry's boundary logits".into(),
6527                );
6528            }
6529            next_pred = Some(if sampled {
6530                let sp = sampling.expect("sampled implies a sampler");
6531                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
6532                match sample_boundary_token(
6533                    e,
6534                    boundary_logits,
6535                    &sp,
6536                    &hist,
6537                    &mut sctr,
6538                    "restore-full-cover",
6539                ) {
6540                    Ok(t) => t,
6541                    // nothing has been mutated on this shape — hand the carrier back and let
6542                    // the hit serve PLAIN (the banked pre-lane path).
6543                    Err(err) => {
6544                        return fail(cache, format!("boundary token draw failed: {err}"));
6545                    }
6546                }
6547            } else {
6548                argmax(boundary_logits) as u32
6549            });
6550        }
6551        Ok(SpecSession {
6552            cache,
6553            scratch,
6554            committed,
6555            last_h: last_h_dev,
6556            next_pred,
6557            sctr,
6558            uctr: 0,
6559            draft_ctx: None,
6560            pending_tok: None,
6561            turn_ckpt: None,
6562            telem: SpecTelemetryCounters::default(),
6563            capture_at: None,
6564            boundary_capture,
6565        })
6566    }
6567
6568    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
6569    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
6570    /// snapshot, or draft-KV row that only corrupts the following round.
6571    pub fn optipipe_compare_session_state(
6572        &self,
6573        e: &Engine,
6574        reference: &SpecSession,
6575        candidate: &SpecSession,
6576    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
6577        fn fail(what: &str) -> Box<dyn std::error::Error> {
6578            format!("optipipe state mismatch: {what}").into()
6579        }
6580        fn same_f32(a: &[f32], b: &[f32]) -> bool {
6581            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
6582        }
6583        fn compare_layers(
6584            es: &Engine,
6585            range: std::ops::Range<usize>,
6586            reference: &SpecSession,
6587            candidate: &SpecSession,
6588            report: &mut OptiForkStateIdentity,
6589        ) -> Result<(), Box<dyn std::error::Error>> {
6590            for il in range {
6591                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
6592                    (Some(a), Some(b)) => {
6593                        if a.len != b.len {
6594                            return Err(fail(&format!(
6595                                "layer {il} host KV len {} != {}",
6596                                a.len, b.len
6597                            )));
6598                        }
6599                        let ad = es.dtoh_i32(&a.len_d)?;
6600                        let bd = es.dtoh_i32(&b.len_d)?;
6601                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
6602                            return Err(fail(&format!(
6603                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
6604                                a.len,
6605                            )));
6606                        }
6607                        let kb = a.len * a.k_tok_bytes;
6608                        let vb = a.len * a.v_tok_bytes;
6609                        if kb > 0 {
6610                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
6611                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
6612                            if ak != bk {
6613                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
6614                                return Err(fail(&format!(
6615                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
6616                                    at / a.k_tok_bytes,
6617                                    at % a.k_tok_bytes,
6618                                    ak[at],
6619                                    bk[at],
6620                                )));
6621                            }
6622                        }
6623                        if vb > 0 {
6624                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
6625                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
6626                            if av != bv {
6627                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
6628                                return Err(fail(&format!(
6629                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
6630                                    at / a.v_tok_bytes,
6631                                    at % a.v_tok_bytes,
6632                                    av[at],
6633                                    bv[at],
6634                                )));
6635                            }
6636                        }
6637                        report.trunk_kv_bytes += kb + vb;
6638                    }
6639                    (None, None) => {}
6640                    _ => return Err(fail(&format!("layer {il} KV presence"))),
6641                }
6642                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
6643                    (Some(a), Some(b)) => {
6644                        let ac = es.dtoh(&a.conv_state)?;
6645                        let bc = es.dtoh(&b.conv_state)?;
6646                        if !same_f32(&ac, &bc) {
6647                            return Err(fail(&format!("layer {il} conv state")));
6648                        }
6649                        let as_ = es.dtoh(&a.ssm_state)?;
6650                        let bs = es.dtoh(&b.ssm_state)?;
6651                        if !same_f32(&as_, &bs) {
6652                            return Err(fail(&format!("layer {il} SSM state")));
6653                        }
6654                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
6655                    }
6656                    (None, None) => {}
6657                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
6658                }
6659            }
6660            Ok(())
6661        }
6662
6663        if reference.committed != candidate.committed {
6664            return Err(fail("committed token ids"));
6665        }
6666        if reference.cache.pos != candidate.cache.pos
6667            || reference.cache.max_ctx != candidate.cache.max_ctx
6668        {
6669            return Err(fail("cache pos/capacity"));
6670        }
6671        if reference.pending_tok != candidate.pending_tok
6672            || reference.next_pred != candidate.next_pred
6673            || reference.sctr != candidate.sctr
6674            || reference.uctr != candidate.uctr
6675        {
6676            return Err(fail("pending/prediction/counter tail"));
6677        }
6678
6679        let mut report = OptiForkStateIdentity::default();
6680        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
6681            let rt = crate::pp::PpNRt::get(e)?;
6682            for stage in 0..rt.n_stages() {
6683                let _scope = rt.enter(stage);
6684                compare_layers(
6685                    rt.engine(stage, e),
6686                    fence[stage]..fence[stage + 1],
6687                    reference,
6688                    candidate,
6689                    &mut report,
6690                )?;
6691            }
6692        } else {
6693            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
6694        }
6695
6696        let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
6697        if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
6698            return Err(fail("draft scratch length"));
6699        }
6700        let kb = a.len * a.k_tok_bytes;
6701        let vb = a.len * a.v_tok_bytes;
6702        if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
6703            return Err(fail("draft scratch K bytes"));
6704        }
6705        if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
6706            return Err(fail("draft scratch V bytes"));
6707        }
6708        report.scratch_kv_bytes = kb + vb;
6709
6710        match (&reference.last_h, &candidate.last_h) {
6711            (Some(a), Some(b)) => {
6712                let ah = e.dtoh(a)?;
6713                let bh = e.dtoh(b)?;
6714                if !same_f32(&ah, &bh) {
6715                    return Err(fail("last hidden/seed bytes"));
6716                }
6717                report.hidden_bytes = ah.len() * 4;
6718            }
6719            (None, None) => {}
6720            _ => return Err(fail("last hidden/seed presence")),
6721        }
6722        Ok(report)
6723    }
6724
6725    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
6726    /// retained prompt-end checkpoint, so a request whose prompt matches
6727    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
6728    ///
6729    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
6730    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
6731    /// restored from the device copy taken there, draft scratch length reset, `committed`
6732    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
6733    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
6734    /// every burst after it are identical to a cold run of the same token stream — the
6735    /// committed-tokens-authoritative contract.
6736    ///
6737    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
6738    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
6739    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
6740    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
6741    /// (the scratch KV, the resident embedding), none of which the rewind moves.
6742    ///
6743    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
6744    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
6745    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
6746    pub fn spec_rewind_to_checkpoint(
6747        &self,
6748        e: &Engine,
6749        sess: &mut SpecSession,
6750    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
6751        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
6752            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
6753        }) {
6754            return Err(
6755                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
6756            );
6757        }
6758        let Some(ckpt) = sess.turn_ckpt.take() else {
6759            return Ok(None);
6760        };
6761        assert!(
6762            ckpt.pos <= sess.committed.len(),
6763            "checkpoint past committed ({} > {})",
6764            ckpt.pos,
6765            sess.committed.len()
6766        );
6767        // Restore through each layer's owning engine. A single primary-engine rollback is not
6768        // sufficient when the serving cache is stage-owned under cross-device PP.
6769        crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
6770        debug_assert_eq!(
6771            sess.cache.pos, ckpt.pos,
6772            "rollback landed off the checkpoint"
6773        );
6774        sess.scratch.set_len(e, ckpt.pos)?;
6775        sess.committed.truncate(ckpt.pos);
6776        sess.last_h = Some(ckpt.last_h);
6777        sess.next_pred = None;
6778        sess.pending_tok = None;
6779        Ok(Some(ckpt.pos))
6780    }
6781
6782    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
6783    /// checkpoint without re-priming the checkpoint prefix.
6784    ///
6785    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
6786    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
6787    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
6788    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
6789    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
6790    ///
6791    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
6792    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
6793    pub fn spec_grow_and_rewind_to_checkpoint(
6794        &self,
6795        e: &Engine,
6796        sess: &mut SpecSession,
6797        target_cap: usize,
6798    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
6799        if target_cap <= sess.cache.max_ctx {
6800            return self.spec_rewind_to_checkpoint(e, sess);
6801        }
6802        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
6803            return Ok(None);
6804        };
6805        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
6806            return Err(format!(
6807                "checkpoint pos {} outside committed length {}",
6808                ckpt.pos,
6809                sess.committed.len(),
6810            )
6811            .into());
6812        }
6813        if ckpt.pos > target_cap {
6814            return Err(format!(
6815                "checkpoint pos {} exceeds grown capacity {target_cap}",
6816                ckpt.pos,
6817            )
6818            .into());
6819        }
6820
6821        let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
6822        let mut grown_scratch = MtpScratch::new(
6823            e,
6824            &self.cfg,
6825            target_cap,
6826            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6827        )?;
6828        crate::pp::restore_cache_checkpoint(
6829            e,
6830            &self.cfg,
6831            Some(&sess.cache),
6832            &mut grown_cache,
6833            &ckpt.snap,
6834        )?;
6835
6836        let src = &sess.scratch.kv;
6837        let dst = &mut grown_scratch.kv;
6838        if ckpt.pos > src.len
6839            || src.kv_dim_k != dst.kv_dim_k
6840            || src.kv_dim_v != dst.kv_dim_v
6841            || src.k_tok_bytes != dst.k_tok_bytes
6842            || src.v_tok_bytes != dst.v_tok_bytes
6843        {
6844            return Err(format!(
6845                "checkpoint draft layout mismatch (pos {}, source len {})",
6846                ckpt.pos, src.len,
6847            )
6848            .into());
6849        }
6850        let kb = ckpt.pos * src.k_tok_bytes;
6851        let vb = ckpt.pos * src.v_tok_bytes;
6852        if kb > 0 {
6853            e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
6854        }
6855        if vb > 0 {
6856            e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
6857        }
6858        grown_scratch.set_len(e, ckpt.pos)?;
6859        // The old scratch is dropped immediately after publication below. Bound its D2D reads
6860        // first; growth happens once per rewritten turn, outside the decode hot loop.
6861        e.stream().synchronize()?;
6862
6863        let ckpt = sess
6864            .turn_ckpt
6865            .take()
6866            .expect("checkpoint remained present through transactional grow");
6867        let pos = ckpt.pos;
6868        sess.cache = grown_cache;
6869        sess.scratch = grown_scratch;
6870        sess.committed.truncate(pos);
6871        sess.last_h = Some(ckpt.last_h);
6872        sess.next_pred = None;
6873        sess.pending_tok = None;
6874        sess.draft_ctx = None;
6875        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
6876        debug_assert_eq!(
6877            sess.scratch.kv.len, pos,
6878            "grown draft rewind landed off checkpoint"
6879        );
6880        Ok(Some(pos))
6881    }
6882
6883    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
6884    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
6885    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
6886    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
6887    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
6888    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
6889    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
6890    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
6891    /// park-time flush is a future request whose sampler is not knowable here (residual
6892    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
6893    pub fn spec_flush_pending(
6894        &self,
6895        e: &Engine,
6896        sess: &mut SpecSession,
6897        sampling: Option<SpecSampling>,
6898    ) -> Result<(), Box<dyn std::error::Error>> {
6899        let Some(b) = sess.pending_tok.take() else {
6900            return Ok(());
6901        };
6902        let mtp = self
6903            .mtp
6904            .as_ref()
6905            .expect("pending carry requires an MTP head");
6906        let n_embd = self.cfg.n_embd as usize;
6907        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
6908        let embd_gpu = if spec_host_embd() {
6909            None
6910        } else {
6911            Some(
6912                self.embd_gpu
6913                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
6914            )
6915        };
6916        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
6917        let pos_b = sess.cache.pos;
6918        sess.scratch.set_len(e, pos_b)?;
6919        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
6920        sess.next_pred = Some(match sampling {
6921            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
6922                // window includes `b` itself: it is committed by this pass, and the pre-lane
6923                // code never counted a boundary token in the penalty history at all.
6924                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
6925                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
6926            }
6927            _ => argmax(&lg_b) as u32,
6928        });
6929        let anchor = sess
6930            .last_h
6931            .as_ref()
6932            .expect("pending carry requires last_h (the predecessor-row anchor)");
6933        self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
6934        sess.last_h = Some(hb);
6935        sess.committed.push(b);
6936        Ok(())
6937    }
6938
6939    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
6940    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
6941    /// rounds through that same graph. Other model families keep their eager T=1 contract.
6942    fn spec_target_step_h(
6943        &self,
6944        e: &Engine,
6945        token: u32,
6946        cache: &mut Cache,
6947    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6948        if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
6949            return self.decode_step_h(e, token, cache);
6950        }
6951        let pos0 = cache.pos;
6952        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
6953        Ok((e.dtoh(&logits)?, hidden))
6954    }
6955
6956    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
6957    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
6958    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
6959    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
6960    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
6961    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
6962    /// dispatch sites cannot drift apart again.
6963    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
6964    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
6965    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
6966    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
6967    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
6968    /// eligibility sites so they cannot drift (the qwen35_serving_class lesson).
6969    fn mtp_graph_capturable(&self) -> bool {
6970        self.mtp
6971            .as_ref()
6972            .map(|m| match &m.ffn {
6973                crate::hybrid::Ffn::Dense { .. } => true,
6974                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
6975            })
6976            .unwrap_or(false)
6977    }
6978
6979    fn qwen35_serving_class(&self) -> bool {
6980        matches!(
6981            self.cfg.arch,
6982            memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
6983        )
6984    }
6985
6986    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
6987    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
6988    /// session already exist.
6989    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
6990        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
6991            || !spec_devacc()
6992            || spec_replay_env_enabled()
6993            || spec_stream()
6994            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
6995            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
6996            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
6997            || std::env::var("MEMRA_SPEC_PMIN")
6998                .ok()
6999                .and_then(|v| v.parse::<f32>().ok())
7000                .unwrap_or(0.0)
7001                > 0.0
7002            || self.is_gemma4_e4b()
7003            || self.cfg.gemma4.is_some()
7004            || self.mtp.is_none()
7005        {
7006            return false;
7007        }
7008        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
7009            return false;
7010        };
7011        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
7012            return false;
7013        }
7014        crate::pp::PpNRt::get(e)
7015            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
7016            .unwrap_or(false)
7017    }
7018
7019    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
7020    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
7021    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
7022    #[allow(clippy::too_many_arguments)]
7023    pub fn generate_spec_session_pair(
7024        &self,
7025        e: &Engine,
7026        sess_a: &mut SpecSession,
7027        max_new_a: usize,
7028        k_a: usize,
7029        sess_b: &mut SpecSession,
7030        max_new_b: usize,
7031        k_b: usize,
7032    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
7033    {
7034        if !self.spec_pipe_available(e) {
7035            return Err("two-session speculative pipeline is outside its reduced matrix".into());
7036        }
7037        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
7038            return Err(
7039                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
7040            );
7041        }
7042        for sess in [&*sess_a, &*sess_b] {
7043            if sess.committed.is_empty()
7044                || sess.last_h.is_none()
7045                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
7046            {
7047                return Err("two-session speculative pipeline requires warm continuations".into());
7048            }
7049        }
7050
7051        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7052            && !spec_host_embd()
7053            && self.mtp_graph_capturable()
7054            && !crate::model::full_prec_enabled();
7055        let graph_a = graph_ok && k_a + 2 < 96;
7056        let graph_b = graph_ok && k_b + 2 < 96;
7057        let was_tracking = e.ctx().is_event_tracking();
7058        if (graph_a || graph_b) && was_tracking {
7059            unsafe {
7060                e.ctx().disable_event_tracking();
7061            }
7062        }
7063
7064        static LOGGED: std::sync::Once = std::sync::Once::new();
7065        LOGGED.call_once(|| {
7066            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
7067        });
7068        let sync = std::sync::Arc::new(SpecPipeSync::new());
7069        let lane_a = SpecPipeLane {
7070            sync: sync.clone(),
7071            lane: 0,
7072        };
7073        let lane_b = SpecPipeLane { sync, lane: 1 };
7074        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
7075        let (result_a, result_b) = std::thread::scope(|scope| {
7076            let b = scope.spawn(move || {
7077                let mut finish = SpecPipeFinish::new(&lane_b);
7078                let sess_b = unsafe { sess_b_ptr.get_mut() };
7079                let result = e
7080                    .ctx()
7081                    .bind_to_thread()
7082                    .map_err(|err| err.to_string())
7083                    .and_then(|_| {
7084                        self.generate_spec_inner2(
7085                            e,
7086                            &[],
7087                            max_new_b,
7088                            k_b,
7089                            graph_b,
7090                            Some(sess_b),
7091                            None,
7092                            None,
7093                            None,
7094                            None,
7095                            Some(&lane_b),
7096                        )
7097                        .map_err(|err| err.to_string())
7098                    });
7099                finish.close(result.is_err());
7100                result
7101            });
7102            let mut finish = SpecPipeFinish::new(&lane_a);
7103            let result_a = self.generate_spec_inner2(
7104                e,
7105                &[],
7106                max_new_a,
7107                k_a,
7108                graph_a,
7109                Some(sess_a),
7110                None,
7111                None,
7112                None,
7113                None,
7114                Some(&lane_a),
7115            );
7116            finish.close(result_a.is_err());
7117            let result_b = b
7118                .join()
7119                .map_err(|_| "paired speculative session B panicked".to_string())
7120                .and_then(|r| r);
7121            (result_a, result_b)
7122        });
7123
7124        if (graph_a || graph_b) && was_tracking {
7125            unsafe {
7126                e.ctx().enable_event_tracking();
7127            }
7128        }
7129        let result_a = result_a?;
7130        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
7131        Ok((result_a, result_b))
7132    }
7133
7134    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
7135    /// message rendered through the chat template continuation). Returns (new tokens emitted,
7136    /// drafted, accepted); session.committed grows by suffix + emitted.
7137    pub fn generate_spec_session(
7138        &self,
7139        e: &Engine,
7140        sess: &mut SpecSession,
7141        suffix: &[u32],
7142        max_new: usize,
7143        k: usize,
7144    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7145        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
7146    }
7147
7148    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
7149    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
7150    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
7151    /// for the filtered target (feat/filtered-spec).
7152    ///
7153    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
7154    /// output — once right after the prime's first token, then once per round commit — so a
7155    /// streaming caller can flush text at round cadence instead of once per burst. The slices
7156    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
7157    /// timing only: token bytes, session state, and exactness are untouched.
7158    ///
7159    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
7160    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
7161    /// the caller's scheduler regains control without waiting the burst out. Burst size is
7162    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
7163    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
7164    /// drains and the defensive tail flush can land with nothing new committed).
7165    #[allow(clippy::too_many_arguments)]
7166    pub fn generate_spec_session_sampled(
7167        &self,
7168        e: &Engine,
7169        sess: &mut SpecSession,
7170        suffix: &[u32],
7171        max_new: usize,
7172        k: usize,
7173        sampling: Option<SpecSampling>,
7174        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7175    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7176        self.generate_spec_session_sampled_prime_split(
7177            e, sess, suffix, max_new, k, sampling, None, on_commit,
7178        )
7179    }
7180
7181    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
7182    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
7183    /// pass `None` and stay on the existing zero-prime path.
7184    #[allow(clippy::too_many_arguments)]
7185    pub fn generate_spec_session_sampled_prime_split(
7186        &self,
7187        e: &Engine,
7188        sess: &mut SpecSession,
7189        suffix: &[u32],
7190        max_new: usize,
7191        k: usize,
7192        sampling: Option<SpecSampling>,
7193        prime_split: Option<usize>,
7194        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7195    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7196        self.generate_spec_session_constrained_prime_split(
7197            e,
7198            sess,
7199            suffix,
7200            max_new,
7201            k,
7202            sampling,
7203            None,
7204            prime_split,
7205            on_commit,
7206        )
7207    }
7208
7209    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
7210    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
7211    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
7212    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
7213    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
7214    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
7215    /// may drop (drafter is unconstrained); that is measured, not hidden.
7216    #[allow(clippy::too_many_arguments)]
7217    pub fn generate_spec_session_constrained(
7218        &self,
7219        e: &Engine,
7220        sess: &mut SpecSession,
7221        suffix: &[u32],
7222        max_new: usize,
7223        k: usize,
7224        sampling: Option<SpecSampling>,
7225        constraint: Option<&mut dyn SpecConstraint>,
7226        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7227    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7228        self.generate_spec_session_constrained_prime_split(
7229            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
7230        )
7231    }
7232
7233    #[allow(clippy::too_many_arguments)]
7234    pub fn generate_spec_session_constrained_prime_split(
7235        &self,
7236        e: &Engine,
7237        sess: &mut SpecSession,
7238        suffix: &[u32],
7239        max_new: usize,
7240        k: usize,
7241        sampling: Option<SpecSampling>,
7242        constraint: Option<&mut dyn SpecConstraint>,
7243        prime_split: Option<usize>,
7244        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7245    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7246        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
7247            return Err(
7248                "constrained spec decode is greedy-only (worker routes sampled \
7249                        constrained to plain decode)"
7250                    .into(),
7251            );
7252        }
7253        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
7254        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
7255        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
7256        // serve continuation case — consume the carry in-loop with zero solo passes.
7257        if sess.pending_tok.is_some()
7258            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
7259        {
7260            self.spec_flush_pending(e, sess, sampling)?;
7261        }
7262
7263        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
7264        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
7265        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
7266        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7267            && !spec_host_embd()
7268            && self.mtp_graph_capturable()
7269            && k + 2 < 96
7270            && !crate::model::full_prec_enabled();
7271        let was_tracking = e.ctx().is_event_tracking();
7272        if graph_draft && was_tracking {
7273            unsafe {
7274                e.ctx().disable_event_tracking();
7275            }
7276        }
7277        let r = self.generate_spec_inner2(
7278            e,
7279            suffix,
7280            max_new,
7281            k,
7282            graph_draft,
7283            Some(sess),
7284            sampling,
7285            constraint,
7286            on_commit,
7287            prime_split,
7288            None,
7289        );
7290        if graph_draft && was_tracking {
7291            unsafe {
7292                e.ctx().enable_event_tracking();
7293            }
7294        }
7295        let (out, d, a) = r?;
7296        Ok((out, d, a))
7297    }
7298
7299    pub fn generate_spec(
7300        &self,
7301        e: &Engine,
7302        prompt: &[u32],
7303        max_new: usize,
7304        k: usize,
7305    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7306        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
7307        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
7308        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7309            && !spec_host_embd()
7310            && self.mtp_graph_capturable()
7311            && k + 2 < 96
7312            && !crate::model::full_prec_enabled();
7313        if !graph_draft {
7314            return self.generate_spec_inner2(
7315                e, prompt, max_new, k, false, None, None, None, None, None, None,
7316            );
7317        }
7318        let was_tracking = e.ctx().is_event_tracking();
7319        if was_tracking {
7320            unsafe {
7321                e.ctx().disable_event_tracking();
7322            }
7323        }
7324        let r = self.generate_spec_inner2(
7325            e, prompt, max_new, k, true, None, None, None, None, None, None,
7326        );
7327        if was_tracking {
7328            unsafe {
7329                e.ctx().enable_event_tracking();
7330            }
7331        }
7332        r
7333    }
7334
7335    fn generate_spec_inner2(
7336        &self,
7337        e: &Engine,
7338        prompt: &[u32],
7339        max_new: usize,
7340        k: usize,
7341        graph_draft: bool,
7342        mut sess: Option<&mut SpecSession>,
7343        sampling: Option<SpecSampling>,
7344        mut constraint: Option<&mut dyn SpecConstraint>,
7345        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7346        prime_split: Option<usize>,
7347        pipe: Option<&SpecPipeLane>,
7348    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7349        assert!(k >= 1, "k must be >= 1");
7350        if let Some(p) = pipe {
7351            p.setup_begin()?;
7352        }
7353        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
7354        let mut flushed = 0usize;
7355        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
7356        // at the next round boundary (same exit as max_new reached — the session tail runs).
7357        // Initialized by the unconditional post-prime flush below.
7358        let mut keep_going;
7359        let mtp = self
7360            .mtp
7361            .as_ref()
7362            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
7363        let n_vocab = self.output.out_features();
7364        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
7365        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
7366        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
7367        let d_vocab = mtp
7368            .shared_head_head
7369            .as_ref()
7370            .unwrap_or(&self.output)
7371            .out_features();
7372        let n_embd = self.cfg.n_embd as usize;
7373        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
7374        // already committed (their state is in the caches); 0 = fresh single-shot call.
7375        let session_mode = sess.is_some();
7376        let max_ctx = match sess.as_ref() {
7377            Some(s) => s.cache.max_ctx,
7378            None => prompt.len() + max_new + k + 8,
7379        };
7380        let mut own_cache;
7381        let mut own_scratch;
7382        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
7383        // (requested split, destination slot). Single-shot per burst; fresh calls have none.
7384        let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
7385        let (
7386            cache,
7387            scratch,
7388            mut sess_tail,
7389            mut sess_draft_slot,
7390            mut sess_pending_slot,
7391            sess_ckpt_slot,
7392            sess_telem,
7393        ): (
7394            &mut Cache,
7395            &mut MtpScratch,
7396            Option<(
7397                &mut Vec<u32>,
7398                &mut Option<CudaSlice<f32>>,
7399                &mut Option<u32>,
7400                &mut u32,
7401                &mut u32,
7402            )>,
7403            Option<&mut Option<DraftGraphCtx>>,
7404            Option<&mut Option<u32>>,
7405            Option<&mut Option<SpecCheckpoint>>,
7406            Option<&SpecTelemetryCounters>,
7407        ) = match sess.take() {
7408            Some(sr) => {
7409                let SpecSession {
7410                    cache,
7411                    scratch,
7412                    committed,
7413                    last_h,
7414                    next_pred,
7415                    sctr: s_sctr,
7416                    uctr: s_uctr,
7417                    draft_ctx,
7418                    pending_tok,
7419                    turn_ckpt,
7420                    telem,
7421                    capture_at,
7422                    boundary_capture,
7423                } = sr;
7424                sess_capture = Some((capture_at.take(), boundary_capture));
7425                (
7426                    cache,
7427                    scratch,
7428                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
7429                    Some(draft_ctx),
7430                    Some(pending_tok),
7431                    Some(turn_ckpt),
7432                    Some(telem),
7433                )
7434            }
7435            None => {
7436                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
7437                // `Cache::new` verbatim.
7438                own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
7439                // Persistent scratch = max_ctx rows (~2KB/token quantized).
7440                own_scratch = MtpScratch::new(
7441                    e,
7442                    &self.cfg,
7443                    max_ctx,
7444                    self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7445                )?;
7446                (
7447                    &mut own_cache,
7448                    &mut own_scratch,
7449                    None,
7450                    None,
7451                    None,
7452                    None,
7453                    None,
7454                )
7455            }
7456        };
7457        let base = cache.pos;
7458        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
7459        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
7460        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
7461        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
7462        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
7463        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
7464        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
7465        // acceptance-only — exactness is verify's job either way).
7466        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
7467        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
7468        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
7469        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
7470        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
7471        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
7472        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
7473        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
7474        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
7475        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
7476        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
7477        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
7478        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
7479        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
7480        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
7481        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
7482        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
7483        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
7484        // + fallback seam).
7485        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
7486        // bar — the retained verify-state commit proven equivalent to sequential serving —
7487        // was waiting on this arch running the serving batched verify class, which the
7488        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
7489        // replay-free commit consumes is now produced by the SAME serving-class verify that
7490        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
7491        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
7492        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
7493        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
7494        // rollback + A/B seam.
7495        let spec_replay = spec_replay_env_enabled();
7496        if constraint.is_some() && spec_replay {
7497            return Err(
7498                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
7499                        (legacy replay commits an unmasked bonus)"
7500                    .into(),
7501            );
7502        }
7503        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
7504        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
7505        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
7506        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
7507
7508        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
7509        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
7510        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
7511        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
7512        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
7513        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
7514        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
7515        // generation exactly where the last turn stopped — no prime at all. The stashed
7516        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
7517        // committed.last() by the same rule this entry applies to a cold prime's last row —
7518        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
7519        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
7520        // where the sampler and the session's Philox counters were live). `last_h` seeds the
7521        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
7522        let continuation = prompt.is_empty();
7523        if continuation {
7524            assert!(session_mode, "empty prompt requires a session");
7525            assert!(
7526                sess_tail
7527                    .as_ref()
7528                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
7529                        && lh.is_some()
7530                        && (np.is_some() || carried_pending.is_some())),
7531                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
7532            );
7533        }
7534        let mut prime_logits;
7535        let mut prompt_h: Option<CudaSlice<f32>> = None;
7536        let t_prime = std::time::Instant::now();
7537        let batched_prime = !continuation
7538            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
7539            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7540            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
7541        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
7542        if prime_split.is_some() && (continuation || base != 0) {
7543            return Err("spec prime split is cold-session-only".into());
7544        }
7545        if continuation {
7546            prime_logits = Vec::new();
7547        } else if let Some(split) = prime_split {
7548            if split < crate::hybrid_forward::PRIME_MIN_T {
7549                return Err(format!(
7550                    "spec prime split {split} is below PRIME_MIN_T {}",
7551                    crate::hybrid_forward::PRIME_MIN_T,
7552                )
7553                .into());
7554            }
7555            // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
7556            // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
7557            // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
7558            // Retain every hidden row so the draft scratch fill remains one coherent prompt.
7559            let mut h_all = e.uninit(prompt.len() * n_embd)?;
7560            let (l, _, h_prefix) =
7561                self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
7562            e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
7563            prime_logits = l;
7564            // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
7565            // are about to be advanced in place by the tail prime, so this is the ONLY moment
7566            // the boundary's recurrent state exists. Capture iff the worker requested exactly
7567            // this split. cache.pos == split here (the prefix prime just finished). A failed
7568            // snapshot is silent (turn_ckpt convention) — publication is an optimization,
7569            // never a correctness dependency.
7570            if let Some((requested, slot)) = sess_capture.as_mut() {
7571                if *requested == Some(split) {
7572                    debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
7573                    if let Ok(snap) = cache.snapshot(e) {
7574                        **slot = Some(SpecBoundaryCapture {
7575                            snap,
7576                            pos: split,
7577                            logits: prime_logits.clone(),
7578                            // rows [0..split) of h_all are the prefix prime's hiddens — copied
7579                            // just above, before the tail prime overwrites nothing (append-only).
7580                            last_h: capture_boundary_hidden(e, &h_all, split, n_embd),
7581                        });
7582                    }
7583                }
7584            }
7585            let tail = &prompt[split..];
7586            if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
7587                && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7588                && !e.frozen_cpu_experts_prefer_tokenwise_prime()
7589            {
7590                let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
7591                e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
7592                prime_logits = l;
7593            } else {
7594                for (i, &tok) in tail.iter().enumerate() {
7595                    let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
7596                    e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
7597                    prime_logits = l;
7598                }
7599            }
7600            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7601                eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
7602            }
7603            prompt_h = Some(h_all);
7604        } else if batched_prime {
7605            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
7606            prime_logits = l;
7607            prompt_h = Some(hiddens);
7608        } else {
7609            prime_logits = Vec::new();
7610            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
7611            for (i, &tok) in prompt.iter().enumerate() {
7612                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
7613                if let Some(ph) = prompt_h.as_mut() {
7614                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
7615                }
7616                prime_logits = l;
7617            }
7618        }
7619        e.stream().synchronize()?;
7620        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
7621        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
7622        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
7623        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
7624        // prime_split. The mid-prompt capture above already consumed the request if it matched.
7625        if !continuation && base == 0 {
7626            if let Some((requested, slot)) = sess_capture.as_mut() {
7627                if *requested == Some(prompt.len()) && slot.is_none() {
7628                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
7629                    if let Ok(snap) = cache.snapshot(e) {
7630                        **slot = Some(SpecBoundaryCapture {
7631                            snap,
7632                            pos: prompt.len(),
7633                            logits: prime_logits.clone(),
7634                            last_h: prompt_h
7635                                .as_ref()
7636                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
7637                                .unwrap_or_default(),
7638                        });
7639                    }
7640                }
7641            }
7642        }
7643        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
7644        // prime-subtraction hack.
7645        crate::PRIME_NANOS.store(
7646            t_prime.elapsed().as_nanos() as u64,
7647            std::sync::atomic::Ordering::Relaxed,
7648        );
7649
7650        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7651        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
7652        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
7653        let host_embd = spec_host_embd();
7654        let embd_gpu = if host_embd {
7655            None
7656        } else {
7657            Some(
7658                self.embd_gpu
7659                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7660            )
7661        };
7662        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7663        if host_embd {
7664            eprintln!(
7665                "[spec] host-row embedding: {} bytes kept off HBM",
7666                self.embd.raw.len()
7667            );
7668        }
7669        let mut out: Vec<u32> = Vec::with_capacity(max_new);
7670        let mut total_drafted = 0usize;
7671        let mut total_accepted = 0usize;
7672
7673        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
7674        // The sampler config, the session's Philox counters and the penalty window are parsed
7675        // HERE, above the boundary-token selection, because the boundary token must be drawn
7676        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
7677        // selection, which is the whole mechanical reason the boundary token was an argmax:
7678        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
7679        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
7680        // below takes the argmax path it always took).
7681        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
7682        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
7683        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
7684        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
7685        let sp = sampling.unwrap_or_else(|| SpecSampling {
7686            temp: std::env::var("MEMRA_SPEC_TEMP")
7687                .ok()
7688                .and_then(|v| v.parse().ok())
7689                .unwrap_or(0.0),
7690            seed: std::env::var("MEMRA_SEED")
7691                .ok()
7692                .and_then(|v| v.parse().ok())
7693                .unwrap_or(42),
7694            top_k: std::env::var("MEMRA_TOP_K")
7695                .ok()
7696                .and_then(|v| v.parse().ok())
7697                .unwrap_or(0),
7698            top_p: std::env::var("MEMRA_TOP_P")
7699                .ok()
7700                .and_then(|v| v.parse().ok())
7701                .unwrap_or(1.0),
7702            min_p: std::env::var("MEMRA_MIN_P")
7703                .ok()
7704                .and_then(|v| v.parse().ok())
7705                .unwrap_or(0.0),
7706            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
7707                .ok()
7708                .and_then(|v| v.parse().ok())
7709                .unwrap_or(0),
7710            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
7711                .ok()
7712                .and_then(|v| v.parse().ok())
7713                .unwrap_or(1.0),
7714            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
7715                .ok()
7716                .and_then(|v| v.parse().ok())
7717                .unwrap_or(0.0),
7718            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
7719                .ok()
7720                .and_then(|v| v.parse().ok())
7721                .unwrap_or(0.0),
7722        });
7723        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
7724        let sampled = sp_temp > 0.0;
7725        // Counters resume from the session (burst continuity: randomness must never repeat
7726        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
7727        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
7728        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
7729        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
7730        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
7731        // for the penalized+filtered target). History = generated tokens, host-tracked window.
7732        let pen_on = sampled
7733            && sp.penalty_last_n > 0
7734            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
7735        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
7736        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
7737        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
7738        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
7739        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
7740        // which is what the API contract says and what the plain sampler's own `history` does.
7741        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
7742        let mut pen_hist: Vec<u32> = if pen_on {
7743            let sess_hist: &[u32] = if spec_pen_session_on() {
7744                sess_tail
7745                    .as_ref()
7746                    .map(|(c, ..)| c.as_slice())
7747                    .unwrap_or(&[])
7748            } else {
7749                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
7750            };
7751            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
7752        } else {
7753            Vec::new()
7754        };
7755        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
7756        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
7757        // request's own filtered/penalized target through the session's Philox stream
7758        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
7759        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
7760        // Emit it, then FEED it to establish the loop invariant below.
7761        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
7762        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
7763        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
7764        // prompt's last logits (plain constrained-greedy identity); a continuation without
7765        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
7766        // worker never resumes constrained sessions from the pool, so this cannot fire).
7767        if let Some(c) = constraint.as_deref_mut() {
7768            if continuation && carried_pending.is_none() {
7769                return Err("constrained spec continuation requires a carried pending \
7770                            (pool resume is unconstrained-only)"
7771                    .into());
7772            }
7773            if !continuation {
7774                c.mask_logits(&mut prime_logits)
7775                    .map_err(|e2| format!("constraint: {e2}"))?;
7776            }
7777        }
7778        let mut last_token = if let Some(b) = carried_pending {
7779            b
7780        } else if continuation {
7781            // A continuation's boundary token was DRAWN by the burst that stashed it (the
7782            // session tail below), or by `spec_session_from_restored` for a converted
7783            // prefix-cache hit — in both cases from the correct logits row with this same
7784            // session's Philox stream, which is why it can be consumed here as-is.
7785            sess_tail.as_ref().unwrap().2.unwrap()
7786        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
7787            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
7788        } else {
7789            // greedy (byte contract), the rollback door, or constrained (masked-argmax
7790            // identity — the worker routes sampled+constrained to the plain path, and this
7791            // function refuses the combination outright above).
7792            argmax(&prime_logits) as u32
7793        };
7794        if pen_on {
7795            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
7796            // emitted token into its penalty history, and pre-lane the burst's first token
7797            // was invisible to penalties forever (never pushed, and never in `committed`
7798            // until this burst's tail). Covers the carry/continuation seeds too — neither is
7799            // in `committed` yet.
7800            pen_hist.push(last_token);
7801        }
7802        if carried_pending.is_none() {
7803            out.push(last_token);
7804            // grammar advances with every emitted token (carried pendings were consumed
7805            // by the burst that emitted them).
7806            if let Some(c) = constraint.as_deref_mut() {
7807                c.consume(last_token)
7808                    .map_err(|e2| format!("constraint: {e2}"))?;
7809            }
7810        }
7811        if continuation {
7812            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
7813            // overhang so the chain's first append lands at slot base (== committed.len()).
7814            scratch.set_len(e, base)?;
7815        }
7816        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
7817        // concatenating to the full `out`). Called after the prime's first token and after each
7818        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
7819        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
7820        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
7821        fn flush_commit(
7822            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
7823            out: &[u32],
7824            flushed: &mut usize,
7825        ) -> bool {
7826            if let Some(f) = cb.as_mut() {
7827                let keep = f(&out[*flushed..]);
7828                *flushed = out.len();
7829                keep
7830            } else {
7831                true
7832            }
7833        }
7834        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
7835        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
7836        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
7837        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
7838        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
7839        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
7840        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
7841        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
7842        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
7843        // those, so their residual mass is p(x), correct by construction).
7844        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
7845            match &mtp.d2t {
7846                Some(map) => Some(e.htod_u32_v(map)?),
7847                None => None,
7848            }
7849        } else {
7850            None
7851        };
7852        let mut q_full_buf: Option<CudaSlice<f32>> = None;
7853        // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
7854        let host_u01 = |seed: u64, ctr: u32| -> f32 {
7855            let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
7856            let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
7857            let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
7858            for _ in 0..10 {
7859                let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
7860                let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
7861                let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
7862                c0 = n0;
7863                c1 = n1;
7864                c2 = n2;
7865                c3 = n3;
7866                k0 = k0.wrapping_add(0x9E3779B9);
7867                k1 = k1.wrapping_add(0xBB67AE85);
7868            }
7869            (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
7870        };
7871        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
7872        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
7873        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
7874        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
7875        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
7876        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
7877        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
7878        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
7879        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
7880        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
7881        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
7882        let t_ent = std::time::Instant::now();
7883
7884        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
7885        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
7886        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
7887        // the one that matters (a history-rewriting client mutates what the session GENERATED,
7888        // so the next turn's prompt agrees with this one up to exactly here).
7889        //
7890        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
7891        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
7892        // hold exactly `base + prompt.len()` rows and nothing generated.
7893        //
7894        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
7895        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
7896        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
7897        // `<think>` block the client strips, so every later turn's diff diverged exactly one
7898        // token below the checkpoint and affinity declined 100% of the time. Measured on the
7899        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
7900        // whole mechanism inert while looking, from the outside, like a working
7901        // correctness-declines-safely path — hence the decline log carries the offsets.
7902        //
7903        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
7904        // state (the reason a spec session could not rewind before). The draft scratch needs no
7905        // copy: rows below the boundary are rewritten by the next turn's own fill.
7906        //
7907        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
7908        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
7909        // checkpoint rather than replacing it with a strictly worse one.
7910        //
7911        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
7912        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
7913        // fail the burst that is already running — so the error is swallowed, loud only under
7914        // MEMRA_DEBUG_SPEC.
7915        if let Some(slot) = sess_ckpt_slot {
7916            if !continuation {
7917                let pos = cache.pos;
7918                debug_assert_eq!(
7919                    pos,
7920                    base + prompt.len(),
7921                    "turn checkpoint must sit at the prompt end, before the init feed"
7922                );
7923                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
7924                    if let Some(ph) = &prompt_h {
7925                        // hidden of the LAST primed row = the predecessor anchor at this
7926                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
7927                        // last_h, and what the next prime's fill reads for its first row).
7928                        let np = prompt.len();
7929                        e.uninit(n_embd).and_then(|mut a| {
7930                            e.copy_view_into(
7931                                &mut a,
7932                                0,
7933                                &ph.slice((np - 1) * n_embd..np * n_embd),
7934                                n_embd,
7935                            )?;
7936                            Ok(a)
7937                        })
7938                    } else {
7939                        Err("no prompt hiddens".into())
7940                    };
7941                match (cache.snapshot(e), anchor) {
7942                    (Ok(snap), Ok(last_h)) => {
7943                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
7944                    }
7945                    (s, a) => {
7946                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
7947                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
7948                            let err = s
7949                                .err()
7950                                .map(|e| e.to_string())
7951                                .or_else(|| a.err().map(|e| e.to_string()))
7952                                .unwrap_or_default();
7953                            eprintln!(
7954                                "[spec] turn checkpoint skipped ({err}); \
7955                                       next turn re-primes in full"
7956                            );
7957                        }
7958                    }
7959                }
7960            }
7961        }
7962        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
7963        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
7964        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
7965        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
7966        let mut last_pred = 0u32;
7967        let mut last_col_logits: Option<CudaSlice<f32>> = None;
7968        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
7969        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
7970        let mut init_logits_host: Option<Vec<f32>> = None;
7971        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
7972            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
7973            last_pred = argmax(&init_logits) as u32;
7974            if constraint.is_some() {
7975                init_logits_host = Some(init_logits.clone());
7976            }
7977            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
7978            if sampled {
7979                last_col_logits = Some(e.htod(&init_logits)?);
7980            }
7981            h
7982        } else {
7983            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
7984            let lh = sess_tail
7985                .as_ref()
7986                .unwrap()
7987                .1
7988                .as_ref()
7989                .expect("pending carry requires last_h");
7990            e.clone_dtod(lh)?
7991        };
7992        let t_init = t_ent.elapsed();
7993        let mut last_col_stats: Option<(f32, f32, f32)> = None;
7994        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
7995        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
7996        // stable pointer for the graph-draft round-start copy.
7997        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
7998        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
7999        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
8000        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
8001        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
8002        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
8003        // overwritten below).
8004        let mut fill_prev = e.clone_dtod(&h_seed0)?;
8005        {
8006            if let Some(ph) = &prompt_h {
8007                let np = prompt.len();
8008                e.copy_view_into(
8009                    &mut h_seed_buf,
8010                    0,
8011                    &ph.slice((np - 1) * n_embd..np * n_embd),
8012                    n_embd,
8013                )?;
8014            } else if continuation {
8015                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8016                    if let Some(lh) = lh.as_ref() {
8017                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
8018                    }
8019                }
8020            }
8021        }
8022        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
8023        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
8024
8025        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
8026        let fork_mode = OptiForkGateMode::configured();
8027        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
8028        // the end. Metric normalization vs the reference engine: BOTH engines count
8029        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
8030        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
8031        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
8032        let mut st_drafted = vec![0usize; k];
8033        let mut st_accepted = vec![0usize; k];
8034        let mut st_len_hist = vec![0usize; k + 1];
8035        let mut st_full = 0usize;
8036        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
8037        // stop the draft chain early when the head's softmax confidence in its own pick drops
8038        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
8039        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
8040        let p_min = *PMIN.get_or_init(|| {
8041            std::env::var("MEMRA_SPEC_PMIN")
8042                .ok()
8043                .and_then(|v| v.parse().ok())
8044                .unwrap_or(0.0)
8045        });
8046        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
8047        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
8048        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
8049        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
8050        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
8051        // verify batch is not); the j==0 exemption stays for pending-less rounds.
8052        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
8053            .map(|v| v == "1")
8054            .unwrap_or(false);
8055
8056        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
8057        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
8058        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
8059        // cuBLAS path in an exotic head) falls back to the eager draft chain.
8060        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
8061        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
8062        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
8063        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
8064        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
8065        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
8066        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
8067        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
8068        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
8069            Some(c) => c,
8070            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
8071        };
8072        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
8073        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
8074        if sampled && dctx.g_q.len() < d_vocab {
8075            dctx.g_q = e.zeros(d_vocab)?;
8076            dctx.g_perturb = e.zeros(d_vocab)?;
8077        }
8078        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
8079        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
8080        // truncation (the correctness backstop) stops cutting every tight-schema round.
8081        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
8082        // shape, so a parked graph of the other shape is dropped and recaptured.
8083        let dmask_on = constraint
8084            .as_deref()
8085            .is_some_and(|c| c.draft_mask_enabled());
8086        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
8087        if dmask_on && dctx.g_dmask.len() < dmask_words {
8088            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
8089            dctx.graph = None; // the old capture baked the old (or no) mask pointer
8090            dctx.failed.clear_greedy();
8091            dctx.keeper.clear();
8092        }
8093        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
8094            dctx.graph = None;
8095            dctx.failed.clear_greedy();
8096            dctx.keeper.clear();
8097        }
8098        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
8099            let DraftGraphCtx {
8100                g_tok,
8101                g_pos,
8102                g_seed,
8103                g_p,
8104                g_dmask,
8105                ..
8106            } = &mut dctx;
8107            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
8108            // host uploads the position's real words, so the warmups stay grammar-free.
8109            if dmask_on {
8110                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
8111            }
8112            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
8113            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
8114            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
8115            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
8116            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
8117            // passes (and, in serve, other sessions) recycle those addresses and the replay then
8118            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
8119            let cap_res = e.capture_graph_retained(|e| {
8120                self.mtp_head_forward_cap(
8121                    e,
8122                    mtp,
8123                    g_tok,
8124                    g_pos,
8125                    g_seed,
8126                    g_p,
8127                    &mut *scratch,
8128                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
8129                    true,
8130                    embd_gpu.expect("graph draft requires resident embedding"),
8131                    embd_qt,
8132                    embd_rb,
8133                    d_vocab,
8134                    None,
8135                    None,
8136                    if dmask_on {
8137                        Some((g_dmask_ro, dmask_words))
8138                    } else {
8139                        None
8140                    },
8141                )
8142            });
8143            match cap_res {
8144                Ok((g, keep)) => {
8145                    scratch.set_len(e, base)?;
8146                    dctx.graph = Some(g);
8147                    dctx.graph_masked = dmask_on;
8148                    dctx.keeper = keep;
8149                }
8150                Err(err) => {
8151                    scratch.set_len(e, base)?;
8152                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
8153                    // silent. Once per flip — mark returns None on an already-failed ctx.
8154                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
8155                        eprintln!("{line}");
8156                    }
8157                }
8158            }
8159        }
8160        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
8161        // graph object, built only when sampled && graph-eligible — the greedy capture above is
8162        // untouched (and skipped when sampled: its graph would never be launched). Same head
8163        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
8164        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
8165        // once per round); the raw head logits land in the persistent g_q for the host's
8166        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
8167        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
8168        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
8169        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
8170        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
8171        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
8172        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
8173        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
8174        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
8175        // this compare misses at most ONCE per resumed request — the first burst recaptures
8176        // and every later burst in that request replays. A client that wants the parked graph
8177        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
8178        // stable across its whole conversation.
8179        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
8180        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
8181        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
8182        // force the eager draft (which computes stats/penalties per row).
8183        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
8184        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
8185        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
8186        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
8187        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
8188        // the request shape the vendor-default flip makes the majority).
8189        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
8190        let pure_temp = s_key.pure_temp();
8191        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
8192            dctx.graph_s = None;
8193            dctx.failed.clear_sampled();
8194            dctx.s_key = None;
8195            dctx.q_slots.clear();
8196            dctx.keeper_s.clear();
8197        }
8198        if graph_draft
8199            && sampled
8200            && pure_temp
8201            && dctx.graph_s.is_none()
8202            && !dctx.failed.sampled_failed()
8203        {
8204            let DraftGraphCtx {
8205                g_tok,
8206                g_pos,
8207                g_seed,
8208                g_p,
8209                g_ctr,
8210                g_perturb,
8211                g_q,
8212                ..
8213            } = &mut dctx;
8214            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
8215            let cap_res = e.capture_graph_retained(|e| {
8216                self.mtp_head_forward_cap(
8217                    e,
8218                    mtp,
8219                    g_tok,
8220                    g_pos,
8221                    g_seed,
8222                    g_p,
8223                    &mut *scratch,
8224                    p_min > 0.0,
8225                    true,
8226                    embd_gpu.expect("graph draft requires resident embedding"),
8227                    embd_qt,
8228                    embd_rb,
8229                    d_vocab,
8230                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
8231                    None,
8232                    None, // constrained spec is greedy-only — sampled never carries a hook
8233                )
8234            });
8235            match cap_res {
8236                Ok((g, keep)) => {
8237                    scratch.set_len(e, base)?;
8238                    for _ in 0..k {
8239                        dctx.q_slots.push(e.zeros(d_vocab)?);
8240                    }
8241                    dctx.graph_s = Some(g);
8242                    dctx.s_key = Some(s_key);
8243                    dctx.keeper_s = keep;
8244                }
8245                Err(err) => {
8246                    scratch.set_len(e, base)?;
8247                    // LOUD flip (audit Q2): same contract as the greedy capture above.
8248                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
8249                        eprintln!("{line}");
8250                    }
8251                }
8252            }
8253        }
8254        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
8255        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
8256        // captured under this request's exact regime, and capture requires `pure_temp` — so a
8257        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
8258        // the graph arm, so it is asserted here rather than assumed: a future change that widens
8259        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
8260        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
8261        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
8262        // rather than launching it; the launch site re-tests `pure_temp` independently.
8263        if sampled && !pure_temp && dctx.graph_s.is_some() {
8264            debug_assert!(
8265                false,
8266                "sampled draft graph parked under {:?} survived into a FILTERED request \
8267                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
8268                 softmax, so the verify's filtered q would test a distribution the draft was \
8269                 never sampled from",
8270                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
8271            );
8272            eprintln!(
8273                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
8274                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
8275                 EAGER — the key must carry every field that shapes q",
8276                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
8277            );
8278            dctx.graph_s = None;
8279            dctx.s_key = None;
8280            dctx.q_slots.clear();
8281            dctx.keeper_s.clear();
8282        }
8283        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
8284        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
8285        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
8286        // arms below print which chain actually ran, so the probe never restates the condition.
8287        if skey_probe() {
8288            eprintln!(
8289                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
8290                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
8291                sampled as u8,
8292                pure_temp as u8,
8293                sp_temp,
8294                sp.top_k,
8295                sp.top_p,
8296                sp.min_p,
8297                pen_on as u8,
8298                k,
8299                graph_draft as u8,
8300                dctx.graph_s.is_some() as u8,
8301                dctx.s_key,
8302            );
8303        }
8304        let t_cap = t_ent.elapsed();
8305        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
8306        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
8307        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
8308        // fill: the first chain step processes it and appends its entry at slot prompt.len().
8309        if let Some(ph) = &prompt_h {
8310            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
8311            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
8312            // global positions [base..base+tp). Fresh call: base==0, identical to before.
8313            scratch.set_len(e, base)?;
8314            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
8315            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
8316            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
8317            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
8318            let tp = prompt.len();
8319            let fill_chunk: usize = if crate::cache::swa_ring_on() {
8320                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
8321            } else {
8322                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
8323                // meaning one monolithic fill.
8324                std::env::var("MEMRA_PRIME_CHUNK")
8325                    .ok()
8326                    .and_then(|v| v.parse().ok())
8327                    .unwrap_or(4096)
8328            };
8329            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
8330            let mut start = 0usize;
8331            while start < tp {
8332                let end = (start + fill_chunk).min(tp);
8333                let tc = end - start;
8334                {
8335                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
8336                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
8337                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
8338                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
8339                    let mut phs = e.zeros(tc * n_embd)?;
8340                    let (src_lo, dst_off) = if start == 0 {
8341                        (0, n_embd)
8342                    } else {
8343                        ((start - 1) * n_embd, 0)
8344                    };
8345                    let n_copy = if start == 0 {
8346                        (tc - 1) * n_embd
8347                    } else {
8348                        tc * n_embd
8349                    };
8350                    if start == 0 {
8351                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8352                            if let Some(lh) = lh.as_ref() {
8353                                e.copy_into(&mut phs, 0, lh, n_embd)?;
8354                            }
8355                        }
8356                    }
8357                    if n_copy > 0 {
8358                        e.copy_view_into(
8359                            &mut phs,
8360                            dst_off,
8361                            &ph.slice(src_lo..src_lo + n_copy),
8362                            n_copy,
8363                        )?;
8364                    }
8365                    self.mtp_kv_fill(
8366                        e,
8367                        mtp,
8368                        &prompt[start..end],
8369                        &phs,
8370                        base + start,
8371                        &mut *scratch,
8372                        embd_dev,
8373                    )?;
8374                }
8375                start = end;
8376            }
8377        }
8378        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
8379        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
8380        // (=1 brackets the whole call in run_spec.rs, prime included.)
8381        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
8382            unsafe extern "C" {
8383                fn cudaProfilerStart() -> i32;
8384            }
8385            unsafe {
8386                cudaProfilerStart();
8387            }
8388        }
8389        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
8390        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
8391        // consume each other's device outputs; the host drains the ring every M rounds. v1
8392        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
8393        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
8394        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
8395        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
8396        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
8397        let stream_on = crate::spec::spec_stream()
8398            && !sampled
8399            && !spec_replay
8400            && constraint.is_none()
8401            && !session_mode
8402            && embd_gpu.is_some()
8403            && !crate::model::full_prec_enabled()
8404            && k + 2 < 96;
8405        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
8406        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
8407        if stream_on {
8408            let cap = e.capture_graph(|e| {
8409                for j in 0..k.max(1) {
8410                    self.mtp_head_forward_cap(
8411                        e,
8412                        mtp,
8413                        &mut dctx.g_tok,
8414                        &mut dctx.g_pos,
8415                        &mut dctx.g_seed,
8416                        &mut dctx.g_p,
8417                        &mut *scratch,
8418                        true,
8419                        true,
8420                        embd_gpu.expect("round stream requires resident embedding"),
8421                        embd_qt,
8422                        embd_rb,
8423                        d_vocab,
8424                        None,
8425                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
8426                        None, // round-stream requires constraint.is_none() (see stream_on)
8427                    )?;
8428                }
8429                Ok(())
8430            });
8431            match cap {
8432                Ok(g) => {
8433                    scratch.set_len(e, 0)?;
8434                    stream_graph = Some(g);
8435                }
8436                Err(err) => {
8437                    scratch.set_len(e, 0)?;
8438                    if debug_spec {
8439                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
8440                    }
8441                }
8442            }
8443        }
8444        let stream_active = stream_on && stream_graph.is_some();
8445        if debug_spec {
8446            eprintln!(
8447                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
8448                crate::spec::spec_stream(),
8449                dctx.graph.is_some(),
8450                stream_graph.is_some()
8451            );
8452        }
8453        let t_v_s = k + 1;
8454        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
8455        // module (extracted 2026-07-12; the gemma burst reuses them).
8456        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
8457        let crate::round_stream::StreamBufs {
8458            mut vtok_d,
8459            mut brk_d,
8460            mut pend_d,
8461            last_pred_d,
8462            mut pos_ctr,
8463            mut pos_start_d,
8464            mut ring_d,
8465            acc_d: mut stream_acc,
8466            m_rounds,
8467            k: _,
8468        } = sb;
8469        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
8470            Some(crate::round_stream::kv_len_ptr_table(
8471                e,
8472                cache,
8473                Some(&pos_ctr),
8474            )?)
8475        } else {
8476            None
8477        };
8478
8479        let t_fill = t_ent.elapsed();
8480        let mut round = 0usize;
8481        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
8482        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
8483        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
8484        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
8485        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
8486        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
8487        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
8488        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
8489        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
8490        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
8491        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
8492        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
8493        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
8494        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
8495        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
8496        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
8497        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
8498        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
8499        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
8500        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
8501        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
8502        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
8503        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
8504        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
8505        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
8506        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
8507        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
8508        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
8509        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
8510        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
8511            .ok()
8512            .and_then(|v| v.parse().ok());
8513        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
8514            4
8515        } else if self.cfg.n_embd as usize >= 2500 {
8516            2
8517        } else {
8518            1
8519        };
8520        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
8521        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
8522        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
8523        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
8524        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
8525            .ok()
8526            .and_then(|v| v.parse().ok())
8527            .unwrap_or(1024);
8528        let floor_at = |pos: usize| -> usize {
8529            if adapt_floor_env.is_some() || pos < floor_ctx {
8530                adapt_floor
8531            } else if adapt_floor >= 4 {
8532                1
8533            } else {
8534                adapt_floor
8535            }
8536        };
8537        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
8538        // fixed-K default path is untouched by this whole block.
8539        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
8540            .ok()
8541            .and_then(|v| v.parse().ok())
8542            .unwrap_or(7);
8543        let k_cap = k.min(cap_max).max(1);
8544        let mut kc = k_cap;
8545        let mut opti_fork: Option<OptiForkState> = None;
8546        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
8547        if fork_mode != OptiForkGateMode::Disabled {
8548            let fence = crate::pp::pp_cuts(self.layers.len());
8549            let refusal = if !session_mode {
8550                Some("not-session")
8551            } else if k != 1 || adapt {
8552                Some("requires-fixed-k1")
8553            } else if sampled || constraint.is_some() || spec_replay {
8554                Some("sampled-constrained-or-replay")
8555            } else if pipe.is_some() {
8556                Some("two-session-pipeline")
8557            } else if !spec_devacc() {
8558                Some("requires-device-accept")
8559            } else if stream_active || crate::spec::spec_stream() {
8560                Some("round-stream")
8561            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
8562                Some("swa-ring")
8563            } else if crate::pp::pp_host_bounce_active() {
8564                Some("host-bounce")
8565            } else if fork_mode == OptiForkGateMode::Controller
8566                && cache.recur.iter().any(Option::is_some)
8567            {
8568                Some("controller-requires-zero-recurrent-state")
8569            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
8570                Some("requires-pp2")
8571            } else {
8572                None
8573            };
8574            if let Some(reason) = refusal {
8575                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8576                eprintln!("[opti-fork] refused reason={reason}");
8577            } else {
8578                let fence = fence.expect("validated PP-2 fence");
8579                let rt = crate::pp::PpNRt::get(e)?;
8580                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
8581                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
8582                let primary_supported =
8583                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
8584                if !rt.cross_device() || !primary_supported {
8585                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8586                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
8587                } else {
8588                    // Both recurrent snapshots and both seed generations are allocated before
8589                    // the first fork, each through its owning PP stage. Allocation failure
8590                    // therefore happens before any optimistic state mutation can occur.
8591                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
8592                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
8593                    let fork = OptiForkState::new(
8594                        e,
8595                        cache,
8596                        fork_mode,
8597                        alternate_snapshot,
8598                        &h_seed_buf,
8599                        &fill_prev,
8600                        rt,
8601                        fence[1],
8602                        self.layers.len(),
8603                    )?;
8604                    eprintln!(
8605                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
8606                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
8607                        fence[1],
8608                        fork.logical_payload_bytes[0],
8609                        fork.logical_payload_bytes[1],
8610                        fork.controller.map_or(0.0, |policy| policy.threshold),
8611                    );
8612                    fork_snapshot = Some(current_snapshot);
8613                    opti_fork = Some(fork);
8614                }
8615            }
8616        }
8617        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
8618        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
8619        let mut snap = match fork_snapshot {
8620            Some(snapshot) => snapshot,
8621            None => cache.snapshot(e)?,
8622        };
8623        let mut carried_opti: Option<OptiControllerTicket> = None;
8624        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
8625        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
8626        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
8627            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
8628        } else {
8629            None
8630        };
8631        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
8632        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
8633        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
8634        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
8635        // pass of any kind). Verify still
8636        // checks every emitted token against the target -> exactness holds by construction; only
8637        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
8638        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
8639        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
8640        let mut pending: Option<u32> = carried_pending;
8641        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
8642        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
8643        // the verify accept readback). Printed once at loop end via spec-stats.
8644        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
8645        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
8646        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
8647        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
8648        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
8649        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
8650        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
8651        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
8652        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
8653        let mut ph_wait = 0f64;
8654        let mut ph_commit = 0f64;
8655        let mut ph_t = std::time::Instant::now();
8656        let mut ph_mark = |acc: &mut f64, on: bool| {
8657            if on {
8658                let now = std::time::Instant::now();
8659                *acc += (now - ph_t).as_secs_f64();
8660                ph_t = now;
8661            }
8662        };
8663        if let Some(p) = pipe {
8664            p.setup_end();
8665        }
8666        while keep_going && out.len() < max_new {
8667            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
8668            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
8669            if let (true, Some(sg), Some(ptrs)) = (
8670                stream_active && round >= 1 && pending.is_some(),
8671                &stream_graph,
8672                &stream_ptrs,
8673            ) {
8674                if debug_spec {
8675                    static ONCE: std::sync::Once = std::sync::Once::new();
8676                    ONCE.call_once(|| {
8677                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
8678                    });
8679                }
8680                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
8681                e.set_u32_one(&mut pend_d, pending.unwrap())?;
8682                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
8683                for _mi in 0..m_rounds {
8684                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
8685                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
8686                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
8687                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
8688                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
8689                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8690                    sg.launch()?;
8691                    e.spec_assemble_verify(
8692                        &g_tokp2k,
8693                        &pend_d,
8694                        d2t_dev.as_ref(),
8695                        &mut vtok_d,
8696                        &mut brk_d,
8697                        p_min,
8698                        k,
8699                        pmin0,
8700                    )?;
8701                    let mut ck = VerifyCkpt::new(self.layers.len());
8702                    let dummy = vec![0u32; t_v_s];
8703                    let (tl_d, vx) = self.decode_step_t_core_stream(
8704                        e,
8705                        &dummy,
8706                        0,
8707                        &mut *cache,
8708                        embd_dev,
8709                        Some(&mut ck),
8710                        Some((&vtok_d, &pos_ctr)),
8711                        None,
8712                    )?;
8713                    for j in 0..t_v_s {
8714                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
8715                    }
8716                    e.spec_accept_greedy_dc(
8717                        &preds_d,
8718                        &vtok_d,
8719                        &last_pred_d,
8720                        &brk_d,
8721                        &mut stream_acc,
8722                    )?;
8723                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
8724                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
8725                    self.commit_verified_prefix_stream(
8726                        e,
8727                        &mut *cache,
8728                        &snap,
8729                        &ck,
8730                        &stream_acc,
8731                        1,
8732                        t_v_s,
8733                    )?;
8734                    e.spec_rollback_stream(
8735                        ptrs,
8736                        &pos_start_d,
8737                        &stream_acc,
8738                        1,
8739                        self.layers.len() + 1,
8740                    )?;
8741                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
8742                }
8743                e.stream().synchronize()?;
8744                let ring_h = e.dtoh_u32(&ring_d)?;
8745                let cnt = ring_h[0] as usize;
8746                for i in 0..cnt {
8747                    if out.len() < max_new {
8748                        out.push(ring_h[1 + i]);
8749                    }
8750                }
8751                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
8752                for il in 0..self.layers.len() {
8753                    if let Some(kvl) = cache.kv[il].as_mut() {
8754                        kvl.len = pos_h;
8755                    }
8756                }
8757                cache.pos = pos_h;
8758                scratch.kv.len = pos_h;
8759                pending = Some(ring_h[cnt]); // last drained token = the live bonus
8760                last_token = ring_h[cnt];
8761                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
8762                total_accepted += cnt.saturating_sub(m_rounds);
8763                if let Some(t) = sess_telem {
8764                    // totals only — the burst's per-round accept counts stayed on device
8765                    // (that is the point of the round-stream arm). pos_* untouched.
8766                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
8767                }
8768                round += m_rounds;
8769                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
8770                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
8771                continue;
8772            }
8773            let pipe_draft = match pipe {
8774                Some(p) => Some(p.draft_begin(round)?),
8775                None => None,
8776            };
8777            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
8778            let mut current_opti = carried_opti.take();
8779            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
8780                match opti_fork.as_mut() {
8781                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
8782                    None => None,
8783                    Some(_) => None,
8784                }
8785            } else {
8786                None
8787            };
8788            if current_opti.is_none() {
8789                if let Some(fork) = opti_fork.as_ref() {
8790                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
8791                } else {
8792                    cache.snapshot_into(e, &mut snap)?;
8793                }
8794            } else if snap.pos != pos {
8795                return Err(format!(
8796                    "optipipe carried snapshot pos {} != current pos {pos}",
8797                    snap.pos
8798                )
8799                .into());
8800            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
8801            ph_mark(&mut ph_rest, phase_on);
8802
8803            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
8804            // p-min semantics (both paths): stop the chain early when the head's confidence in
8805            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
8806            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
8807            let base0 = if pending.is_some() { 1usize } else { 0usize };
8808            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
8809            // accepted run + 1 (the gemma law — see the setup block above the loop).
8810            let k_this = if adapt { kc } else { k };
8811            let mut draft: Vec<u32> = Vec::with_capacity(k);
8812            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
8813            let mut controller_draft_prob: Option<f32> = None;
8814            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
8815            if let Some(ticket) = current_opti.as_mut() {
8816                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
8817                if ticket.verify_tokens[0] != carried_pending {
8818                    return Err(format!(
8819                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
8820                        ticket.verify_tokens[0],
8821                    )
8822                    .into());
8823                }
8824                draft.push(ticket.verify_tokens[1]);
8825                controller_draft_prob = Some(ticket.draft_prob);
8826                controller_eager_state = ticket
8827                    .take_eager_seed()
8828                    .map(|seed| (ticket.verify_tokens[1], seed));
8829            } else {
8830                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
8831                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
8832                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
8833                // rejected drafts and p-min extras via the len mechanism).
8834                scratch.set_len(e, pos + base0 - 1)?;
8835                if pen_on {
8836                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
8837                    // device dedup: `penalty_last_n` is usize::MAX for any serve request with
8838                    // a penalty, so without the cap this grew with the whole session.
8839                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
8840                    let w0 = pen_hist.len().saturating_sub(win);
8841                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
8842                }
8843                if sampled {
8844                    draft_logits.clear();
8845                    draft_stats.clear();
8846                }
8847                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
8848                // position's mask is computed on that clone and advanced by the PROPOSED token. The
8849                // real state moves only on emission (verify's job), so the emitted stream is
8850                // unchanged — the mask only removes tokens the verify would have truncated anyway.
8851                let mut dmask_live = dmask_on;
8852                if dmask_live {
8853                    let t_c = std::time::Instant::now();
8854                    constraint
8855                        .as_deref_mut()
8856                        .unwrap()
8857                        .draft_begin()
8858                        .map_err(|e2| format!("constraint: {e2}"))?;
8859                    dm_clone_ns += t_c.elapsed().as_nanos();
8860                    dm_rounds += 1;
8861                }
8862                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
8863                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
8864                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
8865                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
8866                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
8867                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
8868                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8869                    for j in 0..k_this {
8870                        // per-position mask upload (contents only — the graph's baked pointer is
8871                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
8872                        // mask node degrades to a no-op ban instead of needing a second graph.
8873                        if dmask_live
8874                            && !upload_draft_mask(
8875                                e,
8876                                constraint.as_deref_mut().unwrap(),
8877                                &mut dctx.g_dmask,
8878                                mtp.d2t.as_ref(),
8879                                d_vocab,
8880                                dmask_words,
8881                            )?
8882                        {
8883                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
8884                            // genuinely miss the legal set): neutralize the captured mask node and
8885                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
8886                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
8887                            dmask_live = false;
8888                        }
8889                        gr.launch()?;
8890                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
8891                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
8892                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
8893                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
8894                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
8895                        // replay's embed node, and the MMU fault kills the CUDA context for the
8896                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
8897                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
8898                        // buffer (g_seed = the verify-side handoff vs head-side compute).
8899                        if (idx as usize) >= d_vocab {
8900                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
8901                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
8902                            // seed, untouched since the round-start copy — the pair discriminates
8903                            // "seed arrived poisoned" from "head forward produced NaN".
8904                            let seed_h = e.dtoh(&dctx.g_seed)?;
8905                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
8906                            let in_h = e.dtoh(&h_seed_buf)?;
8907                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
8908                            return Err(format!(
8909                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
8910                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
8911                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
8912                             the embed row (#87 trap)"
8913                            )
8914                            .into());
8915                        }
8916                        // trimmed draft vocab -> target token id (identity when no d2t map)
8917                        let d = match &mtp.d2t {
8918                            Some(map) => map[idx as usize],
8919                            None => idx,
8920                        };
8921                        let draft_p = if p_min > 0.0
8922                            || opti_fork
8923                                .as_ref()
8924                                .is_some_and(|fork| fork.controller.is_some())
8925                        {
8926                            Some(e.dtoh(&dctx.g_p)?[0])
8927                        } else {
8928                            None
8929                        };
8930                        if j == 0 {
8931                            controller_draft_prob = draft_p;
8932                        }
8933                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
8934                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
8935                                break;
8936                            }
8937                        }
8938                        draft.push(d);
8939                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
8940                        // index the argmax wrote — patch the persistent token buffer (4B htod).
8941                        if d != idx {
8942                            e.set_u32_one(&mut dctx.g_tok, d)?;
8943                        }
8944                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
8945                        // unmasked drafting for the remaining positions (verify still arbitrates).
8946                        // speculative advance; a chain the grammar can no longer follow (EOS
8947                        // proposed) ends here. The captured mask node always runs, so a dead chain
8948                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
8949                        if dmask_live
8950                            && !constraint
8951                                .as_deref_mut()
8952                                .unwrap()
8953                                .draft_advance(d)
8954                                .map_err(|e2| format!("constraint: {e2}"))?
8955                        {
8956                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
8957                            break;
8958                        }
8959                    }
8960                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
8961                // legal ONLY in the regime it was captured in. The condition used to read
8962                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
8963                // which it could not, because the key omitted the filters. Both halves are now
8964                // enforced: the key drops a stale graph, and this site refuses to launch one.
8965                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
8966                    if skey_probe() {
8967                        eprintln!(
8968                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
8969                             top_p={} min_p={} s_key_parked={:?}",
8970                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
8971                        );
8972                    }
8973                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
8974                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
8975                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
8976                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
8977                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
8978                    // stream. Host sctr advances in lockstep (computed, no readback needed).
8979                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
8980                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
8981                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
8982                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
8983                    for j in 0..k_this {
8984                        gr.launch()?;
8985                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
8986                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
8987                        // counts the p-min-discarded token too)
8988                        // q retention: ONE async D2D of the persistent head-logits buffer into this
8989                        // round's slot j (stream-ordered after the replay, before the next one).
8990                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
8991                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
8992                        // #87 SENTINEL TRAP (see the greedy graph arm above).
8993                        if (idx as usize) >= d_vocab {
8994                            let seed_h = e.dtoh(&dctx.g_seed)?;
8995                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
8996                            return Err(format!(
8997                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
8998                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
8999                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
9000                             (#87 trap)"
9001                            )
9002                            .into());
9003                        }
9004                        let d = match &mtp.d2t {
9005                            Some(map) => map[idx as usize],
9006                            None => idx,
9007                        };
9008                        draft_idx.push(idx);
9009                        if p_min > 0.0 {
9010                            let p = e.dtoh(&dctx.g_p)?[0];
9011                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9012                                break;
9013                            }
9014                        }
9015                        draft.push(d);
9016                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
9017                        if d != idx {
9018                            e.set_u32_one(&mut dctx.g_tok, d)?;
9019                        }
9020                    }
9021                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
9022                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
9023                    for j in 0..draft.len().max(draft_idx.len()) {
9024                        let rows0 = e.htod_i32(&[0])?;
9025                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9026                        e.filter_stats(
9027                            &dctx.q_slots[j],
9028                            d_vocab,
9029                            &rows0,
9030                            &mut th_d,
9031                            &mut z_d,
9032                            &mut mx_d,
9033                            d_vocab,
9034                            1,
9035                            sp_temp,
9036                            sp.top_k,
9037                            sp.top_p,
9038                            sp.min_p,
9039                        )?;
9040                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9041                    }
9042                } else {
9043                    if skey_probe() && sampled {
9044                        eprintln!(
9045                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
9046                             top_p={} min_p={} s_key_parked={:?}",
9047                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9048                        );
9049                    }
9050                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
9051                    let mut e_tok = last_token;
9052                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
9053                    for j in 0..k_this {
9054                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
9055                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
9056                        let mtp_pos = pos + base0 + j;
9057                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
9058                        // A position with no legal draft-vocab row drops to unmasked drafting for
9059                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
9060                        if dmask_live {
9061                            dmask_live = upload_draft_mask(
9062                                e,
9063                                constraint.as_deref_mut().unwrap(),
9064                                &mut dctx.g_dmask,
9065                                mtp.d2t.as_ref(),
9066                                d_vocab,
9067                                dmask_words,
9068                            )?;
9069                        }
9070                        let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9071                            e,
9072                            mtp,
9073                            e_tok,
9074                            &d_seed,
9075                            &mut *scratch,
9076                            mtp_pos,
9077                            embd_dev,
9078                            if dmask_live {
9079                                Some((&dctx.g_dmask, dmask_words))
9080                            } else {
9081                                None
9082                            },
9083                        )?;
9084                        let tok_d = if sampled {
9085                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
9086                            // the filtered softmax (filters off => th=0, exact v1 semantics).
9087                            if perturb_buf.is_none() {
9088                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9089                            }
9090                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
9091                            if pen_on {
9092                                let h = pen_hist_d.as_ref().unwrap();
9093                                let nh = h.len();
9094                                e.penalize_logits(
9095                                    &mut q_row,
9096                                    h,
9097                                    nh,
9098                                    sp.penalty_repeat,
9099                                    sp.penalty_freq,
9100                                    sp.penalty_present,
9101                                    d_vocab,
9102                                )?;
9103                            }
9104                            let rows0 = e.htod_i32(&[0])?;
9105                            let (mut th_d, mut z_d, mut mx_d) =
9106                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9107                            e.filter_stats(
9108                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
9109                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
9110                            )?;
9111                            let (th, z, mx) =
9112                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
9113                            let pb = perturb_buf.as_mut().unwrap();
9114                            e.gumbel_perturb_filtered(
9115                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
9116                            )?;
9117                            sctr += 1;
9118                            draft_logits.push(q_row);
9119                            draft_stats.push((mx, th, z));
9120                            e.argmax_token_device(pb, d_vocab)?
9121                        } else {
9122                            e.argmax_token_device(&dl_d, d_vocab)?
9123                        };
9124                        let idx = e.dtoh_u32_one(&tok_d)?;
9125                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
9126                        // here because the eager chain's operands are all readable: dl_d (the head
9127                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
9128                        if (idx as usize) >= d_vocab {
9129                            let dl_h = e.dtoh(&dl_d)?;
9130                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
9131                            let seed_h = e.dtoh(&d_seed)?;
9132                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9133                            return Err(format!(
9134                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9135                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
9136                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
9137                             embed row (#87 trap)"
9138                            )
9139                            .into());
9140                        }
9141                        let d = match &mtp.d2t {
9142                            Some(map) => map[idx as usize],
9143                            None => idx,
9144                        };
9145                        if sampled {
9146                            draft_idx.push(idx);
9147                        }
9148                        let draft_p = if p_min > 0.0
9149                            || opti_fork
9150                                .as_ref()
9151                                .is_some_and(|fork| fork.controller.is_some())
9152                        {
9153                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
9154                            Some(e.dtoh(&p_d)?[0])
9155                        } else {
9156                            None
9157                        };
9158                        if j == 0 {
9159                            controller_draft_prob = draft_p;
9160                        }
9161                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9162                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9163                                break;
9164                            }
9165                        }
9166                        draft.push(d);
9167                        e_tok = d;
9168                        d_seed = h_nextn;
9169                        // speculative advance; a chain the grammar can no longer follow (EOS
9170                        // proposed) ends here — the prefix already proposed still rides verify.
9171                        if dmask_live
9172                            && !constraint
9173                                .as_deref_mut()
9174                                .unwrap()
9175                                .draft_advance(d)
9176                                .map_err(|e2| format!("constraint: {e2}"))?
9177                        {
9178                            break;
9179                        }
9180                    }
9181                    if opti_fork
9182                        .as_ref()
9183                        .is_some_and(|fork| fork.controller.is_some())
9184                    {
9185                        controller_eager_state = Some((e_tok, d_seed));
9186                    }
9187                }
9188            }
9189            let k_round = draft.len();
9190            if let Some(p) = pipe {
9191                p.draft_end(round);
9192            }
9193            drop(pipe_draft);
9194
9195            ph_mark(&mut ph_draft, phase_on);
9196            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
9197            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
9198            let verify_tokens: Vec<u32> = match pending {
9199                Some(b) => {
9200                    let mut v = Vec::with_capacity(k_round + 1);
9201                    v.push(b);
9202                    v.extend_from_slice(&draft);
9203                    v
9204                }
9205                None => draft.clone(),
9206            };
9207            let base = if pending.is_some() { 1 } else { 0 };
9208            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
9209            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
9210            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
9211                Some(ticket.take_ckpt())
9212            } else if spec_replay {
9213                None
9214            } else {
9215                Some(VerifyCkpt::new(self.layers.len()))
9216            };
9217            let controller_can_probe = base == 1
9218                && k_round == 1
9219                && out.len().saturating_add(2) < max_new
9220                && controller_draft_prob.is_some()
9221                && opti_fork
9222                    .as_ref()
9223                    .and_then(|fork| fork.controller.as_ref())
9224                    .is_some_and(|policy| !policy.breaker_tripped);
9225            let mut successor_attempt: Option<OptiControllerTicket> = None;
9226            let mut rejected_probe: Option<(f32, u32)> = None;
9227            let mut controller_prepared: Option<OptiControllerPrepared> = None;
9228            if controller_can_probe {
9229                // Prepare d2/q and, on admission, d3 before either current verify half is
9230                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
9231                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
9232                // the primary stream after N stage 1 would serialize the supposed pipeline.
9233                let eager_pos = scratch.kv.len + 1;
9234                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
9235                    e,
9236                    mtp,
9237                    &mut dctx,
9238                    &mut *scratch,
9239                    d_vocab,
9240                    &mut controller_eager_state,
9241                    eager_pos,
9242                    embd_dev,
9243                )?;
9244                let first_probability = controller_draft_prob
9245                    .ok_or("optipipe controller probe lost first-token probability")?;
9246                let q_proxy = first_probability * pending_probability;
9247                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9248                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9249                let admitted = opti_fork
9250                    .as_ref()
9251                    .and_then(|fork| fork.controller.as_ref())
9252                    .ok_or("optipipe controller policy disappeared")?
9253                    .admit(q_proxy);
9254                if admitted {
9255                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9256                    let eager_pos = scratch.kv.len + 1;
9257                    let (optimistic_draft, optimistic_draft_probability) = self
9258                        .opti_controller_draft_step(
9259                            e,
9260                            mtp,
9261                            &mut dctx,
9262                            &mut *scratch,
9263                            d_vocab,
9264                            &mut controller_eager_state,
9265                            eager_pos,
9266                            embd_dev,
9267                        )?;
9268                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9269                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
9270                        debug_assert_eq!(token, optimistic_draft);
9271                        seed
9272                    });
9273                    controller_prepared = Some(OptiControllerPrepared {
9274                        verify_tokens: [optimistic_pending, optimistic_draft],
9275                        draft_prob: optimistic_draft_probability,
9276                        eager_seed,
9277                        q_proxy,
9278                        scratch_len: scratch.kv.len,
9279                    });
9280                } else {
9281                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9282                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9283                    rejected_probe = Some((q_proxy, optimistic_pending));
9284                    eprintln!(
9285                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
9286                        opti_fork
9287                            .as_ref()
9288                            .and_then(|fork| fork.controller.as_ref())
9289                            .expect("controller policy")
9290                            .threshold,
9291                    );
9292                }
9293            }
9294            let fork_attempt = match fork_generation.take() {
9295                Some(generation) if base == 1 && k_round == 1 => Some(generation),
9296                Some(generation) => {
9297                    opti_fork
9298                        .as_mut()
9299                        .expect("fork generation without fork state")
9300                        .retire(generation)?;
9301                    None
9302                }
9303                None => None,
9304            };
9305            let (tlogits_d, vx) = if let Some(p) = pipe {
9306                self.decode_step_t_core_pipelined(
9307                    e,
9308                    &verify_tokens,
9309                    pos,
9310                    &mut *cache,
9311                    embd_dev,
9312                    ckpt.as_mut(),
9313                    p,
9314                    round,
9315                )?
9316            } else if controller_can_probe {
9317                let fence = opti_fork
9318                    .as_ref()
9319                    .ok_or("optipipe controller probe lost fork state")?
9320                    .fence;
9321                let boundary = match current_opti.as_mut() {
9322                    Some(ticket) => ticket.take_boundary(),
9323                    None => self.verify_stage0_issue(
9324                        e,
9325                        &verify_tokens,
9326                        pos,
9327                        &mut *cache,
9328                        embd_dev,
9329                        ckpt.as_mut(),
9330                        None,
9331                        &fence,
9332                        Some(true),
9333                        None,
9334                    )?,
9335                };
9336                if let Some(prepared) = controller_prepared.take() {
9337                    let generation = {
9338                        let fork = opti_fork
9339                            .as_mut()
9340                            .ok_or("optipipe controller admission lost fork state")?;
9341                        let generation = fork.reserve_successor()?;
9342                        let rt = fork.rt;
9343                        let snapshot_fence = fork.fence;
9344                        opti_snapshot_one_stage_owned_into(
9345                            e,
9346                            cache,
9347                            rt,
9348                            &snapshot_fence,
9349                            0,
9350                            fork.successor_snapshot_mut(),
9351                        )?;
9352                        generation
9353                    };
9354                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
9355                    let successor_boundary = self.verify_stage0_issue(
9356                        e,
9357                        &prepared.verify_tokens,
9358                        pos + verify_tokens.len(),
9359                        &mut *cache,
9360                        embd_dev,
9361                        Some(&mut successor_ckpt),
9362                        None,
9363                        &fence,
9364                        Some(false),
9365                        None,
9366                    )?;
9367                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9368                    let fork = opti_fork
9369                        .as_ref()
9370                        .ok_or("optipipe controller ticket lost fork state")?;
9371                    successor_attempt = Some(fork.controller_ticket(
9372                        generation,
9373                        successor_boundary,
9374                        successor_ckpt,
9375                        prepared.verify_tokens,
9376                        prepared.draft_prob,
9377                        prepared.eager_seed,
9378                        prepared.q_proxy,
9379                        prepared.scratch_len,
9380                    ));
9381                    eprintln!(
9382                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
9383                         verify={:?}",
9384                        generation.id,
9385                        prepared.q_proxy,
9386                        fork.controller.expect("controller policy").threshold,
9387                        prepared.verify_tokens,
9388                    );
9389                }
9390                let result = self.verify_stage1_finish(
9391                    e,
9392                    boundary,
9393                    &mut *cache,
9394                    ckpt.as_mut(),
9395                    None,
9396                    &fence,
9397                    successor_attempt.is_none(),
9398                )?;
9399                if let Some(ticket) = current_opti.as_mut() {
9400                    ticket.settle();
9401                }
9402                if successor_attempt.is_some() {
9403                    let fork = opti_fork
9404                        .as_mut()
9405                        .ok_or("optipipe successor snapshot lost fork state")?;
9406                    let rt = fork.rt;
9407                    let snapshot_fence = fork.fence;
9408                    opti_snapshot_one_stage_owned_into(
9409                        e,
9410                        cache,
9411                        rt,
9412                        &snapshot_fence,
9413                        1,
9414                        fork.successor_snapshot_mut(),
9415                    )?;
9416                    // Publish N only after both independent successor-state queues are complete.
9417                    fork.rt.publish_to(1, &e.stream())?;
9418                }
9419                result
9420            } else if let Some(ticket) = current_opti.as_mut() {
9421                let fork = opti_fork
9422                    .as_mut()
9423                    .ok_or("optipipe carried controller ticket lost fork state")?;
9424                let boundary = ticket.take_boundary();
9425                let result = self.verify_stage1_finish(
9426                    e,
9427                    boundary,
9428                    &mut *cache,
9429                    ckpt.as_mut(),
9430                    None,
9431                    &fork.fence,
9432                    true,
9433                )?;
9434                ticket.settle();
9435                result
9436            } else if let Some(generation) = fork_attempt {
9437                let fork = opti_fork
9438                    .as_mut()
9439                    .expect("fork generation without fork state");
9440                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
9441                let action = fork.mode.action(generation.id);
9442                let boundary = self.verify_stage0_issue(
9443                    e,
9444                    &verify_tokens,
9445                    pos,
9446                    &mut *cache,
9447                    embd_dev,
9448                    ckpt.as_mut(),
9449                    None,
9450                    &fork.fence,
9451                    Some(true),
9452                    None,
9453                )?;
9454                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9455                let mut ticket = fork.ticket(generation, boundary);
9456                if action == OptiForkAction::Abort {
9457                    return Err(format!(
9458                        "optipipe forced abort with generation {} stage0 in flight",
9459                        generation.id,
9460                    )
9461                    .into());
9462                }
9463                fork.reconcile(
9464                    e,
9465                    &mut *cache,
9466                    &mut *scratch,
9467                    &snap,
9468                    &mut h_seed_buf,
9469                    &mut fill_prev,
9470                    generation,
9471                    action,
9472                    verify_tokens[0],
9473                )?;
9474                let result = if action == OptiForkAction::Hit {
9475                    let boundary = ticket.take_boundary();
9476                    self.verify_stage1_finish(
9477                        e,
9478                        boundary,
9479                        &mut *cache,
9480                        ckpt.as_mut(),
9481                        None,
9482                        &fork.fence,
9483                        true,
9484                    )?
9485                } else {
9486                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
9487                    // verify only after E_restart published the restored stage-0 state.
9488                    self.decode_step_t_core(
9489                        e,
9490                        &verify_tokens,
9491                        pos,
9492                        &mut *cache,
9493                        embd_dev,
9494                        ckpt.as_mut(),
9495                    )?
9496                };
9497                ticket.settle();
9498                debug_assert_eq!(ticket.generation, generation);
9499                fork.retire(generation)?;
9500                result
9501            } else {
9502                self.decode_step_t_core(
9503                    e,
9504                    &verify_tokens,
9505                    pos,
9506                    &mut *cache,
9507                    embd_dev,
9508                    ckpt.as_mut(),
9509                )?
9510            };
9511            let pipe_accept = match pipe {
9512                Some(p) => Some(p.accept_begin(round)?),
9513                None => None,
9514            };
9515
9516            ph_mark(&mut ph_verify, phase_on);
9517            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
9518            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
9519            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
9520            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
9521            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
9522            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
9523            // (== the bonus), so every index shifts by `base` and last_pred is unused.
9524            let t_v = verify_tokens.len();
9525            let mut preds: Vec<u32> = Vec::new();
9526            if !sampled {
9527                for j in 0..t_v {
9528                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
9529                }
9530                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
9531                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
9532                // next round's last_token = the next chain's embed lookup. Catch it at the
9533                // source with the column named — an all-NaN VERIFY column implicates the
9534                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
9535                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
9536                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
9537                    let mut probe = e.zeros(n_vocab)?;
9538                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
9539                    let col_h = e.dtoh(&probe)?;
9540                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
9541                    return Err(format!(
9542                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
9543                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
9544                         — the stage-split verify produced a poisoned column (#87 trap)",
9545                        preds[bad]
9546                    )
9547                    .into());
9548                }
9549            }
9550            ph_mark(&mut ph_wait, phase_on);
9551            let t_pred = |j: usize| -> u32 {
9552                if j == 0 && base == 0 {
9553                    last_pred
9554                } else {
9555                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
9556                    // used to call this from the sampled arm and panicked the worker; it now goes
9557                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
9558                    // out-of-range pred is a real bug, not something to paper over.
9559                    debug_assert!(
9560                        !sampled,
9561                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
9562                    );
9563                    preds[base + j - 1]
9564                }
9565            };
9566            let mut devacc_seeded = false;
9567            let mut devacc_acc: Option<CudaSlice<u32>> = None;
9568            let (n_acc, bonus) = if !sampled {
9569                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
9570                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
9571                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
9572                // gated on token identity vs the host walk (the arms below are bit-equal rules).
9573                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
9574                {
9575                    let draft_d = e.htod_u32_v(&draft)?;
9576                    let mut acc_out = e.alloc_u32_zeroed(2)?;
9577                    e.spec_accept_greedy(
9578                        &preds_d,
9579                        &draft_d,
9580                        last_pred,
9581                        base,
9582                        k_round,
9583                        &mut acc_out,
9584                    )?;
9585                    devacc_acc = Some(acc_out.clone());
9586                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
9587                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
9588                    // non-replay commit arms skip their host-offset seed copies (guarded below);
9589                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
9590                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
9591                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
9592                    // the update lands after the arms (devacc_seeded guard below).
9593                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
9594                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
9595                    // unified rule; full accept rewrites the verify-left value). Host mirrors
9596                    // update after the readback; commit_verified_prefix skips its len_d writes.
9597                    if let Some(successor) = successor_attempt.as_ref() {
9598                        opti_fork
9599                            .as_mut()
9600                            .ok_or("optipipe successor reconcile lost fork state")?
9601                            .queue_actual_reconcile(
9602                                e,
9603                                &snap,
9604                                &acc_out,
9605                                successor.verify_tokens[0],
9606                                base,
9607                            )?;
9608                    } else if let Some(ptrs) = &kv_len_ptrs {
9609                        let saved: Vec<i32> = (0..self.layers.len())
9610                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
9611                            .collect();
9612                        let saved_d = e.htod_i32(&saved)?;
9613                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
9614                    }
9615                    devacc_seeded = true;
9616                    let ab = e.dtoh_u32(&acc_out)?;
9617                    (ab[0] as usize, ab[1])
9618                } else {
9619                    let mut n_acc = 0usize;
9620                    for j in 0..k_round {
9621                        if t_pred(j) == draft[j] {
9622                            n_acc += 1;
9623                        } else {
9624                            break;
9625                        }
9626                    }
9627                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
9628                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
9629                    (n_acc, t_pred(n_acc))
9630                }
9631            } else {
9632                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
9633                if col_buf.is_none() {
9634                    col_buf = Some(e.zeros(n_vocab)?);
9635                }
9636                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
9637                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
9638                let mut pj = vec![0f32; k_round.max(1)];
9639                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
9640                if k_round > 0 {
9641                    let mut ids: Vec<u32> = Vec::new();
9642                    let mut rows: Vec<i32> = Vec::new();
9643                    for j in 0..k_round {
9644                        if j > 0 || base == 1 {
9645                            ids.push(draft[j]);
9646                            rows.push((base + j) as i32 - 1);
9647                        }
9648                    }
9649                    if !ids.is_empty() {
9650                        let nr = rows.len();
9651                        // penalties: materialize the used columns into one contiguous penalized
9652                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
9653                        // penalties: materialize used columns contiguously, penalize all rows in
9654                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
9655                        let p_rows: Vec<i32> = if pen_on {
9656                            (0..nr as i32).collect()
9657                        } else {
9658                            rows.clone()
9659                        };
9660                        if pen_on {
9661                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
9662                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
9663                            }
9664                            let pc = pcol_buf.as_mut().unwrap();
9665                            for (i2, &r) in rows.iter().enumerate() {
9666                                let c = r as usize;
9667                                e.copy_view_into(
9668                                    pc,
9669                                    i2 * n_vocab,
9670                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
9671                                    n_vocab,
9672                                )?;
9673                            }
9674                            let h = pen_hist_d.as_ref().unwrap();
9675                            let nh = h.len();
9676                            e.penalize_logits_rows(
9677                                pc,
9678                                h,
9679                                nh,
9680                                sp.penalty_repeat,
9681                                sp.penalty_freq,
9682                                sp.penalty_present,
9683                                n_vocab,
9684                                nr,
9685                            )?;
9686                        }
9687                        let p_src: &CudaSlice<f32> = if pen_on {
9688                            pcol_buf.as_ref().unwrap()
9689                        } else {
9690                            &tlogits_d
9691                        };
9692                        let rowsd = e.htod_i32(&p_rows)?;
9693                        let (mut th_d, mut z_d, mut mx_d) =
9694                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
9695                        e.filter_stats(
9696                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
9697                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
9698                        )?;
9699                        let idsd = e.htod_u32_v(&ids)?;
9700                        let mut outd = e.zeros(nr)?;
9701                        e.softmax_gather_filtered(
9702                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
9703                            sp_temp,
9704                        )?;
9705                        let outv = e.dtoh(&outd)?;
9706                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
9707                        let mut oi = 0usize;
9708                        for j in 0..k_round {
9709                            if j > 0 || base == 1 {
9710                                pj[j] = outv[oi];
9711                                oi += 1;
9712                            }
9713                        }
9714                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
9715                    }
9716                    if base == 0 {
9717                        let lc: &CudaSlice<f32> = if pen_on {
9718                            if col_buf.is_none() {
9719                                col_buf = Some(e.zeros(n_vocab)?);
9720                            }
9721                            let cb = col_buf.as_mut().unwrap();
9722                            e.copy_into(
9723                                cb,
9724                                0,
9725                                last_col_logits
9726                                    .as_ref()
9727                                    .expect("sampled: last_col_logits unset"),
9728                                n_vocab,
9729                            )?;
9730                            let h = pen_hist_d.as_ref().unwrap();
9731                            let nh = h.len();
9732                            e.penalize_logits(
9733                                cb,
9734                                h,
9735                                nh,
9736                                sp.penalty_repeat,
9737                                sp.penalty_freq,
9738                                sp.penalty_present,
9739                                n_vocab,
9740                            )?;
9741                            col_buf.as_ref().unwrap()
9742                        } else {
9743                            last_col_logits
9744                                .as_ref()
9745                                .expect("sampled: last_col_logits unset")
9746                        };
9747                        let rows0 = e.htod_i32(&[0])?;
9748                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9749                        e.filter_stats(
9750                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
9751                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
9752                        )?;
9753                        let idsd = e.htod_u32_v(&[draft[0]])?;
9754                        let mut outd = e.zeros(1)?;
9755                        e.softmax_gather_filtered(
9756                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
9757                        )?;
9758                        pj[0] = e.dtoh(&outd)?[0];
9759                        last_col_stats =
9760                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9761                    }
9762                }
9763                // q source: the graph arm retained the head logits in the persistent q_slots;
9764                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
9765                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
9766                // computes them post-replay — graph engages only filter/penalty-free, so the
9767                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
9768                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
9769                    &dctx.q_slots
9770                } else {
9771                    &draft_logits
9772                };
9773                let mut n_acc = 0usize;
9774                for j in 0..k_round {
9775                    let (qmx, qth, qz) = draft_stats[j];
9776                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
9777                    let rowsd = e.htod_i32(&[0])?;
9778                    let thd = e.htod(&[qth])?;
9779                    let zd = e.htod(&[qz])?;
9780                    let _ = qmx;
9781                    let mut outd = e.zeros(1)?;
9782                    e.softmax_gather_filtered(
9783                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
9784                        sp_temp,
9785                    )?;
9786                    let qj = e.dtoh(&outd)?[0];
9787                    let u = host_u01(sp_seed, uctr);
9788                    uctr += 1;
9789                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
9790                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
9791                    // exactness signature (see `skey_probe`). Impossible when the draft was
9792                    // drawn from the same filtered distribution the verify reconstructs here;
9793                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
9794                    if skey_probe() && qj == 0.0 {
9795                        eprintln!(
9796                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
9797                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
9798                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
9799                        );
9800                    }
9801                    if accept {
9802                        n_acc += 1;
9803                    } else {
9804                        break;
9805                    }
9806                }
9807                let bonus = if n_acc == k_round {
9808                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
9809                    let col = base + k_round - 1;
9810                    let cb = col_buf.as_mut().unwrap();
9811                    e.copy_view_into(
9812                        cb,
9813                        0,
9814                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
9815                        n_vocab,
9816                    )?;
9817                    if pen_on {
9818                        let h = pen_hist_d.as_ref().unwrap();
9819                        let nh = h.len();
9820                        e.penalize_logits(
9821                            cb,
9822                            h,
9823                            nh,
9824                            sp.penalty_repeat,
9825                            sp.penalty_freq,
9826                            sp.penalty_present,
9827                            n_vocab,
9828                        )?;
9829                    }
9830                    if perturb_buf.is_none() {
9831                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9832                    }
9833                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
9834                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
9835                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
9836                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
9837                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
9838                    // last gathered column, in both base arms. `th` is a threshold in e-units of
9839                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
9840                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
9841                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
9842                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
9843                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
9844                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
9845                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
9846                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
9847                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
9848                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
9849                    // and row_max is unused once nothing is masked), so this fix is a byte-level
9850                    // no-op for the untruncated serve default. One extra one-block filter_stats
9851                    // per full-accept round is the whole cost.
9852                    let (mx, th) = {
9853                        let rows0 = e.htod_i32(&[0])?;
9854                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9855                        let cb0 = col_buf.as_ref().unwrap();
9856                        e.filter_stats(
9857                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
9858                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
9859                        )?;
9860                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
9861                    };
9862                    let pb = perturb_buf.as_mut().unwrap();
9863                    let cb2 = col_buf.as_ref().unwrap();
9864                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
9865                    sctr += 1;
9866                    let td = e.argmax_token_device(pb, n_vocab)?;
9867                    e.dtoh_u32_one(&td)?
9868                } else {
9869                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
9870                    let cb = col_buf.as_mut().unwrap();
9871                    if n_acc > 0 || base == 1 {
9872                        let col = base + n_acc - 1;
9873                        e.copy_view_into(
9874                            cb,
9875                            0,
9876                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
9877                            n_vocab,
9878                        )?;
9879                    } else {
9880                        let lc = last_col_logits.as_ref().unwrap();
9881                        e.copy_into(cb, 0, lc, n_vocab)?;
9882                    }
9883                    if pen_on {
9884                        let h = pen_hist_d.as_ref().unwrap();
9885                        let nh = h.len();
9886                        e.penalize_logits(
9887                            cb,
9888                            h,
9889                            nh,
9890                            sp.penalty_repeat,
9891                            sp.penalty_freq,
9892                            sp.penalty_present,
9893                            n_vocab,
9894                        )?;
9895                    }
9896                    let cb2 = col_buf.as_ref().unwrap();
9897                    let sc = sctr;
9898                    sctr += 1;
9899                    // p-stats for the reject column: from col_stats when the col was gathered,
9900                    // else (j==0&&base==0) from last_col_stats.
9901                    let p_stats = if n_acc > 0 || base == 1 {
9902                        // col index within the gathered set == number of gathered cols before n_acc
9903                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
9904                        col_stats.get(gi).copied().unwrap_or_else(|| {
9905                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
9906                        })
9907                    } else {
9908                        last_col_stats.expect("sampled: last_col_stats unset at reject")
9909                    };
9910                    let q_stats = draft_stats[n_acc];
9911                    if let Some(map) = &d2t_dev {
9912                        if q_full_buf.is_none() {
9913                            q_full_buf = Some(e.zeros(n_vocab)?);
9914                        }
9915                        let qf = q_full_buf.as_mut().unwrap();
9916                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
9917                        let qf2 = q_full_buf.as_ref().unwrap();
9918                        e.residual_sample_filtered(
9919                            cb2,
9920                            Some(qf2),
9921                            n_vocab,
9922                            sp_temp,
9923                            sp_seed,
9924                            sc,
9925                            p_stats,
9926                            q_stats,
9927                            &mut sample_tok,
9928                        )?;
9929                    } else {
9930                        e.residual_sample_filtered(
9931                            cb2,
9932                            Some(&q_bufs[n_acc]),
9933                            n_vocab,
9934                            sp_temp,
9935                            sp_seed,
9936                            sc,
9937                            p_stats,
9938                            q_stats,
9939                            &mut sample_tok,
9940                        )?;
9941                    }
9942                    e.dtoh_u32(&sample_tok)?[0]
9943                };
9944                (n_acc, bonus)
9945            };
9946            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
9947            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
9948            // ordering). Walk the accepted drafts through the grammar in commit order; the
9949            // first illegal token truncates acceptance at its slot, and that slot's emission
9950            // is recomputed as the MASKED argmax of the target's own verify column — token-
9951            // identical to constrained plain greedy decode (an unmasked argmax that is
9952            // grammar-legal IS the masked argmax: masking only removes competitors). The
9953            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
9954            // measured in acceptance numbers, never hidden.
9955            let (n_acc, bonus) = match constraint.as_deref_mut() {
9956                None => (n_acc, bonus),
9957                Some(c) => {
9958                    fn ce(e2: String) -> Box<dyn std::error::Error> {
9959                        format!("constraint: {e2}").into()
9960                    }
9961                    let mut na = n_acc;
9962                    let mut cut = false;
9963                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
9964                        if c.is_allowed(d).map_err(ce)? {
9965                            c.consume(d).map_err(ce)?;
9966                        } else {
9967                            na = j;
9968                            cut = true;
9969                            dm_cut_tokens += n_acc - j;
9970                            break;
9971                        }
9972                    }
9973                    if cut {
9974                        dm_cuts += 1;
9975                    }
9976                    let mut bo = bonus;
9977                    if cut || !c.is_allowed(bo).map_err(ce)? {
9978                        let mut row = if na == 0 && base == 0 {
9979                            init_logits_host
9980                                .clone()
9981                                .ok_or("constraint: init logits missing (round-0 cut)")?
9982                        } else {
9983                            e.dtoh_view(
9984                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
9985                            )?
9986                        };
9987                        c.mask_logits(&mut row).map_err(ce)?;
9988                        bo = argmax(&row) as u32;
9989                    }
9990                    c.consume(bo).map_err(ce)?;
9991                    (na, bo)
9992                }
9993            };
9994            let mut successor_valid = false;
9995            if let Some((q_proxy, expected_d2)) = rejected_probe {
9996                let v_n = n_acc == 1 && bonus == expected_d2;
9997                eprintln!(
9998                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
9999                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
10000                );
10001            }
10002            if let Some(successor) = successor_attempt.as_ref() {
10003                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
10004                let generation = successor.generation;
10005                let q_proxy = successor.q_proxy;
10006                let expected_pending = successor.verify_tokens[0];
10007                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
10008                let fork = opti_fork
10009                    .as_mut()
10010                    .ok_or("optipipe successor resolution lost fork state")?;
10011                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
10012                if successor_valid {
10013                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10014                } else {
10015                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10016                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10017                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
10018                }
10019                let breaker_tripped = fork
10020                    .controller
10021                    .as_mut()
10022                    .expect("controller policy")
10023                    .resolve(successor_valid);
10024                if breaker_tripped {
10025                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10026                }
10027                eprintln!(
10028                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
10029                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
10030                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
10031                    generation.id, successor_valid, !successor_valid, breaker_tripped,
10032                );
10033                if !successor_valid {
10034                    let mut successor = successor_attempt
10035                        .take()
10036                        .expect("controller successor disappeared on miss");
10037                    successor.settle();
10038                    fork.retire(generation)?;
10039                }
10040            }
10041            total_drafted += k_round;
10042            total_accepted += n_acc;
10043            if let Some(t) = sess_telem {
10044                // Greedy, rejection-sampling, and grammar truncation all converge here after
10045                // the accept decision is already on host. Fixed-size relaxed atomics only.
10046                t.record_round(k_round, n_acc);
10047            }
10048            if spec_stats {
10049                st_len_hist[k_round] += 1;
10050                for j in 0..k_round {
10051                    st_drafted[j] += 1;
10052                }
10053                for j in 0..n_acc {
10054                    st_accepted[j] += 1;
10055                }
10056                if n_acc == k_round {
10057                    st_full += 1;
10058                }
10059            }
10060
10061            if debug_spec {
10062                eprintln!(
10063                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
10064                    out.len(),
10065                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
10066                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
10067                    // the GPU worker thread — a debug flag that killed the exact regime you would
10068                    // set it to investigate. See `debug_t_pred0`.
10069                    debug_t_pred0(sampled, base, last_pred, &preds)
10070                );
10071            }
10072
10073            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
10074            let commit_started = std::time::Instant::now();
10075            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
10076            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
10077            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
10078            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
10079            for j in 0..n_acc {
10080                if !session_mode && out.len() >= max_new {
10081                    break;
10082                }
10083                out.push(draft[j]);
10084            }
10085            if pen_on {
10086                pen_hist.extend_from_slice(&draft[0..n_acc]);
10087                pen_hist.push(bonus);
10088            }
10089            let bonus_emitted = session_mode || out.len() < max_new;
10090            if bonus_emitted {
10091                out.push(bonus);
10092            }
10093            last_token = bonus;
10094
10095            // --- 5. ROLLBACK + advance (§C) ---
10096            if n_acc == k_round && !spec_replay {
10097                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
10098                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
10099                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
10100                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
10101                // last_pred is dead in the pending path (t_pred reads verify col 0).
10102                //
10103                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
10104                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
10105                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
10106                // trunk hidden (the last verify column). set_len first: a p-min break may have
10107                // left one extra chain append at that slot. Partial accepts need NO fill (the
10108                // chain already covered every accepted position; round-start set_len truncates).
10109                let mut vh_seed = e.zeros(n_embd)?;
10110                e.copy_view_into(
10111                    &mut vh_seed,
10112                    0,
10113                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
10114                    n_embd,
10115                )?;
10116                if refresh {
10117                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
10118                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
10119                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
10120                    // the full stack (vx) is already resident from the verify. Replaces both the
10121                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
10122                    // (draft attention quality); exactness stays the verify's job.
10123                    scratch.set_len(e, pos)?;
10124                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
10125                    // (hidden of the last committed row before this verify batch).
10126                    let mut vxs = e.zeros(t_v * n_embd)?;
10127                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10128                    if t_v > 1 {
10129                        e.copy_view_into(
10130                            &mut vxs,
10131                            n_embd,
10132                            &vx.slice(0..(t_v - 1) * n_embd),
10133                            (t_v - 1) * n_embd,
10134                        )?;
10135                    }
10136                    self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
10137                } else {
10138                    scratch.set_len(e, pos + base + k_round - 1)?;
10139                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
10140                    let mut hp = e.zeros(n_embd)?;
10141                    if t_v >= 2 {
10142                        e.copy_view_into(
10143                            &mut hp,
10144                            0,
10145                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
10146                            n_embd,
10147                        )?;
10148                    } else {
10149                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
10150                    }
10151                    self.mtp_kv_fill(
10152                        e,
10153                        mtp,
10154                        &[draft[k_round - 1]],
10155                        &hp,
10156                        pos + base + k_round - 1,
10157                        &mut *scratch,
10158                        embd_dev,
10159                    )?;
10160                }
10161                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
10162                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
10163                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
10164                // col). Saves one MTP-block pass per round on top of the pairing fix.
10165                if !devacc_seeded {
10166                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
10167                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
10168                }
10169                pending = Some(bonus);
10170                if debug_spec {
10171                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
10172                }
10173            } else if !spec_replay && base + n_acc >= 1 {
10174                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
10175                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
10176                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
10177                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
10178                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
10179                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
10180                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
10181                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
10182                // accept (never compounds: the next verify recomputes true hiddens for all
10183                // committed columns).
10184                let j = base + n_acc;
10185                self.commit_verified_prefix(
10186                    e,
10187                    &mut *cache,
10188                    &snap,
10189                    ckpt.as_ref().unwrap(),
10190                    j,
10191                    devacc_seeded,
10192                    if devacc_seeded {
10193                        devacc_acc.as_ref().map(|a| (a, base, t_v))
10194                    } else {
10195                        None
10196                    },
10197                )?;
10198                let mut seed = e.zeros(n_embd)?;
10199                e.copy_view_into(
10200                    &mut seed,
10201                    0,
10202                    &vx.slice((j - 1) * n_embd..j * n_embd),
10203                    n_embd,
10204                )?;
10205                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
10206                // branch); without it the chain entries stand and only the tail truncates. Either
10207                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
10208                // (persistent mode), rope pos+j+1 (chain convention).
10209                if refresh {
10210                    scratch.set_len(e, pos)?;
10211                    let mut vxs = e.zeros(j * n_embd)?;
10212                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10213                    if j > 1 {
10214                        e.copy_view_into(
10215                            &mut vxs,
10216                            n_embd,
10217                            &vx.slice(0..(j - 1) * n_embd),
10218                            (j - 1) * n_embd,
10219                        )?;
10220                    }
10221                    self.mtp_kv_fill(
10222                        e,
10223                        mtp,
10224                        &verify_tokens[0..j],
10225                        &vxs,
10226                        pos,
10227                        &mut *scratch,
10228                        embd_dev,
10229                    )?;
10230                } else {
10231                    scratch.set_len(e, pos + j)?;
10232                }
10233                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
10234                // bonus's predecessor (verify col j-1); no pseudo pass.
10235                if !devacc_seeded {
10236                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
10237                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
10238                }
10239                pending = Some(bonus);
10240                if debug_spec {
10241                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
10242                }
10243            } else if !spec_replay {
10244                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
10245                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
10246                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
10247                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
10248                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
10249                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
10250                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
10251                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
10252                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
10253                cache.rollback(e, &snap, 0)?;
10254                scratch.set_len(e, pos)?;
10255                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
10256                pending = Some(bonus);
10257                if debug_spec {
10258                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
10259                }
10260            } else {
10261                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
10262                // this round survives, only possible before the first pending exists, ~round 0):
10263                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
10264                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
10265                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
10266                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
10267                // trunk hidden.
10268                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
10269                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
10270                if let Some(b) = pending.take() {
10271                    replay.push(b);
10272                }
10273                replay.extend_from_slice(&draft[0..n_acc]);
10274                replay.push(bonus);
10275                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
10276                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
10277                // last col exactly as before (byte-identical to the old _h_emb_dev call).
10278                let (rl_d, rx) = if self.qwen35_serving_class() {
10279                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
10280                    let mut hidden = e.uninit(replay.len() * n_embd)?;
10281                    for (row, &token) in replay.iter().enumerate() {
10282                        let (row_logits, row_hidden) =
10283                            self.spec_target_step_h(e, token, &mut *cache)?;
10284                        logits.extend_from_slice(&row_logits);
10285                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
10286                    }
10287                    (e.htod(&logits)?, hidden)
10288                } else {
10289                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
10290                };
10291                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
10292                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
10293                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
10294                last_pred = e.dtoh_u32(&preds_d)?[0];
10295                if sampled {
10296                    let lr0 = replay.len();
10297                    let lc = last_col_logits
10298                        .as_mut()
10299                        .expect("sampled: last_col_logits unset");
10300                    e.copy_view_into(
10301                        lc,
10302                        0,
10303                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
10304                        n_vocab,
10305                    )?;
10306                }
10307                let lr = replay.len();
10308                if lr >= 2 {
10309                    e.copy_view_into(
10310                        &mut h_seed_buf,
10311                        0,
10312                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
10313                        n_embd,
10314                    )?;
10315                } else {
10316                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
10317                    // last_token, whose own-row hidden fill_prev still holds.
10318                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
10319                }
10320                // the bonus is COMMITTED here — it becomes the last committed row.
10321                let mut rh_last = e.zeros(n_embd)?;
10322                e.copy_view_into(
10323                    &mut rh_last,
10324                    0,
10325                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
10326                    n_embd,
10327                )?;
10328                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
10329                if debug_spec {
10330                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
10331                }
10332            }
10333            if devacc_seeded {
10334                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
10335                // consumed the old value (both slots carry the same value in every non-replay arm).
10336                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
10337            }
10338            if successor_valid {
10339                let optimistic_scratch_len = successor_attempt
10340                    .as_ref()
10341                    .expect("valid controller successor disappeared")
10342                    .scratch_len;
10343                // The normal current-round commit refreshed/truncated the logical scratch tail.
10344                // Its optimistic successor row was already written physically, so restoring only
10345                // the retained logical length makes that row live for the carried round.
10346                scratch.set_len(e, optimistic_scratch_len)?;
10347            }
10348            if let Some(current) = current_opti.take() {
10349                opti_fork
10350                    .as_mut()
10351                    .ok_or("optipipe current retirement lost fork state")?
10352                    .retire(current.generation)?;
10353            }
10354            if successor_valid {
10355                let successor = successor_attempt
10356                    .take()
10357                    .expect("valid controller successor disappeared before promotion");
10358                let generation = successor.generation;
10359                opti_fork
10360                    .as_mut()
10361                    .ok_or("optipipe successor promotion lost fork state")?
10362                    .promote_successor_snapshot(&mut snap, generation);
10363                carried_opti = Some(successor);
10364            }
10365            if anatomy_on {
10366                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
10367                // only for this diagnostic so it does not disappear into the following draft's
10368                // first token readback.
10369                e.stream().synchronize()?;
10370                ph_commit += commit_started.elapsed().as_secs_f64();
10371            }
10372            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
10373            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
10374            // final position — the floor's position key reads the committed depth). Burst
10375            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
10376            // like gemma's burst arm.
10377            if adapt {
10378                let fl_now = floor_at(cache.pos);
10379                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
10380            }
10381            ph_mark(&mut ph_rest, phase_on);
10382            if let Some(p) = pipe {
10383                p.accept_end(round);
10384            }
10385            drop(pipe_accept);
10386            round += 1;
10387            // sse-cadence: this round's accepted drafts + bonus are committed (out is
10388            // append-only past step 4) — flush at round cadence.
10389            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10390        }
10391        if let Some(mut ticket) = carried_opti.take() {
10392            opti_fork
10393                .as_mut()
10394                .ok_or("optipipe tail drain lost fork state")?
10395                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
10396        }
10397        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
10398        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
10399        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
10400
10401        if spec_stats {
10402            let per_slot: Vec<String> = (0..k)
10403                .map(|j| {
10404                    if st_drafted[j] > 0 {
10405                        format!(
10406                            "{}/{}={:.3}",
10407                            st_accepted[j],
10408                            st_drafted[j],
10409                            st_accepted[j] as f64 / st_drafted[j] as f64
10410                        )
10411                    } else {
10412                        "0/0".into()
10413                    }
10414                })
10415                .collect();
10416            let acc = if total_drafted > 0 {
10417                total_accepted as f64 / total_drafted as f64
10418            } else {
10419                0.0
10420            };
10421            eprintln!(
10422                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
10423                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
10424                       tok_per_round={:.3}",
10425                per_slot.join(" "),
10426                (total_accepted + round) as f64 / round.max(1) as f64
10427            );
10428        }
10429        if constraint.is_some() {
10430            eprintln!(
10431                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
10432                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
10433                dm_clone_ns as f64 / 1e6,
10434                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
10435            );
10436        }
10437        if phase_on {
10438            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
10439            eprintln!(
10440                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
10441                ph_draft * 1e3,
10442                ph_draft / tot * 100.0,
10443                ph_verify * 1e3,
10444                ph_verify / tot * 100.0,
10445                ph_wait * 1e3,
10446                ph_wait / tot * 100.0,
10447                ph_rest * 1e3,
10448                ph_rest / tot * 100.0
10449            );
10450        }
10451        if anatomy_on {
10452            let rounds_f = round.max(1) as f64;
10453            let other = (ph_rest - ph_commit).max(0.0);
10454            eprintln!(
10455                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
10456                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
10457                ph_draft * 1e3 / rounds_f,
10458                ph_verify * 1e3 / rounds_f,
10459                ph_wait * 1e3 / rounds_f,
10460                ph_commit * 1e3 / rounds_f,
10461                other * 1e3 / rounds_f,
10462            );
10463        }
10464        let _pipe_tail = pipe.map(|p| p.primary());
10465        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
10466        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
10467        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
10468        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
10469        if let Some(slot) = sess_draft_slot.take() {
10470            *slot = Some(dctx);
10471        }
10472        let t_rounds = t_ent.elapsed();
10473        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
10474            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
10475            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
10476            // HERE, where the sampler, the session Philox counters and the penalty window are
10477            // all live and the boundary logits row still exists — that is the "make the state
10478            // available" half of the fix; the consuming burst then just emits it. `sctr` is
10479            // written to the session BELOW the draws so the advance is never lost.
10480            *next_pred_slot = Some(last_pred);
10481            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
10482            let mut stashed_pending = false;
10483            if let Some(b) = pending.take() {
10484                if !sampled {
10485                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
10486                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
10487                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
10488                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
10489                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
10490                    // OUT of `committed` (cache rows == committed); the consuming call
10491                    // prepends it once its verify commits the row. next_pred is unknowable
10492                    // without the commit pass — None; callers gate on pending_tok too.
10493                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
10494                    if let Some(slot) = sess_pending_slot.take() {
10495                        *slot = Some(b);
10496                    }
10497                    *next_pred_slot = None;
10498                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
10499                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
10500                    *last_h = Some(e.clone_dtod(&fill_prev)?);
10501                    stashed_pending = true;
10502                } else {
10503                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
10504                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
10505                    let pos_b = cache.pos;
10506                    scratch.set_len(e, pos_b)?;
10507                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
10508                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
10509                    // itself — the prediction AFTER the bonus never materialized; it would have
10510                    // been the next round's verify col 0). The commit's logits ARE that
10511                    // prediction — so they are also the row the next burst's boundary token
10512                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
10513                    *next_pred_slot = Some(if sample_boundary {
10514                        sample_boundary_token(
10515                            e,
10516                            &lg_b,
10517                            &sp,
10518                            &pen_hist,
10519                            &mut sctr,
10520                            "burst-tail-commit",
10521                        )?
10522                    } else {
10523                        argmax(&lg_b) as u32
10524                    });
10525                    self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
10526                    *last_h = Some(hb);
10527                }
10528            } else {
10529                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
10530                *last_h = Some(e.clone_dtod(&fill_prev)?);
10531                if sample_boundary {
10532                    // No pending to commit, so the boundary row is the one `last_pred` was
10533                    // argmaxed from and the sampled path keeps it on device: the init feed's
10534                    // logits when the burst ran zero rounds, else the legacy-replay path's
10535                    // last verify column (both predict the token AFTER the last committed
10536                    // row). It is retained precisely because round 0's accept test needs it,
10537                    // so the draw costs no extra D2H of the [n_vocab] row.
10538                    match last_col_logits.as_ref() {
10539                        Some(lc) => {
10540                            *next_pred_slot = Some(sample_boundary_token_dev(
10541                                e,
10542                                lc,
10543                                n_vocab,
10544                                &sp,
10545                                &pen_hist,
10546                                &mut sctr,
10547                                "burst-tail-nopending",
10548                            )?);
10549                        }
10550                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
10551                        // burst always feeds or replays, so the row exists — but if it ever
10552                        // is, the stream takes a greedy token and SAYS so rather than
10553                        // silently regressing to the pre-lane behaviour.
10554                        None => eprintln!(
10555                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
10556                             (reason: no retained boundary logits row)"
10557                        ),
10558                    }
10559                }
10560            }
10561            *sctr_slot = sctr;
10562            *uctr_slot = uctr;
10563            committed.extend_from_slice(prompt);
10564            if let Some(cb) = carried_pending {
10565                // the consumed carry's cache row landed in round 0's verify (every pending
10566                // round commits col 0) — it joins `committed` here, in sequence order.
10567                committed.push(cb);
10568            }
10569            if stashed_pending {
10570                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
10571                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
10572                // 18446744073709551615 out of range for slice of length 0", killing the
10573                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
10574                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
10575                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
10576                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
10577                // did). So a burst that stashes a pending without emitting anything of its own —
10578                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
10579                // guard skipping every token under a tight budget — arrives here with
10580                // out.len() == 0 and stashed_pending == true.
10581                //
10582                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
10583                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
10584                // just above is already accounted. Saturating, not a min/assert: an empty `out`
10585                // here is a legitimate burst shape, not a corrupt state.
10586                let emitted = out.len().saturating_sub(1);
10587                committed.extend_from_slice(&out[..emitted]);
10588            } else {
10589                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
10590            }
10591            debug_assert_eq!(
10592                cache.pos,
10593                committed.len(),
10594                "session invariant: cache rows == committed tokens"
10595            );
10596            if setup_trace {
10597                e.stream().synchronize()?; // bound the async tail fill in the trace
10598                let t_tail = t_ent.elapsed();
10599                eprintln!(
10600                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
10601                    t_init.as_secs_f64() * 1e3,
10602                    (t_cap - t_init).as_secs_f64() * 1e3,
10603                    (t_fill - t_cap).as_secs_f64() * 1e3,
10604                    (t_rounds - t_fill).as_secs_f64() * 1e3,
10605                    (t_tail - t_rounds).as_secs_f64() * 1e3,
10606                    t_tail.as_secs_f64() * 1e3,
10607                    out.len(),
10608                    continuation
10609                );
10610            }
10611            return Ok((out, total_drafted, total_accepted));
10612        }
10613        out.truncate(max_new);
10614        Ok((out, total_drafted, total_accepted))
10615    }
10616
10617    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
10618    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
10619    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
10620    pub fn extract_dspark_anchors(
10621        &self,
10622        e: &Engine,
10623        tokens: &[u32],
10624        anchor_positions: &[usize],
10625        gamma: usize,
10626        top_k: usize,
10627        chunk: usize,
10628        temperature: f32,
10629    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
10630        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
10631            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
10632        }
10633        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
10634            return Err("DSpark anchor positions must be sorted and unique".into());
10635        }
10636        for &position in anchor_positions {
10637            if position == 0 || position + gamma >= tokens.len() {
10638                return Err(format!(
10639                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
10640                    tokens.len()
10641                )
10642                .into());
10643            }
10644        }
10645
10646        let n_vocab = self.output.out_features();
10647        let n_embd = self.cfg.n_embd as usize;
10648        let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
10649        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10650        let embd_gpu = if spec_host_embd() {
10651            None
10652        } else {
10653            Some(
10654                self.embd_gpu
10655                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10656            )
10657        };
10658        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
10659
10660        struct PendingRecord {
10661            position: usize,
10662            hidden: Option<Vec<f32>>,
10663            tokens: Vec<u32>,
10664            target_top_ids: Vec<Option<Vec<u32>>>,
10665            target_top_logits: Vec<Option<Vec<f32>>>,
10666            target_top_probs: Vec<Option<Vec<f32>>>,
10667            target_tail_probs: Vec<Option<f32>>,
10668        }
10669
10670        let mut pending: Vec<PendingRecord> = anchor_positions
10671            .iter()
10672            .map(|&position| PendingRecord {
10673                position,
10674                hidden: None,
10675                tokens: tokens[position..=position + gamma].to_vec(),
10676                target_top_ids: vec![None; gamma],
10677                target_top_logits: vec![None; gamma],
10678                target_top_probs: vec![None; gamma],
10679                target_tail_probs: vec![None; gamma],
10680            })
10681            .collect();
10682
10683        let mut start = 0usize;
10684        while start < tokens.len() {
10685            let end = (start + chunk).min(tokens.len());
10686            let chunk_tokens = &tokens[start..end];
10687            let (target_logits, hidden_rows) =
10688                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
10689            for record in &mut pending {
10690                let hidden_position = record.position - 1;
10691                if hidden_position >= start && hidden_position < end {
10692                    let local = hidden_position - start;
10693                    record.hidden = Some(
10694                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
10695                    );
10696                }
10697                for slot in 0..gamma {
10698                    let target_row = record.position + slot;
10699                    if target_row < start || target_row >= end {
10700                        continue;
10701                    }
10702                    let local = target_row - start;
10703                    let logits =
10704                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
10705                    let (ids, top_logits, probs, tail) =
10706                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
10707                    record.target_top_ids[slot] = Some(ids);
10708                    record.target_top_logits[slot] = Some(top_logits);
10709                    record.target_top_probs[slot] = Some(probs);
10710                    record.target_tail_probs[slot] = Some(tail);
10711                }
10712            }
10713            start = end;
10714        }
10715
10716        pending
10717            .into_iter()
10718            .map(|record| {
10719                let hidden = record
10720                    .hidden
10721                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
10722                let target_top_ids =
10723                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
10724                let target_top_logits = flatten_dspark_rows(
10725                    record.target_top_logits,
10726                    record.position,
10727                    "target logits",
10728                )?;
10729                let target_top_probs =
10730                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
10731                let target_tail_probs = record
10732                    .target_tail_probs
10733                    .into_iter()
10734                    .enumerate()
10735                    .map(|(slot, value)| {
10736                        value.ok_or_else(|| {
10737                            format!("missing DSpark tail at {} slot {slot}", record.position)
10738                        })
10739                    })
10740                    .collect::<Result<Vec<_>, _>>()?;
10741                Ok(DsparkAnchorRecord {
10742                    position: record.position,
10743                    hidden,
10744                    tokens: record.tokens,
10745                    target_top_ids,
10746                    target_top_logits,
10747                    target_top_probs,
10748                    target_tail_probs,
10749                })
10750            })
10751            .collect()
10752    }
10753
10754    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
10755    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
10756    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
10757    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
10758    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
10759    /// quant-induced head/hidden-state mismatch from text drift.
10760    ///
10761    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
10762    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
10763    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
10764    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
10765    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
10766    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
10767    ///              conditions on the corpus — deterministic and arm-comparable by design.
10768    ///
10769    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
10770    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
10771    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
10772    ///
10773    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
10774    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
10775    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
10776    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
10777    /// agreement vs this path — not usable as a training-data source).
10778    pub fn replay_acceptance(
10779        &self,
10780        e: &Engine,
10781        tokens: &[u32],
10782        k: usize,
10783        stride: usize,
10784        chunk: usize,
10785        mut hdump: Option<&mut std::fs::File>,
10786    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
10787        assert!(k >= 1 && stride >= 1 && chunk >= 2);
10788        let mtp = self
10789            .mtp
10790            .as_ref()
10791            .expect("replay_acceptance requires an MTP head");
10792        let n_vocab = self.output.out_features();
10793        let d_vocab = mtp
10794            .shared_head_head
10795            .as_ref()
10796            .unwrap_or(&self.output)
10797            .out_features();
10798        let n_embd = self.cfg.n_embd as usize;
10799        let t_total = tokens.len();
10800        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
10801        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
10802        let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
10803        let mut scratch = MtpScratch::new(
10804            e,
10805            &self.cfg,
10806            t_total + k + 8,
10807            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
10808        )?;
10809        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10810        let embd_gpu = if spec_host_embd() {
10811            None
10812        } else {
10813            Some(
10814                self.embd_gpu
10815                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10816            )
10817        };
10818        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10819
10820        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
10821        let mut bg: Vec<u32> = vec![0; t_total + 1];
10822        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
10823        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
10824        let mut seed_buf = e.zeros(n_embd)?;
10825        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
10826        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
10827        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
10828        let mut s = 0usize;
10829        while s < t_total {
10830            let cend = (s + chunk).min(t_total);
10831            let tc = cend - s;
10832            let ch = &tokens[s..cend];
10833            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
10834            //    the chunk's true hiddens.
10835            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
10836            for j in 0..tc {
10837                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
10838            }
10839            let preds = e.dtoh_u32(&preds_d)?;
10840            for j in 0..tc {
10841                bg[s + j + 1] = preds[j];
10842            }
10843            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
10844            // checkpoint-quality metric (position j's logits score the GOLD next token).
10845            if nll_on {
10846                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
10847                if jmax > 0 {
10848                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
10849                    let rows: Vec<i32> = (0..jmax as i32).collect();
10850                    let idsd = e.htod_u32_v(&ids)?;
10851                    let rowsd = e.htod_i32(&rows)?;
10852                    let mut outd = e.zeros(jmax)?;
10853                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
10854                    for pr in e.dtoh(&outd)? {
10855                        nll_sum += -((pr.max(1e-30)) as f64).ln();
10856                        nll_cnt += 1;
10857                    }
10858                }
10859            }
10860            if let Some(f) = hdump.as_deref_mut() {
10861                use std::io::Write;
10862                let host: Vec<f32> = e.dtoh(&vx)?;
10863                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
10864                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
10865                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
10866                for v in &host[..tc * n_embd] {
10867                    let b = v.to_bits();
10868                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
10869                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
10870                }
10871                f.write_all(&bytes)?;
10872            }
10873            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
10874            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
10875            // per token saved; the forced trunk pass + hdump is all the mode needs).
10876            let chainless = stride > t_total;
10877            if chainless {
10878                e.copy_view_into(
10879                    &mut prev_last_h,
10880                    0,
10881                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
10882                    n_embd,
10883                )?;
10884                s = cend;
10885                continue;
10886            }
10887            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
10888            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
10889            let mut vxs = e.zeros(tc * n_embd)?;
10890            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
10891            if tc > 1 {
10892                e.copy_view_into(
10893                    &mut vxs,
10894                    n_embd,
10895                    &vx.slice(0..(tc - 1) * n_embd),
10896                    (tc - 1) * n_embd,
10897                )?;
10898            }
10899            scratch.set_len(e, s)?;
10900            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
10901            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
10902            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
10903            //    truncates those approximate appends before they can ever be read.
10904            let ps: Vec<usize> = (s..cend)
10905                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
10906                .collect();
10907            for &p in ps.iter().rev() {
10908                scratch.set_len(e, p)?;
10909                if p == s {
10910                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
10911                } else {
10912                    e.copy_view_into(
10913                        &mut seed_buf,
10914                        0,
10915                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
10916                        n_embd,
10917                    )?;
10918                }
10919                let mut e_tok = tokens[p];
10920                let mut d_seed = e.clone_dtod(&seed_buf)?;
10921                let mut drafts: Vec<u32> = Vec::with_capacity(k);
10922                for j in 0..k {
10923                    let (dl_d, h_nextn) = self.mtp_head_forward_dev(
10924                        e,
10925                        mtp,
10926                        e_tok,
10927                        &d_seed,
10928                        &mut scratch,
10929                        p + 1 + j,
10930                        embd_dev,
10931                        None, // acceptance-oracle walk: no grammar
10932                    )?;
10933                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
10934                    let idx = e.dtoh_u32_one(&tok_d)?;
10935                    let d = match &mtp.d2t {
10936                        Some(map) => map[idx as usize],
10937                        None => idx,
10938                    };
10939                    drafts.push(d);
10940                    e_tok = d;
10941                    d_seed = h_nextn;
10942                }
10943                // targets may live in a LATER chunk's bg — resolved after the walk.
10944                rows.push((p, drafts, Vec::new()));
10945            }
10946            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
10947            //    expect scratch.len == cend with exact rows).
10948            scratch.set_len(e, s)?;
10949            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
10950            e.copy_view_into(
10951                &mut prev_last_h,
10952                0,
10953                &vx.slice((tc - 1) * n_embd..tc * n_embd),
10954                n_embd,
10955            )?;
10956            s = cend;
10957        }
10958        for (p, drafts, targets) in rows.iter_mut() {
10959            for j in 0..drafts.len() {
10960                targets.push(bg[*p + 1 + j]);
10961            }
10962        }
10963        rows.sort_by_key(|r| r.0);
10964        if nll_cnt > 0 {
10965            let mean = nll_sum / nll_cnt as f64;
10966            println!(
10967                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
10968                mean.exp()
10969            );
10970        }
10971        Ok((rows, bg))
10972    }
10973}
10974
10975#[cfg(test)]
10976mod dspark_sparse_tests {
10977    use super::dspark_sparse_softmax_topk;
10978
10979    #[test]
10980    fn topk_keeps_full_softmax_mass_and_stable_ties() {
10981        let logits = [1.0f32, 3.0, 3.0, -2.0];
10982        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
10983        assert_eq!(ids, vec![1, 2]);
10984        assert_eq!(top_logits, vec![3.0, 3.0]);
10985        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
10986        let expected = 1.0 / denominator;
10987        assert!((probs[0] - expected).abs() < 1.0e-6);
10988        assert!((probs[1] - expected).abs() < 1.0e-6);
10989        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
10990        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
10991    }
10992}
10993
10994#[cfg(test)]
10995mod spec_replay_env_tests {
10996    use super::spec_replay_env_on;
10997
10998    #[test]
10999    fn replay_requires_literal_one() {
11000        assert!(!spec_replay_env_on(None));
11001        assert!(!spec_replay_env_on(Some("")));
11002        assert!(!spec_replay_env_on(Some("0")));
11003        assert!(!spec_replay_env_on(Some("true")));
11004        assert!(!spec_replay_env_on(Some("2")));
11005        assert!(spec_replay_env_on(Some("1")));
11006    }
11007}
11008
11009#[cfg(test)]
11010mod telem_tests {
11011    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
11012
11013    #[test]
11014    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
11015        let counters = SpecTelemetryCounters::default();
11016        for mask in [
11017            [true, true, true],
11018            [true, true, false],
11019            [true, false, false],
11020            [false, false, false],
11021        ] {
11022            let accepted = mask.iter().take_while(|&&value| value).count();
11023            counters.record_round(mask.len(), accepted);
11024        }
11025
11026        let snapshot = counters.snapshot();
11027        assert_eq!(
11028            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
11029            (4, 12, 6)
11030        );
11031        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
11032        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
11033        assert_eq!(snapshot.tau(), 1.5);
11034        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11035        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
11036    }
11037
11038    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
11039    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
11040    #[test]
11041    fn delta_isolates_burst_contribution() {
11042        let mut t = SpecTelemetry::default();
11043        // "previous request": 2 rounds of k=3, accepts 3 then 1.
11044        for (kr, na) in [(3usize, 3usize), (3, 1)] {
11045            t.rounds += 1;
11046            t.drafted += kr as u64;
11047            t.accepted += na as u64;
11048            for j in 0..kr {
11049                t.pos_drafted[j] += 1;
11050            }
11051            for j in 0..na {
11052                t.pos_accepted[j] += 1;
11053            }
11054        }
11055        let before = t;
11056        // "this burst": 1 round k=3, accepts 2.
11057        t.rounds += 1;
11058        t.drafted += 3;
11059        t.accepted += 2;
11060        for j in 0..3 {
11061            t.pos_drafted[j] += 1;
11062        }
11063        for j in 0..2 {
11064            t.pos_accepted[j] += 1;
11065        }
11066        let d = t.delta_since(&before);
11067        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
11068        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
11069        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
11070        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11071    }
11072
11073    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
11074    /// aggregation invariant.
11075    #[test]
11076    fn merge_accumulates_fieldwise() {
11077        let mut agg = SpecTelemetry::default();
11078        let mut d1 = SpecTelemetry {
11079            rounds: 2,
11080            drafted: 6,
11081            accepted: 4,
11082            ..Default::default()
11083        };
11084        d1.pos_drafted[0] = 2;
11085        d1.pos_accepted[0] = 2;
11086        let mut d2 = SpecTelemetry {
11087            rounds: 1,
11088            drafted: 3,
11089            accepted: 1,
11090            ..Default::default()
11091        };
11092        d2.pos_drafted[0] = 1;
11093        d2.pos_accepted[0] = 1;
11094        d2.pos_drafted[1] = 1;
11095        agg.merge(&d1);
11096        agg.merge(&d2);
11097        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
11098        assert_eq!(agg.pos_drafted[0], 3);
11099        assert_eq!(agg.pos_accepted[0], 3);
11100        assert_eq!(agg.pos_drafted[1], 1);
11101        assert_eq!(agg.pos_accepted[1], 0);
11102    }
11103
11104    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
11105    /// public metrics surface and must never publish a u64-wrapped garbage value.
11106    #[test]
11107    fn delta_saturates_never_wraps() {
11108        let small = SpecTelemetry {
11109            rounds: 1,
11110            drafted: 2,
11111            accepted: 1,
11112            ..Default::default()
11113        };
11114        let big = SpecTelemetry {
11115            rounds: 5,
11116            drafted: 15,
11117            accepted: 9,
11118            ..Default::default()
11119        };
11120        let d = small.delta_since(&big);
11121        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
11122    }
11123}
11124
11125#[cfg(test)]
11126mod opti_fork_tests {
11127    use super::{
11128        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
11129    };
11130
11131    #[test]
11132    fn controller_threshold_and_three_miss_breaker_are_exact() {
11133        let mut policy = OptiControllerPolicy {
11134            threshold: 0.7,
11135            consecutive_misses: 0,
11136            breaker_tripped: false,
11137        };
11138        assert!(!policy.admit(0.699_999));
11139        assert!(policy.admit(0.7));
11140        assert!(!policy.resolve(false));
11141        assert!(!policy.resolve(false));
11142        assert!(policy.resolve(false));
11143        assert!(policy.breaker_tripped);
11144        assert!(!policy.admit(1.0));
11145        assert!(
11146            !policy.resolve(true),
11147            "a resolved hit cannot re-arm a tripped request"
11148        );
11149        assert!(policy.breaker_tripped);
11150    }
11151
11152    #[test]
11153    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
11154        let mut policy = OptiControllerPolicy {
11155            threshold: 0.0,
11156            consecutive_misses: 0,
11157            breaker_tripped: false,
11158        };
11159        for _ in 0..16 {
11160            assert!(policy.admit(0.0));
11161            assert!(!policy.resolve(false));
11162        }
11163        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
11164            assert!(
11165                !policy.admit(invalid),
11166                "invalid q proxy must fail closed: {invalid}"
11167            );
11168        }
11169        assert!(!policy.breaker_tripped);
11170        assert_eq!(policy.consecutive_misses, 0);
11171    }
11172
11173    #[test]
11174    fn alternating_mode_flips_by_generation_not_round_parity() {
11175        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
11176        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
11177        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
11178        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
11179    }
11180
11181    #[test]
11182    fn live_generation_cannot_be_overwritten() {
11183        let mut tracker = OptiForkGenerationTracker::default();
11184        let g0 = tracker.reserve().unwrap();
11185        let g1 = tracker.reserve().unwrap();
11186        let err = tracker.reserve().unwrap_err().to_string();
11187        assert!(
11188            err.contains("still owns generation 0"),
11189            "unexpected error: {err}"
11190        );
11191        tracker.retire(g0).unwrap();
11192        let g2 = tracker.reserve().unwrap();
11193        assert_eq!((g2.id, g2.slot), (2, 0));
11194        tracker.retire(g1).unwrap();
11195        tracker.retire(g2).unwrap();
11196    }
11197
11198    #[test]
11199    fn teardown_rejects_a_stale_generation_tag() {
11200        let mut tracker = OptiForkGenerationTracker::default();
11201        let g0 = tracker.reserve().unwrap();
11202        tracker.retire(g0).unwrap();
11203        let err = tracker.retire(g0).unwrap_err().to_string();
11204        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
11205    }
11206}
11207
11208#[cfg(test)]
11209mod draft_graph_fallback_tests {
11210    use super::DraftGraphFallback;
11211
11212    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
11213    #[test]
11214    fn flip_is_loud_once_and_memoized_after() {
11215        let mut f = DraftGraphFallback::default();
11216        let line = f
11217            .mark_greedy("out of memory")
11218            .expect("first flip must return the warn line");
11219        assert!(
11220            line.contains("WARN"),
11221            "flip line must be warn-level: {line}"
11222        );
11223        assert!(
11224            line.contains("out of memory"),
11225            "flip line must carry the reason: {line}"
11226        );
11227        assert!(f.greedy_failed());
11228        // re-marking an already-failed graph is the memoization: quiet, still failed.
11229        assert!(f.mark_greedy("out of memory").is_none());
11230        assert!(f.greedy_failed());
11231        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
11232        assert!(!f.sampled_failed());
11233        let line_s = f
11234            .mark_sampled("capture unsupported")
11235            .expect("sampled flip is its own flip");
11236        assert!(
11237            line_s.contains("sampled"),
11238            "sampled flip names itself: {line_s}"
11239        );
11240        assert!(f.mark_sampled("capture unsupported").is_none());
11241    }
11242
11243    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
11244    /// and says so exactly when there was something to reset.
11245    #[test]
11246    fn reset_on_resume_clears_flags_and_logs_once() {
11247        let mut f = DraftGraphFallback::default();
11248        // clean session: resume is silent, nothing to reset.
11249        assert!(f.reset_on_resume().is_none());
11250        f.mark_greedy("oom").unwrap();
11251        f.mark_sampled("oom").unwrap();
11252        let note = f
11253            .reset_on_resume()
11254            .expect("a set flag must produce the reset note");
11255        assert!(
11256            note.contains("greedy+sampled"),
11257            "note names what was reset: {note}"
11258        );
11259        assert!(
11260            !f.greedy_failed() && !f.sampled_failed(),
11261            "both flags cleared"
11262        );
11263        // and the NEXT failure after a reset is a fresh flip — loud again.
11264        assert!(f.mark_greedy("oom again").is_some());
11265        let note2 = f.reset_on_resume().expect("greedy-only reset");
11266        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
11267    }
11268
11269    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
11270    /// they precede a fresh capture attempt whose own failure re-flips loudly.
11271    #[test]
11272    fn shape_change_clears_are_silent() {
11273        let mut f = DraftGraphFallback::default();
11274        f.mark_greedy("oom").unwrap();
11275        f.clear_greedy();
11276        assert!(!f.greedy_failed());
11277        f.mark_sampled("oom").unwrap();
11278        f.clear_sampled();
11279        assert!(!f.sampled_failed());
11280        // after a silent clear there is nothing left for resume to report.
11281        assert!(f.reset_on_resume().is_none());
11282    }
11283}
11284
11285/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
11286///
11287/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
11288/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
11289/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
11290/// than remembered.
11291#[cfg(test)]
11292mod sampled_graph_key_tests {
11293    use super::{SampledGraphKey, debug_t_pred0};
11294
11295    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
11296    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
11297        (k.seed, k.temp_bits, k.k)
11298    }
11299
11300    fn pure_temp_key() -> SampledGraphKey {
11301        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
11302        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
11303    }
11304
11305    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
11306    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
11307    #[test]
11308    fn vendor_filters_change_the_key() {
11309        let parked = pure_temp_key();
11310        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
11311        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
11312        assert_eq!(
11313            legacy_key(&parked),
11314            legacy_key(&vendor),
11315            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
11316        );
11317        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
11318        assert!(parked.pure_temp());
11319        assert!(!vendor.pure_temp());
11320    }
11321
11322    /// Each distribution-shaping field alone is enough to drop the parked graph.
11323    #[test]
11324    fn every_filter_field_is_keyed() {
11325        let base = pure_temp_key();
11326        for (what, other) in [
11327            (
11328                "top_k",
11329                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
11330            ),
11331            (
11332                "top_p",
11333                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
11334            ),
11335            (
11336                "min_p",
11337                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
11338            ),
11339            (
11340                "penalties",
11341                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
11342            ),
11343        ] {
11344            assert_ne!(base, other, "{what} must be part of the key");
11345            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
11346            assert_eq!(
11347                legacy_key(&base),
11348                legacy_key(&other),
11349                "{what} was invisible to the pre-fix key",
11350            );
11351        }
11352    }
11353
11354    /// The baked constants stay keyed (this half was always right — regression cover for it).
11355    #[test]
11356    fn baked_constants_stay_keyed() {
11357        let base = pure_temp_key();
11358        assert_ne!(
11359            base,
11360            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
11361            "seed"
11362        );
11363        assert_ne!(
11364            base,
11365            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
11366            "temp"
11367        );
11368        assert_ne!(
11369            base,
11370            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
11371            "k"
11372        );
11373        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
11374        assert_eq!(
11375            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
11376            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
11377        );
11378    }
11379
11380    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
11381    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
11382    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
11383    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
11384    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
11385    ///
11386    /// This test is the other end of that argument, asserted here rather than remembered in a
11387    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
11388    /// would silently become the unsound thing it is documented not to be.
11389    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
11390    #[test]
11391    fn seed_alone_still_rekeys_the_draft_graph() {
11392        let parked = pure_temp_key();
11393        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
11394        assert_ne!(
11395            parked, reseeded,
11396            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
11397             decision not to compare seed rests on exactly this",
11398        );
11399        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
11400        // because of a filter difference.
11401        assert!(parked.pure_temp() && reseeded.pure_temp());
11402    }
11403
11404    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
11405    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
11406    /// agree on the regime, so a graph that survives the drop is legal to launch.
11407    #[test]
11408    fn equal_keys_agree_on_the_regime() {
11409        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
11410        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
11411        assert_eq!(a, b);
11412        assert_eq!(a.pure_temp(), b.pure_temp());
11413        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
11414        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
11415        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
11416        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
11417    }
11418
11419    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
11420    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
11421    #[test]
11422    fn debug_print_survives_the_sampled_arm() {
11423        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
11424        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
11425        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
11426        // round 0 without a pending bonus still reports last_pred, in both arms.
11427        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
11428        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
11429        // greedy keeps the real prediction it always printed.
11430        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
11431        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
11432    }
11433}