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/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
179/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
180/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
181/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
182/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
183/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
184/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
185/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
186/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
187/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
188pub(crate) fn dspark_defer_readback_on() -> bool {
189    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
190    *ON.get_or_init(|| {
191        std::env::var("MEMRA_DSPARK_DEFER_READBACK")
192            .map(|v| v != "0")
193            .unwrap_or(true)
194    })
195}
196/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
197/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
198/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
199/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
200/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
201/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
202/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
203pub(crate) fn state_copy_batch_on() -> bool {
204    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
205    *ON.get_or_init(|| {
206        std::env::var("MEMRA_STATE_COPY_BATCH")
207            .map(|v| v != "0")
208            .unwrap_or(true)
209    })
210}
211/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §5 rank 1, bucketed verify graphs),
212/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: the dspark verify's LINEAR-layer
213/// runs replay per-(segment, vt) CUDA graphs — see [`DsparkVerifyGraphs`]. Requires the
214/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
215///
216/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20): exactness holds (ALL
217/// EXACT, accept lines byte-match the bank) and steady-state prompts gain +0.4..+1.5
218/// tok/s, but the gate-scale MEAN is flat (109.2 vs 109.4): the AUTO_FREE launch scan
219/// costs 25.6 us x 16 graph launches/round (~0.41 ms — most of the eager-launch
220/// savings) and the first generation pays ~80 lazy captures. The residual verify gap
221/// lives in the FULL-ATTENTION per-row section (fa partial memsets + appends), i.e.
222/// the fa exec-update extension, not here. Off until one of: node-count reduction
223/// (ctx-scratch transients), WithParams UPLOAD instantiation, or the full-verify
224/// single-graph — each re-gated by the same battery.
225pub(crate) fn dspark_verify_graph_on() -> bool {
226    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
227    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
228}
229
230/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
231///
232/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
233/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
234/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
235/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
236/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
237/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
238/// the flag crashed precisely the regime it exists to investigate.
239///
240/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
241/// indexing (an out-of-range pred there is a real bug and must still be loud).
242fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
243    if base == 0 {
244        return last_pred.to_string();
245    }
246    match preds.get(base - 1) {
247        Some(p) => p.to_string(),
248        // sampled: the greedy per-column argmax was never run for this round.
249        None => {
250            debug_assert!(
251                sampled,
252                "greedy spec: preds[{}] missing at base {base}",
253                base - 1
254            );
255            "n/a".to_string()
256        }
257    }
258}
259
260/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
261///
262/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
263/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
264/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
265/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
266/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
267/// not believe in — and `u * 0 < p` then accepts it unconditionally.
268///
269/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
270/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
271pub(crate) fn skey_probe() -> bool {
272    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
273    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
274}
275
276/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
277/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
278/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
279/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
280/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
281/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
282/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
283/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
284/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
285pub trait SpecConstraint {
286    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
287    /// masked argmax).
288    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
289    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
290    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
291    /// Is `tok` consumable in the CURRENT state?
292    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
293    /// Advance the state with an emitted token.
294    fn consume(&mut self, tok: u32) -> Result<(), String>;
295
296    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
297    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
298    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
299    // loose, research/constrained-full-20260803). These three methods let the engine mask the
300    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
301    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
302    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
303    // stays the correctness backstop and the emitted stream is unchanged by construction
304    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
305    // argmax; a cut slot is recomputed as the masked argmax either way).
306    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
307
308    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
309    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
310    fn draft_mask_enabled(&self) -> bool {
311        false
312    }
313    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
314    /// slot. Called once per spec round, before the first draft position.
315    fn draft_begin(&mut self) -> Result<(), String> {
316        Ok(())
317    }
318    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
319    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
320    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
321        Ok(None)
322    }
323    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
324    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
325    /// engine stops drafting; the token already pushed still goes through verify.
326    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
327        Ok(false)
328    }
329}
330
331/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
332/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
333/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
334/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
335/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
336/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
337/// verify emits the masked argmax as usual).
338fn upload_draft_mask(
339    e: &Engine,
340    c: &mut dyn SpecConstraint,
341    dst: &mut CudaSlice<u32>,
342    d2t: Option<&Vec<u32>>,
343    d_vocab: usize,
344    words: usize,
345) -> Result<bool, Box<dyn std::error::Error>> {
346    let Some(tw) = c
347        .draft_mask_words()
348        .map_err(|e2| format!("constraint: {e2}"))?
349    else {
350        return Ok(false);
351    };
352    let bit = |t: usize| -> bool {
353        let w = t >> 5;
354        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
355    };
356    let mut buf = vec![0u32; words];
357    match d2t {
358        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
359        Some(map) => {
360            for (i, &t) in map.iter().enumerate().take(d_vocab) {
361                if bit(t as usize) {
362                    buf[i >> 5] |= 1u32 << (i & 31);
363                }
364            }
365        }
366        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
367        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
368        None => {
369            let n = tw.len().min(words);
370            buf[..n].copy_from_slice(&tw[..n]);
371        }
372    }
373    if buf.iter().all(|w| *w == 0) {
374        return Ok(false);
375    }
376    e.htod_u32_into(dst, &buf)?;
377    Ok(true)
378}
379
380/// Keep the full token-embedding table in host memory and upload only the rows needed by each
381/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
382/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
383/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
384pub(crate) fn spec_host_embd() -> bool {
385    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
386    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
387}
388
389/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
390/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
391/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
392/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
393/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
394/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
395/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
396/// run-spec K=1..8 + acceptance identity arbitrate e2e).
397pub(crate) fn spec_fused_t() -> bool {
398    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
399    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
400    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
401    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
402    *F.get_or_init(|| {
403        std::env::var("MEMRA_SPEC_FUSED_T")
404            .map(|v| v != "0")
405            .unwrap_or(true)
406    })
407}
408
409/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
410/// Only call this on such buffers — the lean contract is "identical bytes by construction".
411fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
412    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
413}
414
415/// Scratch KV for the MTP block (one full-attn layer).
416///
417/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
418/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
419/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
420/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
421/// engine's "mtp_update" design). Entries come from two sources:
422///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
423///     hidden chain-approximate — the reference engine accepts the same);
424///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
425///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
426/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
427/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
428/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
429/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
430/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
431/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
432/// committed row across turns (the predecessor-pairing seed + fill anchor).
433/// Per-request sampling config for the sampled-spec serve path.
434#[derive(Clone, Copy, Debug)]
435pub struct SpecSampling {
436    pub temp: f32,
437    pub seed: u64,
438    pub top_k: i32,            // 0 = off
439    pub top_p: f32,            // 1.0 = off
440    pub min_p: f32,            // 0.0 = off
441    pub penalty_last_n: usize, // 0 = penalties off
442    pub penalty_repeat: f32,
443    pub penalty_freq: f32,
444    pub penalty_present: f32,
445}
446
447/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
448/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
449pub const SPEC_TELEM_POS: usize = 8;
450
451/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
452/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
453/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
454/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
455/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
456/// in NEITHER drafted nor accepted.
457#[derive(Clone, Copy, Default, Debug)]
458pub struct SpecTelemetry {
459    /// verify rounds completed (a round-stream burst counts each of its M rounds).
460    pub rounds: u64,
461    /// tokens drafted / accepted across all rounds.
462    pub drafted: u64,
463    pub accepted: u64,
464    /// how often draft position j (0-based within a round's chain) was offered / accepted.
465    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
466    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
467    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
468    pub pos_drafted: [u64; SPEC_TELEM_POS],
469    pub pos_accepted: [u64; SPEC_TELEM_POS],
470}
471
472impl SpecTelemetry {
473    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
474    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
475    /// a wrapped counter.
476    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
477        let mut d = SpecTelemetry {
478            rounds: self.rounds.saturating_sub(prev.rounds),
479            drafted: self.drafted.saturating_sub(prev.drafted),
480            accepted: self.accepted.saturating_sub(prev.accepted),
481            ..Default::default()
482        };
483        for j in 0..SPEC_TELEM_POS {
484            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
485            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
486        }
487        d
488    }
489    /// Fieldwise `self += d` — the worker's per-model aggregation.
490    pub fn merge(&mut self, d: &SpecTelemetry) {
491        self.rounds += d.rounds;
492        self.drafted += d.drafted;
493        self.accepted += d.accepted;
494        for j in 0..SPEC_TELEM_POS {
495            self.pos_drafted[j] += d.pos_drafted[j];
496            self.pos_accepted[j] += d.pos_accepted[j];
497        }
498    }
499
500    /// Mean accepted draft-prefix length per verify round (tau).
501    pub fn tau(&self) -> f64 {
502        if self.rounds > 0 {
503            self.accepted as f64 / self.rounds as f64
504        } else {
505            0.0
506        }
507    }
508}
509
510/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
511/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
512/// launch, synchronization, allocation, or ordering dependency to the numeric path.
513struct SpecTelemetryCounters {
514    rounds: AtomicU64,
515    drafted: AtomicU64,
516    accepted: AtomicU64,
517    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
518    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
519}
520
521impl Default for SpecTelemetryCounters {
522    fn default() -> Self {
523        Self {
524            rounds: AtomicU64::new(0),
525            drafted: AtomicU64::new(0),
526            accepted: AtomicU64::new(0),
527            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
528            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
529        }
530    }
531}
532
533impl SpecTelemetryCounters {
534    fn record_round(&self, drafted: usize, accepted: usize) {
535        debug_assert!(accepted <= drafted);
536        self.rounds.fetch_add(1, Ordering::Relaxed);
537        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
538        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
539        for counter in self.pos_drafted.iter().take(drafted) {
540            counter.fetch_add(1, Ordering::Relaxed);
541        }
542        for counter in self.pos_accepted.iter().take(accepted) {
543            counter.fetch_add(1, Ordering::Relaxed);
544        }
545    }
546
547    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
548    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
549    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
550        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
551        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
552        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
553    }
554
555    fn snapshot(&self) -> SpecTelemetry {
556        SpecTelemetry {
557            rounds: self.rounds.load(Ordering::Relaxed),
558            drafted: self.drafted.load(Ordering::Relaxed),
559            accepted: self.accepted.load(Ordering::Relaxed),
560            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
561            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
562        }
563    }
564}
565
566pub struct SpecSession {
567    pub(crate) cache: Cache,
568    pub(crate) scratch: MtpScratch,
569    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
570    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
571    /// session must count them. Callers render output from this, not from their own echo.
572    pub committed: Vec<u32>,
573    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
574    pub(crate) last_h: Option<CudaSlice<f32>>,
575    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
576    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
577    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
578    pub next_pred: Option<u32>,
579    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
580    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
581    pub sctr: u32,
582    pub uctr: u32,
583    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
584    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
585    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
586    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
587    /// research/spec-serving-20260801). None before the first turn; error paths drop it
588    /// (next burst recaptures — serve retires errored sessions anyway).
589    pub(crate) draft_ctx: Option<DraftGraphCtx>,
590    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
591    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
592    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
593    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
594    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
595    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
596    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
597    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
598    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
599    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
600    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
601    pub pending_tok: Option<u32>,
602    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
603    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
604    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
605    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
606    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
607    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
608    /// accounting the loop already does — no syncs, no allocation. NOTE a
609    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
610    /// diff with [`SpecTelemetry::delta_since`] around each burst.
611    telem: SpecTelemetryCounters,
612    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
613    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
614    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
615    /// prime, result lands in `boundary_capture`.
616    pub capture_at: Option<usize>,
617    /// The capture the last cold prime produced (see [`SpecBoundaryCapture`]). Worker takes it
618    /// post-burst to assemble the prefix entry. A failed capture is silent, like `turn_ckpt` —
619    /// publication just isn't available for that request.
620    pub boundary_capture: Option<SpecBoundaryCapture>,
621}
622impl SpecSession {
623    /// Context capacity of the session's caches (the server's ContextFull guard).
624    pub fn cache_max_ctx(&self) -> usize {
625        self.cache.max_ctx
626    }
627    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
628    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
629    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
630    /// the prime boundary), so no copy was taken at prime time.
631    pub fn cache_ref(&self) -> &Cache {
632        &self.cache
633    }
634    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
635    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
636    /// like the trunk KV — draft rows below the prompt end are append-only for the
637    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
638    /// committed length, never below the prime boundary, and the true-hidden refresh
639    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
640    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
641    /// prefix-addressable; the prefix cache already refuses that class end to end).
642    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
643        if self.scratch.kv.ring.is_some() {
644            return None;
645        }
646        Some((
647            &self.scratch.kv.k,
648            &self.scratch.kv.v,
649            self.scratch.kv.k_tok_bytes,
650            self.scratch.kv.v_tok_bytes,
651        ))
652    }
653    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
654    pub fn telemetry(&self) -> SpecTelemetry {
655        self.telem.snapshot()
656    }
657    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
658    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
659    /// `spec_rewind_to_checkpoint`.
660    pub fn rewind_pos(&self) -> Option<usize> {
661        self.turn_ckpt.as_ref().map(|c| c.pos)
662    }
663    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
664    pub fn rewind_is_resident(&self) -> bool {
665        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
666            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
667        })
668    }
669    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
670    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
671    /// session has never run a turn and has no prediction to hand over.
672    pub fn demote_ready(&self) -> bool {
673        self.pending_tok.is_none() && self.next_pred.is_some()
674    }
675    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
676    pub fn has_pending(&self) -> bool {
677        self.pending_tok.is_some()
678    }
679    /// Committed row count == cache rows (the session invariant), for the caller's own
680    /// `fed`-length cross-check at a handoff boundary.
681    pub fn committed_len(&self) -> usize {
682        self.committed.len()
683    }
684    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
685    /// cache + next-token prediction to the plain batched-decode path.
686    ///
687    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
688    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
689    /// tokenwise prime of the same `committed` sequence would have left it (that is the
690    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
691    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
692    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
693    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
694    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
695    /// a state indistinguishable from one the batched path produced itself: the batched tick
696    /// emits `next_pred`, feeds it into this same cache, and decodes on.
697    ///
698    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
699    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
700    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
701    /// path would silently skip a token.
702    ///
703    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
704    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
705    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
706    /// would mean an `mtp_kv_fill` over the whole committed history).
707    pub fn into_demoted(self) -> Option<(Cache, u32)> {
708        if self.pending_tok.is_some() {
709            return None;
710        }
711        let np = self.next_pred?;
712        debug_assert_eq!(
713            self.cache.pos,
714            self.committed.len(),
715            "demotion handoff: cache rows != committed tokens"
716        );
717        Some((self.cache, np))
718    }
719    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
720    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
721    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
722    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
723    pub fn reset_graph_fallback_on_resume(&mut self) {
724        if let Some(line) = self
725            .draft_ctx
726            .as_mut()
727            .and_then(|c| c.failed.reset_on_resume())
728        {
729            eprintln!("{line}");
730        }
731    }
732}
733
734/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
735///
736/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
737/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
738/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
739/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
740/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
741/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
742///
743/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
744/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
745/// position index, so it must be a real device COPY — that copy is the entire reason a spec
746/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
747/// below the boundary were written by this turn's fill and are never revisited (the per-round
748/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
749/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
750/// predecessor-pairing anchor the next prime's fill reads for its first row.
751///
752/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
753pub(crate) struct SpecCheckpoint {
754    snap: crate::cache::CacheSnapshot,
755    /// Committed length at the boundary (== cache.pos there, the session invariant).
756    pos: usize,
757    /// Pre-output_norm hidden of row `pos - 1`.
758    last_h: CudaSlice<f32>,
759}
760
761/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
762/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
763/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
764/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
765/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
766/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
767/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
768/// so the worker slices those from the live caches post-burst instead of copying at prime time.
769pub struct SpecBoundaryCapture {
770    pub snap: crate::cache::CacheSnapshot,
771    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
772    pub pos: usize,
773    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
774    pub logits: Vec<f32>,
775    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
776    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
777    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
778    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
779    pub last_h: Vec<f32>,
780}
781
782/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
783/// spec boundary capture carries for later restored-session fills. Failure is silent
784/// (`turn_ckpt` convention): the capture publishes without an anchor.
785fn capture_boundary_hidden(
786    e: &Engine,
787    h_rows: &CudaSlice<f32>,
788    pos: usize,
789    n_embd: usize,
790) -> Vec<f32> {
791    if pos == 0 || h_rows.len() < pos * n_embd {
792        return Vec::new();
793    }
794    let Ok(mut row) = e.uninit(n_embd) else {
795        return Vec::new();
796    };
797    if e.copy_view_into(
798        &mut row,
799        0,
800        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
801        n_embd,
802    )
803    .is_err()
804    {
805        return Vec::new();
806    }
807    e.dtoh(&row).unwrap_or_default()
808}
809
810/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
811/// Default ON: the token a burst emits at its own boundary is drawn from the request's
812/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
813/// every boundary) without touching greedy, which is byte-unaffected either way.
814pub fn spec_sampled_boundary_on() -> bool {
815    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
816    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
817}
818
819/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
820/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
821/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
822/// restores the pre-lane posture (each burst restarts the window from its own prompt
823/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
824/// must keep refusing penalized sampled prefix-cache restores, because the restored
825/// session's continuation burst is handed no prompt slice at all.
826pub fn spec_pen_session_on() -> bool {
827    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
828    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
829}
830
831/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
832/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
833/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
834/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
835/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
836/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
837pub fn spec_restore_republish_on() -> bool {
838    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
839    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
840}
841
842/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
843/// the argmax the pre-lane code would have emitted from the same row. This is how the
844/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
845fn spec_boundary_trace() -> bool {
846    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
847    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
848}
849
850/// llama-parity floor for the penalty window when the request does not ask for a bigger
851/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
852/// non-identity penalty, so this floor only matters to explicit small windows and to the
853/// CLI env path.
854const PEN_WINDOW_FLOOR: usize = 64;
855
856/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
857/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
858/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
859/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
860/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
861/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
862/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
863/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
864/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
865/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
866/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
867/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
868const PEN_WINDOW_MAX: usize = 8192;
869
870/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
871/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
872/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
873/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
874/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
875/// client actually asked us to penalize, where the pre-lane code had NOTHING.
876fn pen_window_seed(
877    session_committed: &[u32],
878    burst_prompt: &[u32],
879    penalty_last_n: usize,
880) -> Vec<u32> {
881    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
882    let take_prompt = burst_prompt.len().min(win);
883    let take_sess = (win - take_prompt).min(session_committed.len());
884    let mut hist = Vec::with_capacity(take_sess + take_prompt);
885    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
886    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
887    hist
888}
889
890/// Draw a BOUNDARY token from the target distribution the request asked for
891/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
892/// every burst boundary".
893///
894/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
895/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
896/// row after the last committed token on a continuation burst; the prefix-cache entry's
897/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
898/// regimes, so a sampled stream took a greedy token once per burst — measured, not
899/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
900/// customer asked for a sampled token, so this draws one.
901///
902/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
903/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
904/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
905/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
906/// composition means `sample_check`'s distributional oracle covers this draw too, and the
907/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
908///
909/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
910/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
911/// stream the accept walk uses — never a second, independently seeded stream (which would be
912/// a new distributional bug: two streams from one seed correlate wherever their counters
913/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
914/// to the cold session's own first draw from the same logits row, which is what preserves the
915/// sampled-hit lane's per-seed hit==cold byte identity.
916#[allow(clippy::too_many_arguments)]
917pub fn sample_boundary_token_dev(
918    e: &Engine,
919    logits: &CudaSlice<f32>,
920    n_vocab: usize,
921    sp: &SpecSampling,
922    pen_hist: &[u32],
923    sctr: &mut u32,
924    site: &str,
925) -> Result<u32, Box<dyn std::error::Error>> {
926    debug_assert!(
927        sp.temp > 0.0,
928        "boundary sampling is the sampled regime only"
929    );
930    // Own copy: penalize_logits mutates in place and the caller's row is live state
931    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
932    let mut col = e.zeros(n_vocab)?;
933    e.copy_into(&mut col, 0, logits, n_vocab)?;
934    let pen_on = sp.penalty_last_n > 0
935        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
936    if pen_on && !pen_hist.is_empty() {
937        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
938        let w0 = pen_hist
939            .len()
940            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
941        let hist = &pen_hist[w0..];
942        let hd = e.htod_u32_v(hist)?;
943        e.penalize_logits(
944            &mut col,
945            &hd,
946            hist.len(),
947            sp.penalty_repeat,
948            sp.penalty_freq,
949            sp.penalty_present,
950            n_vocab,
951        )?;
952    }
953    let rows0 = e.htod_i32(&[0])?;
954    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
955    e.filter_stats(
956        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
957        sp.top_p, sp.min_p,
958    )?;
959    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
960    let mut perturb = e.zeros(n_vocab)?;
961    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
962    *sctr = sctr.wrapping_add(1);
963    let td = e.argmax_token_device(&perturb, n_vocab)?;
964    let tok = e.dtoh_u32_one(&td)?;
965    if spec_boundary_trace() {
966        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
967        let raw = e.argmax_token_device(logits, n_vocab)?;
968        let greedy = e.dtoh_u32_one(&raw)?;
969        eprintln!(
970            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
971             deviates={} temp={} sctr={}",
972            (tok != greedy) as u8,
973            sp.temp,
974            sctr.wrapping_sub(1),
975        );
976    }
977    Ok(tok)
978}
979
980/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
981/// host `Vec<f32>`).
982#[allow(clippy::too_many_arguments)]
983pub fn sample_boundary_token(
984    e: &Engine,
985    logits: &[f32],
986    sp: &SpecSampling,
987    pen_hist: &[u32],
988    sctr: &mut u32,
989    site: &str,
990) -> Result<u32, Box<dyn std::error::Error>> {
991    let n_vocab = logits.len();
992    let d = e.htod(logits)?;
993    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
994}
995
996struct SpecPipeTraceClock {
997    pair: usize,
998    started: std::time::Instant,
999}
1000
1001#[derive(Clone)]
1002struct SpecPipeTraceCtx {
1003    clock: std::sync::Arc<SpecPipeTraceClock>,
1004    round: usize,
1005    lane: usize,
1006}
1007
1008struct SpecPipeTraceMarker {
1009    trace: SpecPipeTraceCtx,
1010    phase: &'static str,
1011    edge: &'static str,
1012    slot: Option<usize>,
1013}
1014
1015unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1016    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1017    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1018    let slot = marker
1019        .slot
1020        .map(|v| v.to_string())
1021        .unwrap_or_else(|| "-".into());
1022    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1023    use std::io::Write as _;
1024    let stderr = std::io::stderr();
1025    let mut stderr = stderr.lock();
1026    let _ = writeln!(
1027        stderr,
1028        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1029         slot={slot} t_ms={t_ms:.3}",
1030        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1031    );
1032}
1033
1034fn enqueue_spec_pipe_trace_marker(
1035    stream: &cudarc::driver::CudaStream,
1036    trace: Option<&SpecPipeTraceCtx>,
1037    phase: &'static str,
1038    edge: &'static str,
1039    slot: Option<usize>,
1040) -> Result<(), Box<dyn std::error::Error>> {
1041    let Some(trace) = trace else {
1042        return Ok(());
1043    };
1044    let marker = Box::new(SpecPipeTraceMarker {
1045        trace: trace.clone(),
1046        phase,
1047        edge,
1048        slot,
1049    });
1050    let raw = Box::into_raw(marker);
1051    let result = unsafe {
1052        cudarc::driver::result::stream::launch_host_function(
1053            stream.cu_stream(),
1054            spec_pipe_trace_marker,
1055            raw.cast(),
1056        )
1057    };
1058    if let Err(err) = result {
1059        unsafe {
1060            drop(Box::from_raw(raw));
1061        }
1062        return Err(err.into());
1063    }
1064    Ok(())
1065}
1066
1067#[derive(Default)]
1068struct SpecPipeProgress {
1069    setup_done: [bool; 2],
1070    draft_done: [usize; 2],
1071    stage0_done: [usize; 2],
1072    verify_done: [usize; 2],
1073    accept_done: [usize; 2],
1074    finished: [bool; 2],
1075    aborted: bool,
1076}
1077
1078/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1079/// keeps its existing call stack and round locals; this object only orders phase entry. The
1080/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1081/// cannot be interleaved by the two host threads.
1082struct SpecPipeSync {
1083    progress: std::sync::Mutex<SpecPipeProgress>,
1084    changed: std::sync::Condvar,
1085    primary: std::sync::Mutex<()>,
1086    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1087}
1088
1089impl SpecPipeSync {
1090    fn new() -> Self {
1091        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1092        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1093            std::sync::Arc::new(SpecPipeTraceClock {
1094                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1095                started: std::time::Instant::now(),
1096            })
1097        });
1098        Self {
1099            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1100            changed: std::sync::Condvar::new(),
1101            primary: std::sync::Mutex::new(()),
1102            trace,
1103        }
1104    }
1105}
1106
1107#[derive(Clone)]
1108struct SpecPipeLane {
1109    sync: std::sync::Arc<SpecPipeSync>,
1110    lane: usize,
1111}
1112
1113impl SpecPipeLane {
1114    fn peer(&self) -> usize {
1115        1 - self.lane
1116    }
1117
1118    fn aborted() -> Box<dyn std::error::Error> {
1119        "paired speculative peer aborted".into()
1120    }
1121
1122    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1123        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1124            clock: clock.clone(),
1125            round,
1126            lane: self.lane,
1127        })
1128    }
1129
1130    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1131        let mut p = self.sync.progress.lock().unwrap();
1132        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1133            p = self.sync.changed.wait(p).unwrap();
1134        }
1135        if p.aborted {
1136            Err(Self::aborted())
1137        } else {
1138            Ok(())
1139        }
1140    }
1141
1142    fn setup_end(&self) {
1143        let mut p = self.sync.progress.lock().unwrap();
1144        p.setup_done[self.lane] = true;
1145        self.sync.changed.notify_all();
1146    }
1147
1148    fn draft_begin(
1149        &self,
1150        round: usize,
1151    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1152        let peer = self.peer();
1153        let mut p = self.sync.progress.lock().unwrap();
1154        loop {
1155            if p.aborted {
1156                return Err(Self::aborted());
1157            }
1158            let setup_ready =
1159                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1160            let prior_ready = p.accept_done[self.lane] >= round
1161                && (p.accept_done[peer] >= round || p.finished[peer]);
1162            let turn_ready = if self.lane == 0 {
1163                true
1164            } else {
1165                p.draft_done[0] > round || p.finished[0]
1166            };
1167            if setup_ready && prior_ready && turn_ready {
1168                break;
1169            }
1170            p = self.sync.changed.wait(p).unwrap();
1171        }
1172        drop(p);
1173        Ok(self.sync.primary.lock().unwrap())
1174    }
1175
1176    fn draft_end(&self, round: usize) {
1177        let mut p = self.sync.progress.lock().unwrap();
1178        p.draft_done[self.lane] = round + 1;
1179        self.sync.changed.notify_all();
1180    }
1181
1182    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1183    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1184    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1185        let peer = self.peer();
1186        let mut p = self.sync.progress.lock().unwrap();
1187        loop {
1188            if p.aborted {
1189                return Err(Self::aborted());
1190            }
1191            let ready = if self.lane == 0 {
1192                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1193            } else {
1194                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1195            };
1196            if ready {
1197                return Ok(self.lane == 0 || p.finished[peer]);
1198            }
1199            p = self.sync.changed.wait(p).unwrap();
1200        }
1201    }
1202
1203    fn stage0_end(&self, round: usize) {
1204        let mut p = self.sync.progress.lock().unwrap();
1205        p.stage0_done[self.lane] = round + 1;
1206        self.sync.changed.notify_all();
1207    }
1208
1209    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1210    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1211    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1212        let mut p = self.sync.progress.lock().unwrap();
1213        while !p.aborted
1214            && !(p.stage0_done[self.lane] > round
1215                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1216        {
1217            p = self.sync.changed.wait(p).unwrap();
1218        }
1219        if p.aborted {
1220            Err(Self::aborted())
1221        } else {
1222            Ok(())
1223        }
1224    }
1225
1226    fn verify_end(&self, round: usize) {
1227        let mut p = self.sync.progress.lock().unwrap();
1228        p.verify_done[self.lane] = round + 1;
1229        self.sync.changed.notify_all();
1230    }
1231
1232    fn accept_begin(
1233        &self,
1234        round: usize,
1235    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1236        let mut p = self.sync.progress.lock().unwrap();
1237        loop {
1238            if p.aborted {
1239                return Err(Self::aborted());
1240            }
1241            let ready = if self.lane == 0 {
1242                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1243            } else {
1244                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1245            };
1246            if ready {
1247                break;
1248            }
1249            p = self.sync.changed.wait(p).unwrap();
1250        }
1251        drop(p);
1252        Ok(self.sync.primary.lock().unwrap())
1253    }
1254
1255    fn accept_end(&self, round: usize) {
1256        let mut p = self.sync.progress.lock().unwrap();
1257        p.accept_done[self.lane] = round + 1;
1258        self.sync.changed.notify_all();
1259    }
1260
1261    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1262        self.sync.primary.lock().unwrap()
1263    }
1264
1265    fn finish(&self, failed: bool) {
1266        let mut p = self.sync.progress.lock().unwrap();
1267        p.finished[self.lane] = true;
1268        p.aborted |= failed;
1269        self.sync.changed.notify_all();
1270    }
1271}
1272
1273struct SpecPipeFinish<'a> {
1274    lane: &'a SpecPipeLane,
1275    closed: bool,
1276}
1277
1278impl<'a> SpecPipeFinish<'a> {
1279    fn new(lane: &'a SpecPipeLane) -> Self {
1280        Self {
1281            lane,
1282            closed: false,
1283        }
1284    }
1285
1286    fn close(&mut self, failed: bool) {
1287        self.lane.finish(failed);
1288        self.closed = true;
1289    }
1290}
1291
1292impl Drop for SpecPipeFinish<'_> {
1293    fn drop(&mut self) {
1294        if !self.closed {
1295            self.lane.finish(true);
1296        }
1297    }
1298}
1299
1300/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1301/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1302/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1303/// binds that context before touching the session, joins before returning, and never aliases the
1304/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1305/// session type Send.
1306struct SpecPipeSessionPtr(*mut SpecSession);
1307
1308unsafe impl Send for SpecPipeSessionPtr {}
1309
1310impl SpecPipeSessionPtr {
1311    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1312        unsafe { &mut *self.0 }
1313    }
1314}
1315
1316/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1317/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1318/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1319/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1320/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1321/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1322/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1323/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1324/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1325///
1326/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1327/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1328/// load-bearing:
1329///
1330/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1331///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1332///   This is all the key used to carry.
1333/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1334///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1335///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1336///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1337///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1338///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1339///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1340///
1341/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1342/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1343/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1344/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1345/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1346#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1347pub(crate) struct SampledGraphKey {
1348    seed: u64,
1349    temp_bits: u32,
1350    k: usize,
1351    top_k: i32,
1352    top_p_bits: u32,
1353    min_p_bits: u32,
1354    pen_on: bool,
1355}
1356
1357impl SampledGraphKey {
1358    pub(crate) fn new(
1359        seed: u64,
1360        temp: f32,
1361        k: usize,
1362        top_k: i32,
1363        top_p: f32,
1364        min_p: f32,
1365        pen_on: bool,
1366    ) -> Self {
1367        SampledGraphKey {
1368            seed,
1369            temp_bits: temp.to_bits(),
1370            k,
1371            top_k,
1372            top_p_bits: top_p.to_bits(),
1373            min_p_bits: min_p.to_bits(),
1374            pen_on,
1375        }
1376    }
1377
1378    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1379    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1380    /// the key can never drift apart (they were three separate expressions before this lane, and
1381    /// the launch site simply forgot to ask).
1382    pub(crate) fn pure_temp(&self) -> bool {
1383        self.top_k == 0
1384            && f32::from_bits(self.top_p_bits) >= 1.0
1385            && f32::from_bits(self.min_p_bits) <= 0.0
1386            && !self.pen_on
1387    }
1388}
1389
1390pub(crate) struct DraftGraphCtx {
1391    g_tok: CudaSlice<u32>,
1392    g_pos: CudaSlice<i32>,
1393    g_seed: CudaSlice<f32>,
1394    g_p: CudaSlice<f32>,
1395    g_ctr: CudaSlice<u32>,
1396    g_q: CudaSlice<f32>,
1397    g_perturb: CudaSlice<f32>,
1398    q_slots: Vec<CudaSlice<f32>>,
1399    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1400    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1401    /// per-position contents the host re-uploads before each replay (the graph-promote
1402    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1403    g_dmask: CudaSlice<u32>,
1404    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1405    graph_masked: bool,
1406    graph: Option<cudarc::driver::CudaGraph>,
1407    graph_s: Option<cudarc::driver::CudaGraph>,
1408    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1409    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1410    failed: DraftGraphFallback,
1411    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1412    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1413    s_key: Option<SampledGraphKey>,
1414    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1415    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1416    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1417    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1418    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1419    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1420    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1421    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1422    keeper: Vec<Box<dyn std::any::Any + Send>>,
1423    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1424}
1425
1426/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1427/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1428///
1429/// Three contracts:
1430/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1431///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1432///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1433///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1434///   fallback from paying a doomed capture attempt every burst).
1435/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1436///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1437///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1438///   actually set (quiet on the common clean-resume path).
1439/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1440///   capture attempt whose own failure would re-flip loudly.
1441#[derive(Default)]
1442pub(crate) struct DraftGraphFallback {
1443    greedy: bool,
1444    sampled: bool,
1445}
1446impl DraftGraphFallback {
1447    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1448        if self.greedy {
1449            return None;
1450        }
1451        self.greedy = true;
1452        Some(format!(
1453            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1454        ))
1455    }
1456    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1457        if self.sampled {
1458            return None;
1459        }
1460        self.sampled = true;
1461        Some(format!(
1462            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1463        ))
1464    }
1465    fn greedy_failed(&self) -> bool {
1466        self.greedy
1467    }
1468    fn sampled_failed(&self) -> bool {
1469        self.sampled
1470    }
1471    fn clear_greedy(&mut self) {
1472        self.greedy = false;
1473    }
1474    fn clear_sampled(&mut self) {
1475        self.sampled = false;
1476    }
1477    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1478    /// was set (so clean resumes stay quiet).
1479    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1480        if !self.greedy && !self.sampled {
1481            return None;
1482        }
1483        let which = match (self.greedy, self.sampled) {
1484            (true, true) => "greedy+sampled",
1485            (true, false) => "greedy",
1486            _ => "sampled",
1487        };
1488        self.greedy = false;
1489        self.sampled = false;
1490        Some(format!(
1491            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1492        ))
1493    }
1494}
1495
1496impl DraftGraphCtx {
1497    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1498        Ok(DraftGraphCtx {
1499            g_tok: e.alloc_u32_zeroed(1)?,
1500            g_pos: e.htod_i32(&[0])?,
1501            g_seed: e.zeros(n_embd)?,
1502            g_p: e.zeros(1)?,
1503            g_ctr: e.alloc_u32_zeroed(1)?,
1504            g_q: e.zeros(qlen)?,
1505            g_perturb: e.zeros(qlen)?,
1506            q_slots: Vec::new(),
1507            g_dmask: e.alloc_u32_zeroed(1)?,
1508            graph_masked: false,
1509            graph: None,
1510            graph_s: None,
1511            failed: DraftGraphFallback::default(),
1512            s_key: None,
1513            keeper: Vec::new(),
1514            keeper_s: Vec::new(),
1515        })
1516    }
1517}
1518
1519pub(crate) struct MtpScratch {
1520    kv: KvLayer,
1521    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1522    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1523    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1524    /// smaller host-indexed SWA ring instead.
1525    cap: usize,
1526}
1527
1528fn mtp_scratch_layout(
1529    cfg: &memra_gguf::config::ModelConfig,
1530    geom: Option<&crate::hybrid::DraftGeom>,
1531) -> (usize, usize, usize, usize) {
1532    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1533    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1534    let head_dim_k = cfg.head_dim_k as usize;
1535    let head_dim_v = cfg.head_dim_v as usize;
1536    assert!(
1537        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1538        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1539    );
1540    let kv_dim_k = head_dim_k * n_head_kv;
1541    let kv_dim_v = head_dim_v * n_head_kv;
1542    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1543    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1544    let (kbb, vbb) = crate::kv_blk_bytes();
1545    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1546    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1547    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1548}
1549
1550impl MtpScratch {
1551    fn new(
1552        e: &Engine,
1553        cfg: &memra_gguf::config::ModelConfig,
1554        cap: usize,
1555        geom: Option<&crate::hybrid::DraftGeom>,
1556    ) -> Result<Self, Box<dyn std::error::Error>> {
1557        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1558        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1559        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1560        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1561        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1562        let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1563            let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1564            Some(crate::cache::KvRing::new(
1565                crate::cache::swa_ring_rows(window, cap),
1566                window,
1567            ))
1568        } else {
1569            None
1570        };
1571        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1572        Ok(MtpScratch {
1573            kv: KvLayer {
1574                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1575                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1576                kv_dim_k,
1577                kv_dim_v,
1578                k_tok_bytes,
1579                v_tok_bytes,
1580                len: 0,
1581                ring,
1582                len_d: e.htod_i32(&[0])?,
1583            },
1584            cap,
1585        })
1586    }
1587    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1588    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1589    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1590    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1591        if self
1592            .kv
1593            .ring
1594            .as_ref()
1595            .is_some_and(|ring| !ring.can_rewind_to(n))
1596        {
1597            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1598        }
1599        self.kv.len = n;
1600        e.set_i32_one(&mut self.kv.len_d, n as i32)
1601    }
1602
1603    fn can_rewind_to(&self, n: usize) -> bool {
1604        self.kv
1605            .ring
1606            .as_ref()
1607            .is_none_or(|ring| ring.can_rewind_to(n))
1608    }
1609}
1610
1611/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1612/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1613/// full weight reads per round — recomputing columns the verify had already produced
1614/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1615/// to "after the first j verify columns" WITHOUT re-running the trunk:
1616/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1617///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1618///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1619///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1620///   pure-copy ring rebuild.
1621/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1622///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1623///   target: j <= t-1).
1624/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1625/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1626struct GdnStash {
1627    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1628    q_l2: CudaSlice<f32>,
1629    k_l2: CudaSlice<f32>,
1630    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1631    g_log: CudaSlice<f32>,
1632    beta: CudaSlice<f32>, // [t, num_v]
1633}
1634struct VerifyCkpt {
1635    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1636    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1637}
1638/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1639pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1640
1641/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1642/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1643/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1644/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1645/// layers between full-attention layers are shape-static given vt — no positions, no
1646/// t_kv, state addressed through pointer tables — so runs of them capture per
1647/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1648/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1649///
1650/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1651/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1652/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1653/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1654/// before and restored after — the graph's first real launch starts from the exact
1655/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1656/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1657/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1658pub(crate) struct DsparkVerifyGraphs {
1659    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1660    lin: Vec<usize>,
1661    lin_pos: std::collections::HashMap<usize, usize>,
1662    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1663    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1664    table_all: CudaSlice<u64>,
1665    host_table: Vec<u64>,
1666    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1667    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1668    stash_conv: Vec<CudaSlice<f32>>,
1669    stash_ssm: Vec<CudaSlice<f32>>,
1670    conv_words: usize,
1671    ssm_words: usize,
1672    /// Per-vt input/output staging (stable addresses the graphs bake).
1673    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1674    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1675    /// so the sink buffer must live (and persist) with the graphs, not with the round.
1676    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1677    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1678    /// Warmup-corruption guard scratch: pre-capture conv/ssm of one segment.
1679    save_conv: CudaSlice<f32>,
1680    save_ssm: CudaSlice<f32>,
1681    max_run: usize,
1682    n_embd: usize,
1683    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1684    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1685    pub(crate) round_slab: bool,
1686}
1687
1688struct DsparkSegGraph {
1689    graph: cudarc::driver::CudaGraph,
1690    _keeper: Vec<Box<dyn std::any::Any + Send>>,
1691}
1692
1693// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
1694// no automatic trait; CUDA driver graph handles are context-scoped rather than
1695// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
1696// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
1697// single decode-stream thread.
1698unsafe impl Send for DsparkVerifyGraphs {}
1699
1700impl DsparkVerifyGraphs {
1701    /// Build for this cache's shape. None when there are no linear layers, sizes are
1702    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
1703    pub(crate) fn new(
1704        e: &Engine,
1705        cache: &Cache,
1706        t_max: usize,
1707        n_embd: usize,
1708    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
1709        let lin: Vec<usize> = (0..cache.recur.len())
1710            .filter(|&il| cache.recur[il].is_some())
1711            .collect();
1712        if lin.is_empty() || t_max < 2 {
1713            return Ok(None);
1714        }
1715        let first = cache.recur[lin[0]].as_ref().unwrap();
1716        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
1717        for &il in &lin {
1718            let rl = cache.recur[il].as_ref().unwrap();
1719            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
1720                return Ok(None);
1721            }
1722        }
1723        let n = lin.len();
1724        let mut lin_pos = std::collections::HashMap::with_capacity(n);
1725        for (k, &il) in lin.iter().enumerate() {
1726            lin_pos.insert(il, k);
1727        }
1728        // longest run of consecutive linear layers (save-scratch sizing)
1729        let mut max_run = 1usize;
1730        let mut run = 1usize;
1731        for w in lin.windows(2) {
1732            if w[1] == w[0] + 1 {
1733                run += 1;
1734                max_run = max_run.max(run);
1735            } else {
1736                run = 1;
1737            }
1738        }
1739        let rows = t_max - 1;
1740        let mut stash_conv = Vec::with_capacity(n);
1741        let mut stash_ssm = Vec::with_capacity(n);
1742        for _ in 0..n {
1743            stash_conv.push(e.uninit(rows * conv_words)?);
1744            stash_ssm.push(e.uninit(rows * ssm_words)?);
1745        }
1746        let host_table = vec![0u64; n * 6];
1747        let table_all = e.htod_u64(&host_table)?;
1748        Ok(Some(Self {
1749            lin,
1750            lin_pos,
1751            table_all,
1752            host_table,
1753            stash_conv,
1754            stash_ssm,
1755            conv_words,
1756            ssm_words,
1757            stage: std::collections::HashMap::new(),
1758            tap_bufs: std::collections::HashMap::new(),
1759            graphs: std::collections::HashMap::new(),
1760            save_conv: e.uninit(max_run * conv_words)?,
1761            save_ssm: e.uninit(max_run * ssm_words)?,
1762            max_run,
1763            n_embd,
1764            round_slab: false,
1765        }))
1766    }
1767
1768    /// Rebuild the pointer table from the live handles (once per verify — the gdn
1769    /// ping-pong swaps the canonical/alt handles between rounds; a stale table would
1770    /// read the wrong parity's state).
1771    pub(crate) fn refresh_tables(
1772        &mut self,
1773        e: &Engine,
1774        cache: &Cache,
1775    ) -> Result<(), Box<dyn std::error::Error>> {
1776        use cudarc::driver::DevicePtr;
1777        {
1778            let s = &e.gpu.stream();
1779            for (k, &il) in self.lin.iter().enumerate() {
1780                let rl = cache.recur[il].as_ref().unwrap();
1781                let (pc, _g0) = rl.conv_state.device_ptr(s);
1782                let (p0, _g1) = rl.ssm_state.device_ptr(s);
1783                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
1784                let o = k * 6;
1785                self.host_table[o] = pc as u64;
1786                self.host_table[o + 1] = p0 as u64;
1787                self.host_table[o + 2] = p1 as u64;
1788                self.host_table[o + 3] = pc as u64;
1789                self.host_table[o + 4] = p1 as u64;
1790                self.host_table[o + 5] = p0 as u64;
1791            }
1792        }
1793        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
1794        Ok(())
1795    }
1796
1797    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
1798    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
1799    /// bracketed by a segment state save/restore), launch, then apply the host parity
1800    /// bookkeeping the captured body would have done. Returns the fresh residual.
1801    #[allow(clippy::too_many_arguments)]
1802    fn run_segment(
1803        &mut self,
1804        model: &crate::hybrid::HybridModel,
1805        e: &Engine,
1806        start: usize,
1807        end: usize,
1808        x: &CudaSlice<f32>,
1809        t: usize,
1810        cache: &mut Cache,
1811    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1812        let n_embd = self.n_embd;
1813        debug_assert!(end - start <= self.max_run);
1814        if !self.stage.contains_key(&t) {
1815            let xin = e.uninit(t * n_embd)?;
1816            let xout = e.uninit(t * n_embd)?;
1817            self.stage.insert(t, (xin, xout));
1818        }
1819        // Stage the residual at the bucket's baked input address.
1820        {
1821            let (xin, _) = self.stage.get_mut(&t).unwrap();
1822            e.copy_into(xin, 0, x, t * n_embd)?;
1823        }
1824        let key = (start, t);
1825        if !self.graphs.contains_key(&key) {
1826            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
1827            // ssm of every segment layer first, restore after, so the graph's first real
1828            // launch starts from the exact pre-round state (bytes gated e2e).
1829            for (k, il) in (start..end).enumerate() {
1830                let rl = cache.recur[il].as_ref().unwrap();
1831                e.copy_into(
1832                    &mut self.save_conv,
1833                    k * self.conv_words,
1834                    &rl.conv_state,
1835                    self.conv_words,
1836                )?;
1837                e.copy_into(
1838                    &mut self.save_ssm,
1839                    k * self.ssm_words,
1840                    &rl.ssm_state,
1841                    self.ssm_words,
1842                )?;
1843            }
1844            let (graph, keeper) = {
1845                let table_all = &self.table_all;
1846                let lin_pos = &self.lin_pos;
1847                let stash_conv = &mut self.stash_conv;
1848                let stash_ssm = &mut self.stash_ssm;
1849                let (xin, xout) = self
1850                    .stage
1851                    .get_mut(&t)
1852                    .map(|(a, b)| (&*a, b))
1853                    .expect("stage bucket created above");
1854                let cache_ref: &mut Cache = cache;
1855                // AUTO_FREE_ON_LAUNCH (the retained default): UPLOAD via
1856                // cuGraphInstantiateWithFlags is CUDA_ERROR_INVALID_VALUE (the flag is
1857                // WithParams-only), and the alloc nodes need the auto-free semantics.
1858                // Its launch-time mem-pool scan is the measured limiter — 25.6 us per
1859                // cuGraphLaunch x 16 segments = ~0.41 ms/round, most of the
1860                // eager-launch savings — which is why this door is OPT-IN until the
1861                // node count drops (ctx-scratch transients / fused state chain) or the
1862                // full-verify single-graph (fa exec-update) lands.
1863                e.capture_graph_retained_flags(
1864                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
1865                    move |e| {
1866                    let mut xc: Option<CudaSlice<f32>> = None;
1867                    for il in start..end {
1868                        let k = lin_pos[&il];
1869                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
1870                        let nx = model.qwen35_tparallel_linear_layer(
1871                            e,
1872                            il,
1873                            xr,
1874                            t,
1875                            cache_ref,
1876                            None,
1877                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
1878                            Some((table_all, k * 6)),
1879                        )?;
1880                        xc = Some(nx);
1881                    }
1882                        e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
1883                        Ok(())
1884                    },
1885                )?
1886            };
1887            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
1888            // is odd -> 3 runs = net one swap), then restore the device state the
1889            // warmups consumed. The launch below then behaves exactly like one run.
1890            if t % 2 == 1 {
1891                for il in start..end {
1892                    let rl = cache.recur[il].as_mut().unwrap();
1893                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1894                }
1895            }
1896            for (k, il) in (start..end).enumerate() {
1897                let rl = cache.recur[il].as_mut().unwrap();
1898                let (cw, sw) = (self.conv_words, self.ssm_words);
1899                {
1900                    let sv = e.view(&self.save_conv, self.max_run * cw);
1901                    let win = sv.slice(k * cw..(k + 1) * cw);
1902                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
1903                }
1904                {
1905                    let sv = e.view(&self.save_ssm, self.max_run * sw);
1906                    let win = sv.slice(k * sw..(k + 1) * sw);
1907                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
1908                }
1909            }
1910            self.graphs.insert(
1911                key,
1912                DsparkSegGraph {
1913                    graph,
1914                    _keeper: keeper,
1915                },
1916            );
1917        }
1918        self.graphs[&key].graph.launch()?;
1919        // Host parity bookkeeping for the replayed body (the captured host swaps do not
1920        // re-run at replay).
1921        if t % 2 == 1 {
1922            for il in start..end {
1923                let rl = cache.recur[il].as_mut().unwrap();
1924                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1925            }
1926        }
1927        let (_, xout) = self.stage.get(&t).unwrap();
1928        let mut out = e.uninit(t * n_embd)?;
1929        e.copy_into(&mut out, 0, xout, t * n_embd)?;
1930        Ok(out)
1931    }
1932
1933    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
1934    /// `row` (0-based) of layer `il`. None for non-linear layers.
1935    pub(crate) fn slab_row(
1936        &self,
1937        e: &Engine,
1938        il: usize,
1939        row: usize,
1940    ) -> Option<(u64, u64, usize, usize)> {
1941        use cudarc::driver::DevicePtr;
1942        let k = *self.lin_pos.get(&il)?;
1943        let s = &e.gpu.stream();
1944        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
1945        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
1946        Some((
1947            pc as u64 + (row * self.conv_words * 4) as u64,
1948            ps as u64 + (row * self.ssm_words * 4) as u64,
1949            self.conv_words,
1950            self.ssm_words,
1951        ))
1952    }
1953}
1954
1955impl VerifyCkpt {
1956    fn new(n_layer: usize) -> Self {
1957        VerifyCkpt {
1958            gdn: (0..n_layer).map(|_| None).collect(),
1959            cols: (0..n_layer).map(|_| None).collect(),
1960        }
1961    }
1962}
1963
1964/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
1965/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
1966/// a logical round number.
1967struct VerifyBoundaryTicket {
1968    rt: &'static crate::pp::PpNRt,
1969    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
1970    slot: usize,
1971    pos0: usize,
1972    t: usize,
1973    payload: usize,
1974    n_st: usize,
1975    pipelined: bool,
1976    pp_anatomy: bool,
1977    pp_started: std::time::Instant,
1978    reverse_ms: f64,
1979    stage0_ms: f64,
1980    tx_ms: f64,
1981    trace: Option<SpecPipeTraceCtx>,
1982}
1983
1984/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
1985/// increment-2 controller can also be armed by the server's fresh-process research door.
1986#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1987pub enum OptiForkGateMode {
1988    Disabled,
1989    Hit,
1990    Miss,
1991    Alternate,
1992    Abort,
1993    Controller,
1994}
1995
1996static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
1997static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
1998    std::sync::atomic::AtomicU32::new(0);
1999static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2000static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2001static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2002static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2003static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2004static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2005static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2006static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2007static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2008static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2009    std::sync::atomic::AtomicU64::new(0);
2010static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2011    std::sync::atomic::AtomicU64::new(0);
2012static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2013
2014impl OptiForkGateMode {
2015    fn code(self) -> u8 {
2016        match self {
2017            Self::Disabled => 0,
2018            Self::Hit => 1,
2019            Self::Miss => 2,
2020            Self::Alternate => 3,
2021            Self::Abort => 4,
2022            Self::Controller => 5,
2023        }
2024    }
2025
2026    fn configured() -> Self {
2027        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2028            1 => Self::Hit,
2029            2 => Self::Miss,
2030            3 => Self::Alternate,
2031            4 => Self::Abort,
2032            5 => Self::Controller,
2033            _ => Self::Disabled,
2034        }
2035    }
2036
2037    fn action(self, generation: u64) -> OptiForkAction {
2038        match self {
2039            Self::Hit => OptiForkAction::Hit,
2040            Self::Miss => OptiForkAction::Miss,
2041            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2042            Self::Alternate => OptiForkAction::Miss,
2043            Self::Abort => OptiForkAction::Abort,
2044            Self::Disabled | Self::Controller => {
2045                unreachable!("non-forced mode cannot choose a forced fork action")
2046            }
2047        }
2048    }
2049
2050    fn is_forced(self) -> bool {
2051        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2052    }
2053}
2054
2055/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2056pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2057    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2058}
2059
2060/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2061/// two-token draft-probability product. Serving can call this only through its explicit
2062/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2063pub fn set_optipipe_controller_threshold(threshold: f32) {
2064    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2065    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2066    set_optipipe_gate_mode(OptiForkGateMode::Controller);
2067}
2068
2069#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2070pub struct OptiForkGateStats {
2071    pub attempts: u64,
2072    pub hits: u64,
2073    pub misses: u64,
2074    pub abort_drains: u64,
2075    pub refusals: u64,
2076    pub gate_checks: u64,
2077    pub gate_admits: u64,
2078    pub gate_rejects: u64,
2079    pub reconciles: u64,
2080    pub wasted_draft_tokens: u64,
2081    pub shadow_draft_tokens: u64,
2082    pub breaker_trips: u64,
2083}
2084
2085#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2086pub struct OptiForkStateIdentity {
2087    pub trunk_kv_bytes: usize,
2088    pub recurrent_bytes: usize,
2089    pub scratch_kv_bytes: usize,
2090    pub hidden_bytes: usize,
2091}
2092
2093pub fn reset_optipipe_gate_stats() {
2094    for counter in [
2095        &OPTI_FORK_ATTEMPTS,
2096        &OPTI_FORK_HITS,
2097        &OPTI_FORK_MISSES,
2098        &OPTI_FORK_ABORT_DRAINS,
2099        &OPTI_FORK_REFUSALS,
2100        &OPTI_GATE_CHECKS,
2101        &OPTI_GATE_ADMITS,
2102        &OPTI_GATE_REJECTS,
2103        &OPTI_RECONCILES,
2104        &OPTI_WASTED_DRAFT_TOKENS,
2105        &OPTI_SHADOW_DRAFT_TOKENS,
2106        &OPTI_BREAKER_TRIPS,
2107    ] {
2108        counter.store(0, std::sync::atomic::Ordering::Relaxed);
2109    }
2110}
2111
2112pub fn optipipe_gate_stats() -> OptiForkGateStats {
2113    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2114    OptiForkGateStats {
2115        attempts: load(&OPTI_FORK_ATTEMPTS),
2116        hits: load(&OPTI_FORK_HITS),
2117        misses: load(&OPTI_FORK_MISSES),
2118        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2119        refusals: load(&OPTI_FORK_REFUSALS),
2120        gate_checks: load(&OPTI_GATE_CHECKS),
2121        gate_admits: load(&OPTI_GATE_ADMITS),
2122        gate_rejects: load(&OPTI_GATE_REJECTS),
2123        reconciles: load(&OPTI_RECONCILES),
2124        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2125        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2126        breaker_trips: load(&OPTI_BREAKER_TRIPS),
2127    }
2128}
2129
2130#[derive(Clone, Copy, Debug)]
2131struct OptiControllerPolicy {
2132    threshold: f32,
2133    consecutive_misses: u8,
2134    breaker_tripped: bool,
2135}
2136
2137impl OptiControllerPolicy {
2138    fn configured() -> Self {
2139        Self {
2140            threshold: f32::from_bits(
2141                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2142            ),
2143            consecutive_misses: 0,
2144            breaker_tripped: false,
2145        }
2146    }
2147
2148    fn admit(&self, q_proxy: f32) -> bool {
2149        q_proxy.is_finite()
2150            && (0.0..=1.0).contains(&q_proxy)
2151            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2152    }
2153
2154    /// Returns true exactly when this resolution newly trips the three-miss breaker.
2155    fn resolve(&mut self, hit: bool) -> bool {
2156        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2157        // every optimistic opportunity, so the safety breaker is measured separately and must
2158        // not silently turn this arm into "three attempts then serial".
2159        if self.threshold == 0.0 {
2160            self.consecutive_misses = 0;
2161            return false;
2162        }
2163        if hit {
2164            self.consecutive_misses = 0;
2165            return false;
2166        }
2167        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2168        if !self.breaker_tripped && self.consecutive_misses >= 3 {
2169            self.breaker_tripped = true;
2170            return true;
2171        }
2172        false
2173    }
2174}
2175
2176#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2177enum OptiForkAction {
2178    Hit,
2179    Miss,
2180    Abort,
2181}
2182
2183#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2184struct OptiForkGeneration {
2185    id: u64,
2186    slot: usize,
2187}
2188
2189#[derive(Default)]
2190struct OptiForkGenerationTracker {
2191    next: u64,
2192    live: [Option<u64>; 2],
2193}
2194
2195impl OptiForkGenerationTracker {
2196    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2197        let generation = OptiForkGeneration {
2198            id: self.next,
2199            slot: (self.next & 1) as usize,
2200        };
2201        if let Some(live) = self.live[generation.slot] {
2202            return Err(format!(
2203                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2204                generation.slot,
2205            )
2206            .into());
2207        }
2208        self.next += 1;
2209        self.live[generation.slot] = Some(generation.id);
2210        Ok(generation)
2211    }
2212
2213    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2214        match self.live[generation.slot] {
2215            Some(id) if id == generation.id => {
2216                self.live[generation.slot] = None;
2217                Ok(())
2218            }
2219            other => Err(format!(
2220                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2221                generation.id, generation.slot,
2222            )
2223            .into()),
2224        }
2225    }
2226}
2227
2228struct OptiForkSeedGeneration {
2229    h_seed: CudaSlice<f32>,
2230    fill_prev: CudaSlice<f32>,
2231    scratch_len: usize,
2232}
2233
2234/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2235/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2236/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2237/// device ownership.
2238fn opti_snapshot_stage_owned(
2239    e: &Engine,
2240    cache: &Cache,
2241    rt: &'static crate::pp::PpNRt,
2242    fence: &[usize],
2243) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2244    let n = cache.kv.len();
2245    let mut snapshot = crate::cache::CacheSnapshot {
2246        kv_len: vec![None; n],
2247        conv: (0..n).map(|_| None).collect(),
2248        ssm: (0..n).map(|_| None).collect(),
2249        pos: cache.pos,
2250    };
2251    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
2252    Ok(snapshot)
2253}
2254
2255fn opti_snapshot_stage_owned_into(
2256    e: &Engine,
2257    cache: &Cache,
2258    rt: &'static crate::pp::PpNRt,
2259    fence: &[usize],
2260    snapshot: &mut crate::cache::CacheSnapshot,
2261) -> Result<(), Box<dyn std::error::Error>> {
2262    if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
2263        return Err("optipipe stage-owned snapshot shape mismatch".into());
2264    }
2265    for stage in 0..rt.n_stages() {
2266        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
2267    }
2268    snapshot.pos = cache.pos;
2269    Ok(())
2270}
2271
2272/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
2273/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
2274/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
2275/// either point would capture one side of the fork at the wrong generation.
2276fn opti_snapshot_one_stage_owned_into(
2277    e: &Engine,
2278    cache: &Cache,
2279    rt: &'static crate::pp::PpNRt,
2280    fence: &[usize],
2281    stage: usize,
2282    snapshot: &mut crate::cache::CacheSnapshot,
2283) -> Result<(), Box<dyn std::error::Error>> {
2284    if fence.len() != rt.n_stages() + 1
2285        || snapshot.kv_len.len() != cache.kv.len()
2286        || stage >= rt.n_stages()
2287    {
2288        return Err("optipipe single-stage snapshot shape mismatch".into());
2289    }
2290    let _scope = rt.enter(stage);
2291    let owner = rt.engine(stage, e);
2292    for il in fence[stage]..fence[stage + 1] {
2293        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
2294        match &cache.recur[il] {
2295            Some(recur) => {
2296                match snapshot.conv[il].as_mut() {
2297                    Some(dst) => {
2298                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
2299                    }
2300                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
2301                }
2302                match snapshot.ssm[il].as_mut() {
2303                    Some(dst) => {
2304                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
2305                    }
2306                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
2307                }
2308            }
2309            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
2310                return Err(
2311                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
2312                );
2313            }
2314            None => {}
2315        }
2316    }
2317    snapshot.pos = cache.pos;
2318    Ok(())
2319}
2320
2321/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
2322/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
2323/// resolve, so the reconcile tables and conditional restores are stage-local.
2324struct OptiForkState {
2325    mode: OptiForkGateMode,
2326    controller: Option<OptiControllerPolicy>,
2327    generations: OptiForkGenerationTracker,
2328    active_snapshot_slot: usize,
2329    alternate_snapshot: crate::cache::CacheSnapshot,
2330    seeds: [OptiForkSeedGeneration; 2],
2331    rt: &'static crate::pp::PpNRt,
2332    fence: [usize; 3],
2333    split: usize,
2334    len_ptrs: CudaSlice<u64>,
2335    saved_lens: CudaSlice<i32>,
2336    forced_acc: CudaSlice<u32>,
2337    valid: CudaSlice<u32>,
2338    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2339    logical_payload_bytes: [usize; 2],
2340}
2341
2342struct OptiForkTicket {
2343    generation: OptiForkGeneration,
2344    boundary: Option<VerifyBoundaryTicket>,
2345    drain: std::sync::Arc<cudarc::driver::CudaStream>,
2346    settled: bool,
2347}
2348
2349struct OptiControllerTicket {
2350    generation: OptiForkGeneration,
2351    boundary: Option<VerifyBoundaryTicket>,
2352    ckpt: Option<VerifyCkpt>,
2353    verify_tokens: [u32; 2],
2354    draft_prob: f32,
2355    eager_seed: Option<CudaSlice<f32>>,
2356    q_proxy: f32,
2357    scratch_len: usize,
2358    issued_at: std::time::Instant,
2359    drain: std::sync::Arc<cudarc::driver::CudaStream>,
2360    settled: bool,
2361}
2362
2363struct OptiControllerPrepared {
2364    verify_tokens: [u32; 2],
2365    draft_prob: f32,
2366    eager_seed: Option<CudaSlice<f32>>,
2367    q_proxy: f32,
2368    scratch_len: usize,
2369}
2370
2371impl OptiControllerTicket {
2372    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2373        self.boundary
2374            .take()
2375            .expect("controller boundary ticket already consumed")
2376    }
2377
2378    fn take_ckpt(&mut self) -> VerifyCkpt {
2379        self.ckpt
2380            .take()
2381            .expect("controller verify checkpoint already consumed")
2382    }
2383
2384    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
2385        self.eager_seed.take()
2386    }
2387
2388    fn settle(&mut self) {
2389        self.settled = true;
2390    }
2391}
2392
2393impl Drop for OptiControllerTicket {
2394    fn drop(&mut self) {
2395        if !self.settled {
2396            let _ = self.drain.synchronize();
2397            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2398        }
2399    }
2400}
2401
2402impl OptiForkTicket {
2403    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2404        self.boundary
2405            .take()
2406            .expect("fork ticket boundary already consumed")
2407    }
2408
2409    fn settle(&mut self) {
2410        self.settled = true;
2411    }
2412}
2413
2414impl Drop for OptiForkTicket {
2415    fn drop(&mut self) {
2416        if !self.settled {
2417            let _ = self.drain.synchronize();
2418            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2419        }
2420    }
2421}
2422
2423impl OptiForkState {
2424    #[allow(clippy::too_many_arguments)]
2425    fn new(
2426        e: &Engine,
2427        cache: &Cache,
2428        mode: OptiForkGateMode,
2429        alternate_snapshot: crate::cache::CacheSnapshot,
2430        h_seed: &CudaSlice<f32>,
2431        fill_prev: &CudaSlice<f32>,
2432        rt: &'static crate::pp::PpNRt,
2433        split: usize,
2434        n_layer: usize,
2435    ) -> Result<Self, Box<dyn std::error::Error>> {
2436        let fence = [0, split, n_layer];
2437        let mut logical_payload_bytes = [0usize; 2];
2438        for stage in 0..2 {
2439            for il in fence[stage]..fence[stage + 1] {
2440                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
2441                    .as_ref()
2442                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2443                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
2444                    .as_ref()
2445                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2446            }
2447        }
2448        let seeds = [
2449            OptiForkSeedGeneration {
2450                h_seed: e.clone_dtod(h_seed)?,
2451                fill_prev: e.clone_dtod(fill_prev)?,
2452                scratch_len: 0,
2453            },
2454            OptiForkSeedGeneration {
2455                h_seed: e.clone_dtod(h_seed)?,
2456                fill_prev: e.clone_dtod(fill_prev)?,
2457                scratch_len: 0,
2458            },
2459        ];
2460        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
2461            let _stage = rt.enter(0);
2462            let e0 = rt.engine(0, e);
2463            (
2464                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
2465                e0.htod_i32(&vec![0; split])?,
2466                e0.alloc_u32_zeroed(2)?,
2467                e0.alloc_u32_zeroed(1)?,
2468                e0.stream(),
2469            )
2470        };
2471        logical_payload_bytes[0] += seeds
2472            .iter()
2473            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
2474            .sum::<usize>();
2475        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
2476            + saved_lens.len() * std::mem::size_of::<i32>()
2477            + forced_acc.len() * std::mem::size_of::<u32>()
2478            + valid.len() * std::mem::size_of::<u32>();
2479        Ok(Self {
2480            mode,
2481            controller: (mode == OptiForkGateMode::Controller)
2482                .then(OptiControllerPolicy::configured),
2483            generations: OptiForkGenerationTracker::default(),
2484            active_snapshot_slot: 0,
2485            alternate_snapshot,
2486            seeds,
2487            rt,
2488            fence,
2489            split,
2490            len_ptrs,
2491            saved_lens,
2492            forced_acc,
2493            valid,
2494            stage0_stream,
2495            logical_payload_bytes,
2496        })
2497    }
2498
2499    fn reserve(
2500        &mut self,
2501        current_snapshot: &mut crate::cache::CacheSnapshot,
2502    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2503        let generation = self.generations.reserve()?;
2504        if generation.slot != self.active_snapshot_slot {
2505            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2506            self.active_snapshot_slot = generation.slot;
2507        }
2508        Ok(generation)
2509    }
2510
2511    fn capture_seed(
2512        &mut self,
2513        e: &Engine,
2514        generation: OptiForkGeneration,
2515        h_seed: &CudaSlice<f32>,
2516        fill_prev: &CudaSlice<f32>,
2517        scratch_len: usize,
2518    ) -> Result<(), Box<dyn std::error::Error>> {
2519        let seed = &mut self.seeds[generation.slot];
2520        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
2521        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
2522        seed.scratch_len = scratch_len;
2523        Ok(())
2524    }
2525
2526    fn ticket(
2527        &self,
2528        generation: OptiForkGeneration,
2529        boundary: VerifyBoundaryTicket,
2530    ) -> OptiForkTicket {
2531        OptiForkTicket {
2532            generation,
2533            boundary: Some(boundary),
2534            drain: self.stage0_stream.clone(),
2535            settled: false,
2536        }
2537    }
2538
2539    #[allow(clippy::too_many_arguments)]
2540    fn controller_ticket(
2541        &self,
2542        generation: OptiForkGeneration,
2543        boundary: VerifyBoundaryTicket,
2544        ckpt: VerifyCkpt,
2545        verify_tokens: [u32; 2],
2546        draft_prob: f32,
2547        eager_seed: Option<CudaSlice<f32>>,
2548        q_proxy: f32,
2549        scratch_len: usize,
2550    ) -> OptiControllerTicket {
2551        OptiControllerTicket {
2552            generation,
2553            boundary: Some(boundary),
2554            ckpt: Some(ckpt),
2555            verify_tokens,
2556            draft_prob,
2557            eager_seed,
2558            q_proxy,
2559            scratch_len,
2560            issued_at: std::time::Instant::now(),
2561            drain: self.stage0_stream.clone(),
2562            settled: false,
2563        }
2564    }
2565
2566    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2567        self.generations.reserve()
2568    }
2569
2570    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
2571        &mut self.alternate_snapshot
2572    }
2573
2574    fn promote_successor_snapshot(
2575        &mut self,
2576        current_snapshot: &mut crate::cache::CacheSnapshot,
2577        generation: OptiForkGeneration,
2578    ) {
2579        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
2580        self.active_snapshot_slot = generation.slot;
2581    }
2582
2583    fn queue_actual_reconcile(
2584        &mut self,
2585        e: &Engine,
2586        snapshot: &crate::cache::CacheSnapshot,
2587        acc: &CudaSlice<u32>,
2588        optimistic_pending: u32,
2589        base: usize,
2590    ) -> Result<(), Box<dyn std::error::Error>> {
2591        let saved: Vec<i32> = (0..self.split)
2592            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2593            .collect();
2594        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
2595        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
2596        // the validity/reconcile kernels must never peer-read acc before it is written. The
2597        // increment-1 harness uses primary stage 0, where stream order already provides this.
2598        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
2599            self.rt.fence_stages_behind(&e.stream())?;
2600        }
2601        let _stage = self.rt.enter(0);
2602        let e0 = self.rt.engine(0, e);
2603        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2604        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
2605        e0.spec_fork_reconcile_kv(
2606            &self.len_ptrs,
2607            &self.saved_lens,
2608            acc,
2609            &self.valid,
2610            base,
2611            self.split,
2612        )
2613    }
2614
2615    fn finish_actual_reconcile(
2616        &mut self,
2617        e: &Engine,
2618        cache: &mut Cache,
2619        snapshot: &crate::cache::CacheSnapshot,
2620        n_acc: usize,
2621        base: usize,
2622        hit: bool,
2623    ) -> Result<(), Box<dyn std::error::Error>> {
2624        if hit {
2625            return Ok(());
2626        }
2627        let len_delta = base + n_acc;
2628        for il in 0..self.split {
2629            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2630                kv.len = saved + len_delta;
2631            }
2632        }
2633        {
2634            let _stage = self.rt.enter(1);
2635            let e1 = self.rt.engine(1, e);
2636            for il in self.split..self.fence[2] {
2637                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2638                    kv.len = saved + len_delta;
2639                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
2640                }
2641            }
2642        }
2643        self.rt.publish_to(0, &e.stream())?;
2644        Ok(())
2645    }
2646
2647    fn cancel_controller_ticket(
2648        &mut self,
2649        e: &Engine,
2650        cache: &mut Cache,
2651        scratch: &mut MtpScratch,
2652        snapshot: &crate::cache::CacheSnapshot,
2653        ticket: &mut OptiControllerTicket,
2654    ) -> Result<(), Box<dyn std::error::Error>> {
2655        {
2656            let _stage = self.rt.enter(0);
2657            let e0 = self.rt.engine(0, e);
2658            for il in 0..self.split {
2659                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2660                    kv.len = saved;
2661                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
2662                }
2663            }
2664        }
2665        scratch.set_len(e, snapshot.pos)?;
2666        ticket.settle();
2667        self.generations.retire(ticket.generation)?;
2668        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2669        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
2670        eprintln!(
2671            "[opti-controller] tail-drain generation={} slot={}",
2672            ticket.generation.id, ticket.generation.slot,
2673        );
2674        Ok(())
2675    }
2676
2677    #[allow(clippy::too_many_arguments)]
2678    fn reconcile(
2679        &mut self,
2680        e: &Engine,
2681        cache: &mut Cache,
2682        scratch: &mut MtpScratch,
2683        snapshot: &crate::cache::CacheSnapshot,
2684        h_seed: &mut CudaSlice<f32>,
2685        fill_prev: &mut CudaSlice<f32>,
2686        generation: OptiForkGeneration,
2687        action: OptiForkAction,
2688        optimistic_pending: u32,
2689    ) -> Result<(), Box<dyn std::error::Error>> {
2690        debug_assert!(action != OptiForkAction::Abort);
2691        let miss_started = std::time::Instant::now();
2692        let keep = action == OptiForkAction::Hit;
2693        let saved: Vec<i32> = (0..self.split)
2694            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
2695            .collect();
2696        let seed = &self.seeds[generation.slot];
2697        {
2698            let _stage = self.rt.enter(0);
2699            let e0 = self.rt.engine(0, e);
2700            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
2701            let forced = if keep {
2702                [1u32, optimistic_pending]
2703            } else {
2704                [0u32, optimistic_pending]
2705            };
2706            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
2707            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
2708            e0.spec_fork_reconcile_kv(
2709                &self.len_ptrs,
2710                &self.saved_lens,
2711                &self.forced_acc,
2712                &self.valid,
2713                0,
2714                self.split,
2715            )?;
2716            for il in 0..self.split {
2717                if let Some(recur) = cache.recur[il].as_mut() {
2718                    let conv = snapshot.conv[il]
2719                        .as_ref()
2720                        .ok_or("optipipe stage0 snapshot missing conv state")?;
2721                    let ssm = snapshot.ssm[il]
2722                        .as_ref()
2723                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
2724                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
2725                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
2726                }
2727            }
2728            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
2729            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
2730        }
2731
2732        if keep {
2733            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2734            return Ok(());
2735        }
2736
2737        for il in 0..self.split {
2738            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
2739                kv.len = saved;
2740            }
2741        }
2742        scratch.set_len(e, seed.scratch_len)?;
2743        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
2744        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
2745        let caller = e.stream();
2746        self.rt.publish_to(0, &caller)?;
2747        caller.synchronize()?;
2748        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
2749        eprintln!(
2750            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
2751            generation.id, generation.slot,
2752        );
2753        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2754        Ok(())
2755    }
2756
2757    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2758        self.generations.retire(generation)
2759    }
2760}
2761
2762impl HybridModel {
2763    fn opti_graph_draft_step(
2764        &self,
2765        e: &Engine,
2766        mtp: &MtpHead,
2767        dctx: &mut DraftGraphCtx,
2768        scratch: &mut MtpScratch,
2769        d_vocab: usize,
2770    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2771        dctx.graph
2772            .as_ref()
2773            .ok_or("optipipe controller requires the greedy draft graph")?
2774            .launch()?;
2775        scratch.kv.len += 1;
2776        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
2777        if (idx as usize) >= d_vocab {
2778            return Err(
2779                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
2780            );
2781        }
2782        let probability = e.dtoh(&dctx.g_p)?[0];
2783        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2784            return Err(format!("optipipe draft probability is invalid: {probability}").into());
2785        }
2786        let token = match &mtp.d2t {
2787            Some(map) => map[idx as usize],
2788            None => idx,
2789        };
2790        if token != idx {
2791            e.set_u32_one(&mut dctx.g_tok, token)?;
2792        }
2793        Ok((token, probability))
2794    }
2795
2796    #[allow(clippy::too_many_arguments)]
2797    fn opti_controller_draft_step(
2798        &self,
2799        e: &Engine,
2800        mtp: &MtpHead,
2801        dctx: &mut DraftGraphCtx,
2802        scratch: &mut MtpScratch,
2803        d_vocab: usize,
2804        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
2805        eager_pos: usize,
2806        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2807    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
2808        if dctx.graph.is_some() {
2809            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
2810        }
2811        let (input_token, input_seed) = eager_state
2812            .take()
2813            .ok_or("optipipe eager continuation seed is unavailable")?;
2814        let (logits, next_seed) = self.mtp_head_forward_dev(
2815            e,
2816            mtp,
2817            input_token,
2818            &input_seed,
2819            scratch,
2820            eager_pos,
2821            embd_dev,
2822            None,
2823        )?;
2824        let token_d = e.argmax_token_device(&logits, d_vocab)?;
2825        let idx = e.dtoh_u32_one(&token_d)?;
2826        if (idx as usize) >= d_vocab {
2827            return Err(format!(
2828                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
2829            )
2830            .into());
2831        }
2832        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
2833        let probability = e.dtoh(&probability_d)?[0];
2834        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
2835            return Err(
2836                format!("optipipe eager draft probability is invalid: {probability}").into(),
2837            );
2838        }
2839        let token = match &mtp.d2t {
2840            Some(map) => map[idx as usize],
2841            None => idx,
2842        };
2843        *eager_state = Some((token, next_seed));
2844        Ok((token, probability))
2845    }
2846
2847    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
2848    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
2849    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
2850    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
2851    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
2852    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
2853    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
2854    /// transfer + host argmax per draft token from the K-token draft chain.
2855    #[allow(clippy::too_many_arguments)]
2856    fn mtp_head_forward_dev(
2857        &self,
2858        e: &Engine,
2859        mtp: &MtpHead,
2860        e_tok: u32,
2861        h_seed: &CudaSlice<f32>,
2862        scratch: &mut MtpScratch,
2863        mtp_pos: usize,
2864        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
2865        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
2866        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
2867        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
2868        mask: Option<(&CudaSlice<u32>, usize)>,
2869    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
2870        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
2871        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
2872        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
2873        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
2874        static ANAT_NS: [AtomicU64; 5] = [
2875            AtomicU64::new(0),
2876            AtomicU64::new(0),
2877            AtomicU64::new(0),
2878            AtomicU64::new(0),
2879            AtomicU64::new(0),
2880        ];
2881        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
2882        let anat = {
2883            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2884            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
2885        };
2886        if anat {
2887            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
2888        }
2889        let t_all = std::time::Instant::now();
2890        let mut t_ph = std::time::Instant::now();
2891        let mut anat_mark = |i: usize,
2892                             e: &Engine,
2893                             t: &mut std::time::Instant|
2894         -> Result<(), Box<dyn std::error::Error>> {
2895            if anat {
2896                e.stream().synchronize()?;
2897                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
2898                *t = std::time::Instant::now();
2899            }
2900            Ok(())
2901        };
2902        let cfg = &self.cfg;
2903        let n_embd = cfg.n_embd as usize;
2904        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
2905        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
2906        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
2907        let eps = cfg.rms_eps;
2908        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
2909
2910        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
2911        // expands this one row on CPU and transfers n_embd f32 values instead.
2912        let e_emb = match embd_dev {
2913            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
2914            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
2915        };
2916
2917        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
2918        let mut e_norm = e.zeros(n_embd)?;
2919        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
2920        let mut h_norm = e.zeros(n_embd)?;
2921        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
2922
2923        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
2924        let mut concat = e.zeros(2 * n_embd)?;
2925        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
2926        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
2927
2928        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
2929        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
2930
2931        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
2932        let mut a_norm = e.zeros(di)?;
2933        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
2934        anat_mark(0, e, &mut t_ph)?;
2935
2936        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
2937        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
2938        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
2939        // advances only the device counter).
2940        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
2941            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
2942            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
2943            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
2944            // whose host-side mirror the caller does).
2945            (Mixer::Full(fa), Some(g)) => {
2946                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
2947            }
2948            (Mixer::Full(fa), None) => {
2949                let out =
2950                    self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
2951                scratch.kv.len += 1;
2952                out
2953            }
2954            (Mixer::Linear(_), _) => {
2955                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
2956            }
2957            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
2958        };
2959        anat_mark(1, e, &mut t_ph)?;
2960
2961        // op 7: x1 = inpSA + attn_out
2962        let mut x1 = e.zeros(di)?;
2963        e.add(&inp_sa, &attn_out, &mut x1, di)?;
2964
2965        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
2966        let mut z = e.zeros(di)?;
2967        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
2968
2969        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
2970        let ffn_out = match &mtp.ffn {
2971            crate::hybrid::Ffn::Dense {
2972                ffn_gate,
2973                ffn_up,
2974                ffn_down,
2975            } => {
2976                let n_ff = ffn_gate.out_features();
2977                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
2978                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
2979                    (
2980                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
2981                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
2982                    )
2983                } else {
2984                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
2985                };
2986                let mut act = e.zeros(n_ff)?;
2987                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
2988                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
2989                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
2990                // passes None, which is `ffn_act`'s dispatch verbatim.
2991                Self::ffn_act_lim(
2992                    e,
2993                    &self.cfg,
2994                    &gate,
2995                    &up,
2996                    1.0,
2997                    1.0,
2998                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
2999                    &mut act,
3000                    n_ff,
3001                )?;
3002                e.matmul(ffn_down, &act, 1)?
3003            }
3004            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3005            // so they never alias trunk layer 0's cache keys.
3006            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3007        };
3008        anat_mark(2, e, &mut t_ph)?;
3009
3010        // op 10: h_nextn = x1 + ffn_out (at di)
3011        let mut h_inner = e.zeros(di)?;
3012        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3013
3014        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3015        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3016        let h_nextn = match mtp.geom.as_ref() {
3017            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3018            None => h_inner,
3019        };
3020
3021        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3022        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3023        let mut final_h = e.zeros(n_embd)?;
3024        e.rms_norm(
3025            &h_nextn,
3026            final_norm.float_data(),
3027            &mut final_h,
3028            n_embd,
3029            1,
3030            eps,
3031        )?;
3032
3033        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3034        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3035        let mut logits = e.matmul(head, &final_h, 1)?;
3036        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3037        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3038        if let Some((mask_d, mw)) = mask {
3039            let d_vocab = head.out_features();
3040            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3041        }
3042        anat_mark(3, e, &mut t_ph)?;
3043        if anat {
3044            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3045            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3046            if n % 128 == 0 {
3047                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3048                eprintln!(
3049                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3050                    us(0),
3051                    us(1),
3052                    us(2),
3053                    us(3),
3054                    us(4)
3055                );
3056            }
3057        }
3058        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3059        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3060        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3061    }
3062
3063    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3064    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3065    /// the dc path, and all three are properties of this arch's MTP block:
3066    ///
3067    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3068    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3069    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3070    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3071    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3072    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3073    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
3074    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3075    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3076    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3077    ///    resolved `Step35MtpGeom`, never from `cfg`.
3078    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3079    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3080    ///    fused-into-wq `q_gate_split` form the dc arm handles.
3081    ///
3082    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3083    /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3084    /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3085    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3086    ///
3087    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3088    /// caller must not mirror.
3089    fn mtp_step35_attn(
3090        &self,
3091        e: &Engine,
3092        fa: &FullAttnLayer,
3093        g: &crate::hybrid::Step35MtpGeom,
3094        h: &CudaSlice<f32>,
3095        pos_d: &CudaSlice<i32>,
3096        scratch: &mut MtpScratch,
3097    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3098        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3099        let eps = self.cfg.rms_eps;
3100        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3101        let n_embd = self.cfg.n_embd as usize;
3102        let gw = fa
3103            .attn_gate
3104            .as_ref()
3105            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3106
3107        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3108            && e.uses_q8_1_fast(&fa.wk)
3109            && e.uses_q8_1_fast(&fa.wv)
3110            && e.uses_q8_1_fast(gw)
3111        {
3112            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3113            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3114                Some(t3) => t3,
3115                None => (
3116                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3117                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3118                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3119                ),
3120            };
3121            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3122        } else {
3123            (
3124                e.matmul(&fa.wq, h, 1)?,
3125                e.matmul(&fa.wk, h, 1)?,
3126                e.matmul(&fa.wv, h, 1)?,
3127                e.matmul(gw, h, 1)?,
3128            )
3129        };
3130
3131        let mut q = e.uninit(nh * hd)?;
3132        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3133        let mut k = e.uninit(nkv * hd)?;
3134        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3135        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3136        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3137        // the resolved flag, not the constant, so an all-full sibling stays correct.
3138        let ff = if g.swa {
3139            None
3140        } else {
3141            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3142        };
3143        #[cfg(debug_assertions)]
3144        if let Some(ff) = ff {
3145            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3146        }
3147        e.rope_neox2(
3148            &mut q,
3149            &mut k,
3150            pos_d,
3151            hd,
3152            g.n_rot,
3153            nh,
3154            nkv,
3155            1,
3156            g.rope_base,
3157            1.0,
3158            ff,
3159        )?;
3160
3161        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3162        // length on the host anyway, and the windowed view below needs it there to compute the
3163        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3164        // dc-family consumer of this scratch still agree.
3165        let kv = &mut scratch.kv;
3166        assert!(
3167            kv.len < scratch.cap,
3168            "step35 MTP scratch overflow ({} >= {})",
3169            kv.len,
3170            scratch.cap
3171        );
3172        let next_len = kv.len + 1;
3173        let (off, t_kv) = if g.swa && next_len > g.window {
3174            (next_len - g.window, g.window)
3175        } else {
3176            (0, next_len)
3177        };
3178        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
3179        e.append_kv_quantized(
3180            &k,
3181            &v0,
3182            &mut kv.k,
3183            &mut kv.v,
3184            write_row,
3185            kv.kv_dim_k,
3186            kv.kv_dim_v,
3187            kv.k_tok_bytes,
3188            kv.v_tok_bytes,
3189            false,
3190        )?;
3191        kv.len = next_len;
3192        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3193        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
3194        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
3195        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
3196        // therefore live, not theoretical.
3197        let physical = kv.physical_rows(off, off + t_kv)?;
3198        let k_view = e.view_u8_range(
3199            &kv.k,
3200            physical.start * kv.k_tok_bytes,
3201            physical.end * kv.k_tok_bytes,
3202        );
3203        let v_view = e.view_u8_range(
3204            &kv.v,
3205            physical.start * kv.v_tok_bytes,
3206            physical.end * kv.v_tok_bytes,
3207        );
3208        let mut attn = e.uninit(nh * hd)?;
3209        e.fa_decode_kvmod(
3210            &q,
3211            &k_view,
3212            &v_view,
3213            &mut attn,
3214            hd,
3215            nh,
3216            nkv,
3217            t_kv,
3218            scale,
3219            kv.k_tok_bytes,
3220            kv.v_tok_bytes,
3221            false,
3222        )?;
3223
3224        let mut ag = e.uninit(nh * hd)?;
3225        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
3226        Ok(e.matmul(&fa.wo, &ag, 1)?)
3227    }
3228
3229    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
3230    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
3231    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
3232    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
3233    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
3234    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
3235    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
3236    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
3237    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
3238    fn mtp_full_attn_dc(
3239        &self,
3240        e: &Engine,
3241        fa: &FullAttnLayer,
3242        h: &CudaSlice<f32>,
3243        pos_d: &CudaSlice<i32>,
3244        scratch: &mut MtpScratch,
3245        geom: Option<&crate::hybrid::DraftGeom>,
3246    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3247        let cfg = &self.cfg;
3248        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3249        let geometry = cfg.full_attention_geometry_at(mtp_il);
3250        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
3251        let n_head_kv = geom
3252            .map(|g| g.n_head_kv)
3253            .unwrap_or(geometry.n_head_kv as usize);
3254        let head_dim = geometry.head_dim_k as usize;
3255        let eps = cfg.rms_eps;
3256        let scale = geometry.attention_scale();
3257        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
3258        let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
3259
3260        let (qf, mut k, v) =
3261            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
3262                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
3263                (
3264                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
3265                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
3266                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
3267                )
3268            } else {
3269                (
3270                    e.matmul(&fa.wq, h, 1)?,
3271                    e.matmul(&fa.wk, h, 1)?,
3272                    e.matmul(&fa.wv, h, 1)?,
3273                )
3274            };
3275        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
3276        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3277        let (mut q, gate) = if gated {
3278            let mut q = e.zeros(n_head * head_dim)?;
3279            let mut gate = e.zeros(n_head * head_dim)?;
3280            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
3281            (q, Some(gate))
3282        } else {
3283            (qf, None)
3284        };
3285
3286        let mut qn = e.zeros(n_head * head_dim)?;
3287        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
3288        q = qn;
3289        let mut kn = e.zeros(n_head_kv * head_dim)?;
3290        e.rms_norm(
3291            &k,
3292            fa.k_norm.float_data(),
3293            &mut kn,
3294            head_dim,
3295            n_head_kv,
3296            eps,
3297        )?;
3298        k = kn;
3299        let rope_dims = geometry.n_rot as usize;
3300        e.rope_neox(
3301            &mut q,
3302            pos_d,
3303            head_dim,
3304            rope_dims,
3305            n_head,
3306            1,
3307            geometry.rope_base,
3308            1.0,
3309        )?;
3310        e.rope_neox(
3311            &mut k,
3312            pos_d,
3313            head_dim,
3314            rope_dims,
3315            n_head_kv,
3316            1,
3317            geometry.rope_base,
3318            1.0,
3319        )?;
3320
3321        let kv = &mut scratch.kv;
3322        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
3323        e.append_kv_quantized_dc(
3324            &k,
3325            &v,
3326            &mut kv.k,
3327            &mut kv.v,
3328            &kv.len_d,
3329            kv.kv_dim_k,
3330            kv.kv_dim_v,
3331            kv.k_tok_bytes,
3332            kv.v_tok_bytes,
3333            false,
3334        )?;
3335        e.inc_seqlen(&mut kv.len_d)?;
3336        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
3337        // key range from the device counter.
3338        let k_view = e.view_u8(&kv.k, kv.k.len());
3339        let v_view = e.view_u8(&kv.v, kv.v.len());
3340        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
3341        let mut attn = e.zeros(n_head * head_dim)?;
3342        e.fa_decode_dc(
3343            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
3344            scale, ktb, vtb, false,
3345        )?;
3346
3347        let attn_g = match &gate {
3348            Some(gate) => {
3349                let mut gsig = e.zeros(n_head * head_dim)?;
3350                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
3351                let mut ag = e.zeros(n_head * head_dim)?;
3352                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
3353                ag
3354            }
3355            None => attn,
3356        };
3357        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
3358    }
3359
3360    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
3361    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
3362    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
3363    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
3364    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
3365    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
3366    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
3367    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
3368    #[allow(clippy::too_many_arguments)]
3369    fn mtp_kv_fill(
3370        &self,
3371        e: &Engine,
3372        mtp: &MtpHead,
3373        tokens: &[u32],
3374        h: &CudaSlice<f32>,
3375        pos0: usize,
3376        scratch: &mut MtpScratch,
3377        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3378    ) -> Result<(), Box<dyn std::error::Error>> {
3379        let cfg = &self.cfg;
3380        let n_embd = cfg.n_embd as usize;
3381        let eps = cfg.rms_eps;
3382        let t = tokens.len();
3383        assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
3384        assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
3385        let Mixer::Full(fa) = &mtp.mixer else {
3386            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3387        };
3388        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
3389        let pos_d = e.htod_i32(&pos_vec)?;
3390
3391        // ops A/1/2: embed + the two input norms, T-wide.
3392        let e_emb = match embd_dev {
3393            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3394            None => e.htod(&self.embd.gather(n_embd, tokens))?,
3395        };
3396        let mut e_norm = e.zeros(t * n_embd)?;
3397        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
3398        let mut h_norm = e.zeros(t * n_embd)?;
3399        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
3400
3401        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
3402        let mut concat = e.zeros(t * 2 * n_embd)?;
3403        for i in 0..t {
3404            e.copy_view_into(
3405                &mut concat,
3406                i * 2 * n_embd,
3407                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
3408                n_embd,
3409            )?;
3410            e.copy_view_into(
3411                &mut concat,
3412                i * 2 * n_embd + n_embd,
3413                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
3414                n_embd,
3415            )?;
3416        }
3417
3418        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
3419        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3420        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
3421        let mut a_norm = e.zeros(t * di)?;
3422        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
3423
3424        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
3425        // the fill only has to leave correct K/V rows behind for later chains to attend over.
3426        let n_head_kv = mtp
3427            .geom
3428            .as_ref()
3429            .map(|g| g.n_head_kv)
3430            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
3431            .unwrap_or_else(|| {
3432                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3433                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
3434            });
3435        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3436        let geometry = cfg.full_attention_geometry_at(mtp_il);
3437        let head_dim = geometry.head_dim_k as usize;
3438        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
3439        let v = e.matmul(&fa.wv, &a_norm, t)?;
3440        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
3441        e.rms_norm(
3442            &k,
3443            fa.k_norm.float_data(),
3444            &mut kn,
3445            head_dim,
3446            n_head_kv * t,
3447            eps,
3448        )?;
3449        k = kn;
3450        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
3451        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
3452        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
3453        // writes K rows the attention arm then re-derives at a different theta: correct-looking
3454        // output with dead acceptance, invisible to the exactness gates.
3455        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
3456            Some(s) => (
3457                s.n_rot,
3458                s.rope_base,
3459                if s.swa {
3460                    None
3461                } else {
3462                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3463                },
3464            ),
3465            None => (geometry.n_rot as usize, geometry.rope_base, None),
3466        };
3467        #[cfg(debug_assertions)]
3468        if let Some(ff) = ff {
3469            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
3470        }
3471        match ff {
3472            Some(f) => e.rope_neox_ff(
3473                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
3474            )?,
3475            None => e.rope_neox(
3476                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
3477            )?,
3478        }
3479
3480        let kv = &mut scratch.kv;
3481        // Match the trunk prime contract: a chunk may need the aligned window immediately before
3482        // its first row, so preserve that prefix when the physical tail rebases at wrap.
3483        let retain_from = kv
3484            .ring
3485            .as_ref()
3486            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
3487            .unwrap_or(0);
3488        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
3489        for i in 0..t {
3490            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
3491            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
3492            e.append_kv_quantized_view(
3493                &k_row,
3494                &v_row,
3495                &mut kv.k,
3496                &mut kv.v,
3497                write_row + i,
3498                kv.kv_dim_k,
3499                kv.kv_dim_v,
3500                kv.k_tok_bytes,
3501                kv.v_tok_bytes,
3502                false,
3503            )?;
3504        }
3505        kv.len = pos0 + t;
3506        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3507        Ok(())
3508    }
3509
3510    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
3511    /// every varying input device-resident —
3512    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
3513    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
3514    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
3515    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
3516    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
3517    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
3518    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
3519    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
3520    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
3521    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
3522    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
3523    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
3524    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
3525    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
3526    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
3527    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
3528    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
3529    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
3530    #[allow(clippy::too_many_arguments)]
3531    fn mtp_head_forward_cap(
3532        &self,
3533        e: &Engine,
3534        mtp: &MtpHead,
3535        tok_d: &mut CudaSlice<u32>,
3536        pos_d: &mut CudaSlice<i32>,
3537        h_seed_d: &mut CudaSlice<f32>,
3538        p_d: &mut CudaSlice<f32>,
3539        scratch: &mut MtpScratch,
3540        with_prob: bool,
3541        with_head: bool,
3542        embd_gpu: &CudaSlice<u8>,
3543        embd_qt: i32,
3544        embd_rb: usize,
3545        d_vocab: usize,
3546        sampled_cap: Option<(
3547            &mut CudaSlice<u32>,
3548            &mut CudaSlice<f32>,
3549            &mut CudaSlice<f32>,
3550            u64,
3551            f32,
3552        )>,
3553        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
3554        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
3555        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
3556        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
3557        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
3558        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
3559        mask_cap: Option<(&CudaSlice<u32>, usize)>,
3560    ) -> Result<(), Box<dyn std::error::Error>> {
3561        let cfg = &self.cfg;
3562        let n_embd = cfg.n_embd as usize;
3563        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
3564        // whose device-counter key bound always starts at row 0 — it cannot express this block's
3565        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
3566        // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
3567        // refuses step35 heads explicitly (SWA refusal), so the eager chain
3568        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
3569        // panic) is what the two capture sites and the round-stream capture already handle by
3570        // degrading to eager / stream-off.
3571        if mtp.step35.is_some() {
3572            return Err(
3573                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
3574                        block's SWA view offset; same root cause as the dc decode refusal) — the \
3575                        eager draft chain serves this arch"
3576                    .into(),
3577            );
3578        }
3579        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
3580        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3581        let eps = cfg.rms_eps;
3582        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
3583        let mut e_norm = e.zeros(n_embd)?;
3584        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3585        let mut h_norm = e.zeros(n_embd)?;
3586        e.rms_norm(
3587            &*h_seed_d,
3588            mtp.hnorm.float_data(),
3589            &mut h_norm,
3590            n_embd,
3591            1,
3592            eps,
3593        )?;
3594        let mut concat = e.zeros(2 * n_embd)?;
3595        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3596        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3597        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3598        let mut a_norm = e.zeros(di)?;
3599        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3600        let attn_out = match &mtp.mixer {
3601            Mixer::Full(fa) => {
3602                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
3603            }
3604            Mixer::Linear(_) => {
3605                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3606            }
3607            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
3608        };
3609        let mut x1 = e.zeros(di)?;
3610        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3611        let mut z = e.zeros(di)?;
3612        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3613        let ffn_out = match &mtp.ffn {
3614            crate::hybrid::Ffn::Dense {
3615                ffn_gate,
3616                ffn_up,
3617                ffn_down,
3618            } => {
3619                let n_ff = ffn_gate.out_features();
3620                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3621                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3622                    (
3623                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3624                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3625                    )
3626                } else {
3627                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3628                };
3629                let mut act = e.zeros(n_ff)?;
3630                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
3631                e.matmul(ffn_down, &act, 1)?
3632            }
3633            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
3634            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
3635            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
3636            // error arm degrades the caller to eager/stream-off.
3637            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
3638                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
3639            }
3640            crate::hybrid::Ffn::Moe(_) => {
3641                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
3642            }
3643        };
3644        let mut h_inner = e.zeros(di)?;
3645        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3646        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
3647        let h_nextn = match mtp.geom.as_ref() {
3648            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3649            None => h_inner,
3650        };
3651        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
3652        let final_h = if with_head || spec_hpost() {
3653            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3654            let mut fh = e.zeros(n_embd)?;
3655            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
3656            Some(fh)
3657        } else {
3658            None
3659        };
3660        if with_head {
3661            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3662            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
3663            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
3664            // before the argmax — proposals become legal by construction. Contents-only
3665            // per-replay upload keeps the capture valid.
3666            if let Some((mask_d, mw)) = mask_cap {
3667                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3668            }
3669            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
3670                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
3671                // own buffer is pool-recycled after the capture body returns, so it can't be the
3672                // retention target), bump the device event counter, gumbel-perturb reading it,
3673                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
3674                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
3675                e.sctr_inc(ctr_d)?;
3676                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
3677                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
3678                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
3679                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
3680                if with_prob {
3681                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3682                }
3683            } else {
3684                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
3685                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
3686                // p-min under a draft mask reads the MASKED row: confidence relative to the
3687                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
3688                // is the right semantics for "does the drafter know what comes next here" and
3689                // the same row the pick came from. Draft-quality only — verify arbitrates.
3690                if with_prob {
3691                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
3692                }
3693            }
3694        }
3695        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
3696        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
3697        if let Some((out, slot, d2t)) = stream_pack {
3698            e.pack_tok_p(tok_d, p_d, out, slot)?;
3699            if let Some(map) = d2t {
3700                e.tok_map_u32(tok_d, map)?;
3701            }
3702        }
3703        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
3704        if spec_hpost() {
3705            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
3706        } else {
3707            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
3708        }
3709        // advance the draft rope position in-graph.
3710        e.inc_seqlen(pos_d)?;
3711        Ok(())
3712    }
3713
3714    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
3715    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
3716    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
3717    /// Advances `cache.pos` by T.
3718    pub fn decode_step_t(
3719        &self,
3720        e: &Engine,
3721        tokens: &[u32],
3722        pos0: usize,
3723        cache: &mut Cache,
3724    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
3725        if self.is_gemma4_e4b() {
3726            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
3727        }
3728        if self.cfg.gemma4.is_some() {
3729            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
3730        }
3731        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
3732    }
3733
3734    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
3735    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
3736    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
3737    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
3738    pub fn decode_step_t_h(
3739        &self,
3740        e: &Engine,
3741        tokens: &[u32],
3742        pos0: usize,
3743        cache: &mut Cache,
3744    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3745        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
3746    }
3747
3748    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
3749    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
3750    pub fn decode_step_t_h_emb(
3751        &self,
3752        e: &Engine,
3753        tokens: &[u32],
3754        pos0: usize,
3755        cache: &mut Cache,
3756        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3757    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3758        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
3759        Ok((e.dtoh(&logits_d)?, h_seed))
3760    }
3761
3762    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
3763    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
3764    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
3765    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
3766    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
3767    pub fn decode_step_t_h_emb_dev(
3768        &self,
3769        e: &Engine,
3770        tokens: &[u32],
3771        pos0: usize,
3772        cache: &mut Cache,
3773        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3774    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3775        let n_embd = self.cfg.n_embd as usize;
3776        let t = tokens.len();
3777        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
3778        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
3779        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
3780        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
3781        Ok((logits, hs))
3782    }
3783
3784    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
3785    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
3786    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
3787    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
3788    /// retains/copies — they never change what any kernel computes).
3789    fn decode_step_t_core(
3790        &self,
3791        e: &Engine,
3792        tokens: &[u32],
3793        pos0: usize,
3794        cache: &mut Cache,
3795        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3796        mut ckpt: Option<&mut VerifyCkpt>,
3797    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3798        self.decode_step_t_core_stream(
3799            e,
3800            tokens,
3801            pos0,
3802            cache,
3803            embd_dev,
3804            ckpt.take(),
3805            None,
3806            None,
3807            None,
3808            None,
3809        )
3810    }
3811
3812    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
3813    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
3814    fn decode_step_t_core_pipelined(
3815        &self,
3816        e: &Engine,
3817        tokens: &[u32],
3818        pos0: usize,
3819        cache: &mut Cache,
3820        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3821        mut ckpt: Option<&mut VerifyCkpt>,
3822        pipe: &SpecPipeLane,
3823        round: usize,
3824    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3825        let fence = crate::pp::pp_cuts(self.layers.len())
3826            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
3827        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
3828            return Err("two-session speculative pipeline requires the PP verify split".into());
3829        }
3830        let interval_fence = pipe.stage0_begin(round)?;
3831        let ticket = self.verify_stage0_issue(
3832            e,
3833            tokens,
3834            pos0,
3835            cache,
3836            embd_dev,
3837            ckpt.as_deref_mut(),
3838            None,
3839            &fence,
3840            Some(interval_fence),
3841            pipe.trace(round),
3842        )?;
3843        pipe.stage0_end(round);
3844        pipe.stage1_begin(round)?;
3845        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
3846        pipe.verify_end(round);
3847        Ok(result)
3848    }
3849
3850    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
3851    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
3852    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
3853    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
3854    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
3855    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
3856    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
3857    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
3858    #[allow(clippy::too_many_arguments)]
3859    fn decode_step_t_core_stream(
3860        &self,
3861        e: &Engine,
3862        tokens: &[u32],
3863        pos0: usize,
3864        cache: &mut Cache,
3865        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3866        mut ckpt: Option<&mut VerifyCkpt>,
3867        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
3868        pp_pipe: Option<bool>,
3869        vtok_dev: Option<&CudaSlice<u32>>,
3870        graphs: Option<&mut DsparkVerifyGraphs>,
3871    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3872        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
3873        // exactly as the eager and batched steps do. This is the single funnel every verify
3874        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
3875        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
3876        // is untouched.
3877        //
3878        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
3879        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
3880        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
3881        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
3882        // or a placement whose PpNRt fails to build — so a config that would still walk the
3883        // whole trunk on one stream refuses instead of regressing 28x.
3884        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
3885            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
3886                if vtok_dev.is_some() {
3887                    return Err(
3888                        "device-token dspark verify (slice-2 deferred readback) has no PP \
3889                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
3890                         route on one device"
3891                            .into(),
3892                    );
3893                }
3894                return self.decode_step_t_core_ppn(
3895                    e,
3896                    tokens,
3897                    pos0,
3898                    cache,
3899                    embd_dev,
3900                    ckpt.take(),
3901                    stream,
3902                    &fence,
3903                    pp_pipe,
3904                );
3905            }
3906        }
3907        crate::pp::refuse_unsplit_if_remote(
3908            "decode_step_t (spec verify)",
3909            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
3910             split (decode_step_t_core_ppn); or run spec on one device",
3911        )?;
3912        let cfg = &self.cfg;
3913        let n_embd = cfg.n_embd as usize;
3914        let eps = cfg.rms_eps;
3915        let t = tokens.len();
3916        let pos_d = match stream {
3917            Some((_, ctr)) => {
3918                let mut p = e.alloc_uninit::<i32>(t)?;
3919                e.pos_iota(ctr, &mut p, t)?;
3920                p
3921            }
3922            None => {
3923                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
3924                e.htod_i32(&pos_vec)?
3925            }
3926        };
3927
3928        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
3929        let x = match (stream, embd_dev) {
3930            (Some((vtok, _)), Some((g, qt, rb))) => {
3931                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
3932            }
3933            (None, Some((g, qt, rb))) => match vtok_dev {
3934                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
3935                // bit-identical rows to the host-token arm (same per-dtype deq).
3936                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
3937                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3938            },
3939            _ => {
3940                assert!(
3941                    vtok_dev.is_none(),
3942                    "device-token verify requires the resident embed table (embd_dev)"
3943                );
3944                e.htod(&self.embd.gather(n_embd, tokens))?
3945            }
3946        };
3947
3948        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
3949        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
3950        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
3951        let x = self.verify_layers(
3952            e,
3953            x,
3954            0,
3955            self.layers.len(),
3956            &pos_d,
3957            pos0,
3958            t,
3959            cache,
3960            ckpt.take(),
3961            stream,
3962            graphs,
3963        )?;
3964
3965        let mut hn = vbuf(e, t * n_embd)?;
3966        let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
3967        let logits = if serving_head {
3968            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
3969            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
3970            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
3971            // serve one batched numeric class at every live width, including B=1. Keep the
3972            // verify head in that same class; other generic families retain the decode-exact
3973            // head that their run-spec contract pins.
3974            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3975            e.matmul(&self.output, &hn, t)?
3976        } else {
3977            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
3978            e.matmul_decode_exact(&self.output, &hn, t)?
3979        };
3980        // stream: the device pos counter owns position; host mirror reconciles at drain.
3981        if stream.is_none() {
3982            cache.pos += t;
3983        }
3984        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
3985        Ok((logits, if spec_hpost() { hn } else { x }))
3986    }
3987
3988    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
3989    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
3990    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
3991    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
3992    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
3993    /// the payload).
3994    ///
3995    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
3996    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
3997    /// receipts):
3998    ///
3999    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4000    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4001    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4002    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4003    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
4004    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4005    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4006    ///
4007    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4008    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4009    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4010    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4011    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
4012    ///
4013    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4014    ///    sharded loader leaves the table with stage 0 by construction).
4015    ///
4016    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4017    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4018    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4019    ///    model, every round.
4020    ///
4021    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4022    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4023    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4024    /// through the primary context by UVA — the same read the batched serving epilogue's
4025    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4026    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4027    ///
4028    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4029    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4030    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4031    ///
4032    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4033    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4034    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4035    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4036    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4037    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4038    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4039    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4040    #[allow(clippy::too_many_arguments)]
4041    fn decode_step_t_core_ppn(
4042        &self,
4043        e: &Engine,
4044        tokens: &[u32],
4045        pos0: usize,
4046        cache: &mut Cache,
4047        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4048        mut ckpt: Option<&mut VerifyCkpt>,
4049        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4050        fence: &[usize],
4051        pp_pipe: Option<bool>,
4052    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4053        let ticket = self.verify_stage0_issue(
4054            e,
4055            tokens,
4056            pos0,
4057            cache,
4058            embd_dev,
4059            ckpt.as_deref_mut(),
4060            stream,
4061            fence,
4062            pp_pipe,
4063            None,
4064        )?;
4065        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
4066    }
4067
4068    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
4069    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
4070    #[allow(clippy::too_many_arguments)]
4071    fn verify_stage0_issue(
4072        &self,
4073        e: &Engine,
4074        tokens: &[u32],
4075        pos0: usize,
4076        cache: &mut Cache,
4077        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4078        mut ckpt: Option<&mut VerifyCkpt>,
4079        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4080        fence: &[usize],
4081        pp_pipe: Option<bool>,
4082        trace: Option<SpecPipeTraceCtx>,
4083    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
4084        assert!(
4085            !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
4086            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
4087             (the gemma4 arms have their own decode_step_t twins)"
4088        );
4089        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
4090            return Err(
4091                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
4092                 boundary itself is host-staged, but device-resident verify still peer-reads \
4093                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
4094                 serving on this host class; spec requires local per-stage inputs first."
4095                    .into(),
4096            );
4097        }
4098        let rt = crate::pp::PpNRt::get(e)?;
4099        let n_st = fence.len() - 1;
4100        assert_eq!(
4101            rt.n_stages(),
4102            n_st,
4103            "PpNRt stage count {} != fence stages {n_st}",
4104            rt.n_stages()
4105        );
4106        let n_embd = self.cfg.n_embd as usize;
4107        let t = tokens.len();
4108        let payload = t * n_embd;
4109        if pp_pipe.is_some() {
4110            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
4111        }
4112        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
4113        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
4114        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
4115        // the report below names exactly two stages and must never imply it measured middle ones.
4116        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
4117        let pp_started = std::time::Instant::now();
4118        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
4119        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
4120        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
4121        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
4122        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
4123        // stage stream and the wait would self-order into a no-op.
4124        let caller_stream = e.stream();
4125        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
4126        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
4127        // the primary stream still holds queued reads of them — with event tracking elided,
4128        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
4129        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
4130        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
4131        // stage stream behind the caller before enqueueing new stage work.
4132        let reverse_started = std::time::Instant::now();
4133        if pp_pipe != Some(false) {
4134            rt.fence_stages_behind(&caller_stream)?;
4135        }
4136        if pp_pipe == Some(true) {
4137            // Both session verifies must alternate boundary slots even when the ordinary
4138            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
4139            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
4140            rt.prepare_overlap_slots(0, payload)?;
4141        }
4142        if pp_anatomy {
4143            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
4144            // prices any primary-stream rollback/refresh tail inherited from the prior round.
4145            for s in 0..n_st {
4146                let _st = rt.enter(s);
4147                rt.engine(s, e).stream().synchronize()?;
4148            }
4149            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
4150        }
4151
4152        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
4153        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
4154        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4155            match stream {
4156                Some((_, ctr)) => {
4157                    let mut p = es.alloc_uninit::<i32>(t)?;
4158                    es.pos_iota(ctr, &mut p, t)?;
4159                    Ok(p)
4160                }
4161                None => {
4162                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4163                    es.htod_i32(&pos_vec)
4164                }
4165            }
4166        };
4167
4168        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
4169        let slot = {
4170            let _st0 = rt.enter(0);
4171            let e0 = rt.engine(0, e);
4172            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
4173            let stage0_started = std::time::Instant::now();
4174            let pos_d = stage_pos(e0)?;
4175            let x = match (stream, embd_dev) {
4176                (Some((vtok, _)), Some((g, qt, rb))) => {
4177                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4178                }
4179                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4180                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
4181            };
4182            let x = self.verify_layers(
4183                e0,
4184                x,
4185                fence[0],
4186                fence[1],
4187                &pos_d,
4188                pos0,
4189                t,
4190                cache,
4191                ckpt.as_deref_mut(),
4192                stream,
4193                None,
4194            )?;
4195            if pp_anatomy {
4196                e0.stream().synchronize()?;
4197                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
4198            }
4199            let tx_started = std::time::Instant::now();
4200            let slot = if pp_pipe.is_some() {
4201                rt.tx_pipelined(0, &x, payload)?
4202            } else {
4203                rt.tx(0, &x, payload)?
4204            };
4205            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
4206            if pp_anatomy {
4207                e0.stream().synchronize()?;
4208                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
4209            }
4210            slot
4211            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
4212        };
4213
4214        Ok(VerifyBoundaryTicket {
4215            rt,
4216            caller_stream,
4217            slot,
4218            pos0,
4219            t,
4220            payload,
4221            n_st,
4222            pipelined: pp_pipe.is_some(),
4223            pp_anatomy,
4224            pp_started,
4225            reverse_ms,
4226            stage0_ms,
4227            tx_ms,
4228            trace,
4229        })
4230    }
4231
4232    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
4233    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
4234    #[allow(clippy::too_many_arguments)]
4235    fn verify_stage1_finish(
4236        &self,
4237        e: &Engine,
4238        ticket: VerifyBoundaryTicket,
4239        cache: &mut Cache,
4240        mut ckpt: Option<&mut VerifyCkpt>,
4241        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4242        fence: &[usize],
4243        publish_to_caller: bool,
4244    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4245        let VerifyBoundaryTicket {
4246            rt,
4247            caller_stream,
4248            slot,
4249            pos0,
4250            t,
4251            payload,
4252            n_st,
4253            pipelined,
4254            pp_anatomy,
4255            pp_started,
4256            reverse_ms,
4257            stage0_ms,
4258            tx_ms,
4259            trace,
4260        } = ticket;
4261        let n_embd = self.cfg.n_embd as usize;
4262        let eps = self.cfg.rms_eps;
4263        let mut slot = slot;
4264        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
4265        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4266            match stream {
4267                Some((_, ctr)) => {
4268                    let mut p = es.alloc_uninit::<i32>(t)?;
4269                    es.pos_iota(ctr, &mut p, t)?;
4270                    Ok(p)
4271                }
4272                None => {
4273                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4274                    es.htod_i32(&pos_vec)
4275                }
4276            }
4277        };
4278
4279        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
4280        for s in 1..n_st - 1 {
4281            let _st = rt.enter(s);
4282            let es = rt.engine(s, e);
4283            let pos_d = stage_pos(es)?;
4284            let x = rt.rx(s - 1, slot, payload)?;
4285            let x = self.verify_layers(
4286                es,
4287                x,
4288                fence[s],
4289                fence[s + 1],
4290                &pos_d,
4291                pos0,
4292                t,
4293                cache,
4294                ckpt.as_deref_mut(),
4295                stream,
4296                None,
4297            )?;
4298            slot = if pipelined {
4299                rt.tx_pipelined(s, &x, payload)?
4300            } else {
4301                rt.tx(s, &x, payload)?
4302            };
4303        }
4304
4305        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
4306        let _stl = rt.enter(n_st - 1);
4307        let el = rt.engine(n_st - 1, e);
4308        let pos_d = stage_pos(el)?;
4309        let rx_started = std::time::Instant::now();
4310        let x = rt.rx(n_st - 2, slot, payload)?;
4311        if pp_anatomy {
4312            el.stream().synchronize()?;
4313            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
4314        }
4315        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
4316        let stage1_started = std::time::Instant::now();
4317        let x = self.verify_layers(
4318            el,
4319            x,
4320            fence[n_st - 1],
4321            fence[n_st],
4322            &pos_d,
4323            pos0,
4324            t,
4325            cache,
4326            ckpt.as_deref_mut(),
4327            stream,
4328            None,
4329        )?;
4330
4331        let mut hn = vbuf(el, payload)?;
4332        let logits = if self.cfg.step35.is_some() {
4333            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
4334            // Verify must not switch numeric class merely because the same session speculates.
4335            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4336            el.matmul(&self.output, &hn, t)?
4337        } else {
4338            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4339            el.matmul_decode_exact(&self.output, &hn, t)?
4340        };
4341        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
4342        if pp_anatomy {
4343            el.stream().synchronize()?;
4344            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
4345        }
4346        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
4347        // stream. Order the caller's stream behind that work before the buffers escape this
4348        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
4349        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
4350        // the following arm's KV in the same process).
4351        if publish_to_caller {
4352            rt.publish_to(n_st - 1, &caller_stream)?;
4353        }
4354        if pp_anatomy {
4355            if publish_to_caller {
4356                caller_stream.synchronize()?;
4357            }
4358            eprintln!(
4359                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
4360                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
4361                pp_started.elapsed().as_secs_f64() * 1e3,
4362            );
4363        }
4364        // stream: the device pos counter owns position; host mirror reconciles at drain.
4365        if stream.is_none() {
4366            cache.pos += t;
4367        }
4368        Ok((logits, if spec_hpost() { hn } else { x }))
4369    }
4370
4371    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
4372    ///
4373    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
4374    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
4375    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
4376    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
4377    /// bytes when a request moves from batched plain serving into speculative verify. Run the
4378    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
4379    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
4380    /// every norm/projection/FFN uses exactly the live serving dispatch.
4381    #[allow(clippy::too_many_arguments)]
4382    fn step35_verify_batch_layers(
4383        &self,
4384        e: &Engine,
4385        mut x: CudaSlice<f32>,
4386        lo: usize,
4387        hi: usize,
4388        pos0: usize,
4389        t: usize,
4390        cache: &mut Cache,
4391    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4392        let n_embd = self.cfg.n_embd as usize;
4393        self.cfg
4394            .step35
4395            .as_ref()
4396            .ok_or("step35 verify batch requires step35 cfg")?;
4397        let mut ph_last = std::time::Instant::now();
4398        for il in lo..hi {
4399            let mut next = e.uninit(t * n_embd)?;
4400            for r in 0..t {
4401                let mut row = e.uninit(n_embd)?;
4402                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4403                // The caller owns this verify's position. During controller overlap, cache.pos
4404                // still describes generation N while this stage-0 walk belongs to N+1.
4405                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4406                let mut one = [&mut *cache];
4407                let out = self.step35_decode_batch_layers(
4408                    e,
4409                    row,
4410                    &mut one,
4411                    &row_pos,
4412                    il,
4413                    il + 1,
4414                    &mut ph_last,
4415                )?;
4416                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4417            }
4418            self.dflash_tap(e, cache, il, &next, t)?;
4419            x = next;
4420        }
4421        Ok(x)
4422    }
4423
4424    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
4425    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
4426    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
4427    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
4428    /// prefix-keep, not all-or-nothing).
4429    pub(crate) fn dspark_verify_t_am(
4430        &self,
4431        e: &Engine,
4432        tokens: &[u32],
4433        pos0: usize,
4434        cache: &mut Cache,
4435    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4436        let (logits, _hn) = self.decode_step_t_core_stream(
4437            e, tokens, pos0, cache, None, None, None, None, None, None,
4438        )?;
4439        let t = tokens.len();
4440        let v = self.output.out_features();
4441        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4442        for r in 0..t {
4443            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4444        }
4445        Ok(e.dtoh_u32(&am_d)?)
4446    }
4447
4448    /// DSpark verify with the MTP column-stash armed: identical forward to
4449    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
4450    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
4451    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
4452    pub(crate) fn dspark_verify_t_am_ckpt(
4453        &self,
4454        e: &Engine,
4455        tokens: &[u32],
4456        pos0: usize,
4457        cache: &mut Cache,
4458    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4459        let mut ck = VerifyCkpt::new(self.layers.len());
4460        let (logits, _hn) = self.decode_step_t_core_stream(
4461            e,
4462            tokens,
4463            pos0,
4464            cache,
4465            None,
4466            Some(&mut ck),
4467            None,
4468            None,
4469            None,
4470            None,
4471        )?;
4472        let t = tokens.len();
4473        let v = self.output.out_features();
4474        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4475        for r in 0..t {
4476            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4477        }
4478        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
4479    }
4480
4481    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
4482    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
4483    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
4484    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
4485    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
4486    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
4487    pub(crate) fn dspark_verify_t_am_ckpt_dev(
4488        &self,
4489        e: &Engine,
4490        vtok: &CudaSlice<u32>,
4491        t: usize,
4492        pos0: usize,
4493        cache: &mut Cache,
4494        embd_dev: (&CudaSlice<u8>, i32, usize),
4495        graphs: Option<&mut DsparkVerifyGraphs>,
4496    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
4497        debug_assert!(
4498            vtok.len() >= t,
4499            "verify window exceeds the device token buffer"
4500        );
4501        // The slab flag is a per-round statement: clear it here so a verify that never
4502        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
4503        // stale `true` steering the commit at slabs the round never wrote.
4504        let mut graphs = graphs;
4505        if let Some(g) = graphs.as_deref_mut() {
4506            g.round_slab = false;
4507        }
4508        let mut ck = VerifyCkpt::new(self.layers.len());
4509        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
4510        // arm's established pattern — spec.rs stream-mode verify does the same).
4511        let dummy = vec![0u32; t];
4512        let (logits, _hn) = self.decode_step_t_core_stream(
4513            e,
4514            &dummy,
4515            pos0,
4516            cache,
4517            Some(embd_dev),
4518            Some(&mut ck),
4519            None,
4520            None,
4521            Some(vtok),
4522            graphs,
4523        )?;
4524        let v = self.output.out_features();
4525        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4526        for r in 0..t {
4527            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4528        }
4529        Ok((am_d, DsparkVerifyCkpt(ck)))
4530    }
4531
4532    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
4533    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
4534    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
4535    pub(crate) fn dspark_commit_prefix(
4536        &self,
4537        e: &Engine,
4538        cache: &mut Cache,
4539        snap: &crate::cache::CacheSnapshot,
4540        ckpt: &DsparkVerifyCkpt,
4541        keep: usize,
4542    ) -> Result<(), Box<dyn std::error::Error>> {
4543        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
4544    }
4545
4546    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
4547    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
4548    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
4549    /// from the stash of column keep-1), slab-addressed and batched into two copy
4550    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
4551    pub(crate) fn dspark_commit_prefix_slab(
4552        &self,
4553        e: &Engine,
4554        cache: &mut Cache,
4555        snap: &crate::cache::CacheSnapshot,
4556        ctx: &DsparkVerifyGraphs,
4557        keep: usize,
4558    ) -> Result<(), Box<dyn std::error::Error>> {
4559        use cudarc::driver::DevicePtr;
4560        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
4561        let mut conv_src: Vec<u64> = Vec::new();
4562        let mut ssm_src: Vec<u64> = Vec::new();
4563        let mut conv_dst: Vec<u64> = Vec::new();
4564        let mut ssm_dst: Vec<u64> = Vec::new();
4565        for il in 0..self.layers.len() {
4566            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
4567                kvl.len = saved + keep;
4568                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
4569            }
4570            if let Some(rl) = cache.recur[il].as_ref() {
4571                let (pc, ps, _cw, _sw) = ctx
4572                    .slab_row(e, il, keep - 1)
4573                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
4574                conv_src.push(pc);
4575                ssm_src.push(ps);
4576                let st = &e.gpu.stream();
4577                let (dc, _g0) = rl.conv_state.device_ptr(st);
4578                let (ds, _g1) = rl.ssm_state.device_ptr(st);
4579                conv_dst.push(dc as u64);
4580                ssm_dst.push(ds as u64);
4581            }
4582        }
4583        let n = conv_src.len();
4584        if n > 0 {
4585            if state_copy_batch_on() {
4586                let mut tt = vec![0u64; 2 * n];
4587                tt[..n].copy_from_slice(&conv_src);
4588                tt[n..].copy_from_slice(&conv_dst);
4589                let ct = e.htod_u64(&tt)?;
4590                tt[..n].copy_from_slice(&ssm_src);
4591                tt[n..].copy_from_slice(&ssm_dst);
4592                let st = e.htod_u64(&tt)?;
4593                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
4594                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
4595            } else {
4596                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
4597                let row = keep - 1;
4598                for il in 0..self.layers.len() {
4599                    let Some(rl) = cache.recur[il].as_mut() else {
4600                        continue;
4601                    };
4602                    let k = ctx.lin_pos[&il];
4603                    {
4604                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
4605                        let win = sv.slice(row * cw..(row + 1) * cw);
4606                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
4607                    }
4608                    {
4609                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
4610                        let win = sv.slice(row * sw..(row + 1) * sw);
4611                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
4612                    }
4613                }
4614            }
4615        }
4616        cache.pos = snap.pos + keep;
4617        Ok(())
4618    }
4619
4620    /// Qwen35-family verify trunk in the live serving numeric class.
4621    ///
4622    /// Serving intentionally keeps this architecture in the generic batched program even at
4623    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
4624    ///
4625    /// Two arms, one numeric class:
4626    /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
4627    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
4628    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
4629    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
4630    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
4631    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
4632    ///   program its isolated serving step would). One weight read per layer per round
4633    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
4634    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
4635    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
4636    ///   serving layer body, preserving single-session autoregressive cache order (the
4637    ///   correctness reference; also the rollback seam for the t-parallel arm).
4638    ///
4639    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
4640    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
4641    #[allow(clippy::too_many_arguments)]
4642    fn qwen35_verify_batch_layers(
4643        &self,
4644        e: &Engine,
4645        x: CudaSlice<f32>,
4646        lo: usize,
4647        hi: usize,
4648        pos0: usize,
4649        t: usize,
4650        cache: &mut Cache,
4651        ckpt: Option<&mut VerifyCkpt>,
4652        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4653        graphs: Option<&mut DsparkVerifyGraphs>,
4654    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4655        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
4656        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
4657        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
4658        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
4659        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
4660        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
4661        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
4662            || !matches!(
4663                self.cfg.arch,
4664                memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
4665            )
4666            || t > 16;
4667        if rowwise {
4668            if stream.is_some() {
4669                // rowwise replays per row with host cache.pos — irreconcilable with a
4670                // device position counter. Burst callers must keep t <= 16 and the
4671                // ROWWISE env unset; refusing beats silently mispositioned rows.
4672                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
4673                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
4674                    .into());
4675            }
4676            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
4677        } else {
4678            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
4679        }
4680    }
4681
4682    /// The per-row correctness reference: replay each verify row through the authoritative
4683    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
4684    #[allow(clippy::too_many_arguments)]
4685    fn qwen35_verify_rowwise(
4686        &self,
4687        e: &Engine,
4688        mut x: CudaSlice<f32>,
4689        lo: usize,
4690        hi: usize,
4691        pos0: usize,
4692        t: usize,
4693        cache: &mut Cache,
4694        mut ckpt: Option<&mut VerifyCkpt>,
4695    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4696        let n_embd = self.cfg.n_embd as usize;
4697        let saved_pos = cache.pos;
4698        let mut ph_last = std::time::Instant::now();
4699        for il in lo..hi {
4700            let mut next = e.uninit(t * n_embd)?;
4701            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
4702                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4703                    Some(Vec::with_capacity(t - 1))
4704                } else {
4705                    None
4706                };
4707            for r in 0..t {
4708                cache.pos = pos0 + r;
4709                let mut row = e.uninit(n_embd)?;
4710                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4711                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4712                let mut one = [&mut *cache];
4713                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
4714                let out = match self.decode_batch_layers(
4715                    e,
4716                    row,
4717                    &mut one,
4718                    &ctx,
4719                    &row_pos,
4720                    &mut ph_last,
4721                ) {
4722                    Ok(out) => out,
4723                    Err(error) => {
4724                        cache.pos = saved_pos;
4725                        return Err(error);
4726                    }
4727                };
4728                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4729                if r + 1 < t {
4730                    if let Some(states) = col_states.as_mut() {
4731                        let recur = cache.recur[il]
4732                            .as_ref()
4733                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
4734                        states.push((
4735                            e.clone_dtod(&recur.conv_state)?,
4736                            e.clone_dtod(&recur.ssm_state)?,
4737                        ));
4738                    }
4739                }
4740            }
4741            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
4742                checkpoint.cols[il] = Some(states);
4743            }
4744            x = next;
4745        }
4746        cache.pos = saved_pos;
4747        Ok(x)
4748    }
4749
4750    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
4751    ///
4752    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
4753    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
4754    /// pins the serving batch tier already carries:
4755    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
4756    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
4757    ///     alone;
4758    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
4759    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
4760    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
4761    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
4762    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
4763    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
4764    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
4765    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
4766    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
4767    /// program its isolated B=1 serving step would.
4768    ///
4769    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
4770    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
4771    #[allow(clippy::too_many_arguments)]
4772    fn qwen35_verify_tparallel(
4773        &self,
4774        e: &Engine,
4775        mut x: CudaSlice<f32>,
4776        lo: usize,
4777        hi: usize,
4778        pos0: usize,
4779        t: usize,
4780        cache: &mut Cache,
4781        mut ckpt: Option<&mut VerifyCkpt>,
4782        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4783        mut graphs: Option<&mut DsparkVerifyGraphs>,
4784    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4785        use cudarc::driver::DevicePtr;
4786        let cfg = &self.cfg;
4787        let n_embd = cfg.n_embd as usize;
4788        let eps = cfg.rms_eps;
4789        let head_dim_global = cfg.head_dim_k as usize;
4790        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
4791        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
4792        let pos_d = match stream {
4793            Some((_, ctr)) => {
4794                let mut p = e.alloc_uninit::<i32>(t)?;
4795                e.pos_iota(ctr, &mut p, t)?;
4796                p
4797            }
4798            None => {
4799                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
4800                e.htod_i32(&pos_host)?
4801            }
4802        };
4803        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
4804        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
4805        let pos_rows: Vec<CudaSlice<i32>> = match stream {
4806            Some((_, ctr)) => (0..t)
4807                .map(|r| {
4808                    let mut b = e.alloc_uninit::<i32>(1)?;
4809                    e.i32_copy_add(ctr, &mut b, r as i32)?;
4810                    Ok(b)
4811                })
4812                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
4813            None => (0..t)
4814                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
4815                .collect::<Result<_, _>>()?,
4816        };
4817        let seqs_append =
4818            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
4819        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
4820
4821        // Engine-bundle slice 3: with a graphs ctx armed, runs of consecutive LINEAR layers
4822        // replay per-(segment, vt) CUDA graphs (captured from the SAME
4823        // `qwen35_tparallel_linear_layer` body the eager arm runs — no second copy of the
4824        // math). The full-attention layers stay eager: their per-row append/fa arm picks are
4825        // t_kv-driven (the straddle law) and belong to the exec-update extension, not this
4826        // slice. Pointer tables are refreshed once per verify (gdn ping-pong moves handles).
4827        // Merge guard (v0.98 train): the ROUND-STREAM arm (lane/draftcost-moe, device
4828        // position counter) and the dspark verify graphs (engine-bundle slice 3) have no
4829        // common caller — stream rides the qwen35moe burst, graphs ride the dspark route.
4830        // If a future caller arms both, refuse loudly instead of silently dropping the
4831        // graphs ctx (the stream linear arm takes linear_attn_verify_t, not the graphed
4832        // segment body).
4833        if stream.is_some() && graphs.is_some() {
4834            return Err(
4835                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
4836                        cannot arm together"
4837                    .into(),
4838            );
4839        }
4840        if let Some(g) = graphs.as_deref_mut() {
4841            g.refresh_tables(e, cache)?;
4842            g.round_slab = false;
4843        }
4844        let mut il = lo;
4845        while il < hi {
4846            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
4847                let mut end = il;
4848                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
4849                    end += 1;
4850                }
4851                let g = graphs.as_deref_mut().expect("checked above");
4852                x = g.run_segment(self, e, il, end, &x, t, cache)?;
4853                g.round_slab = true;
4854                il = end;
4855                continue;
4856            }
4857            let layer = &self.layers[il];
4858            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
4859                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
4860                // Under ROUND-STREAM the linear layers ride the match's stream arm below
4861                // (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
4862                x = self.qwen35_tparallel_linear_layer(
4863                    e,
4864                    il,
4865                    &x,
4866                    t,
4867                    cache,
4868                    ckpt.as_deref_mut(),
4869                    None,
4870                    None,
4871                )?;
4872                il += 1;
4873                continue;
4874            }
4875            // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
4876            let anorm = layer.attn_norm.float_data();
4877            let mut xn = e.uninit(t * n_embd)?;
4878            e.rms_norm(&x, anorm, &mut xn, n_embd, t, eps)?;
4879            let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
4880
4881            let mixed: CudaSlice<f32> = match &layer.mixer {
4882                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4883                Mixer::Full(fa) => {
4884                    let geometry = cfg.full_attention_geometry_at(il as u32);
4885                    let n_head = geometry.n_head as usize;
4886                    let n_head_kv = geometry.n_head_kv as usize;
4887                    let head_dim = geometry.head_dim_k as usize;
4888                    let rope_dims = geometry.n_rot as usize;
4889                    let rope_base = geometry.rope_base;
4890                    let scale = geometry.attention_scale();
4891                    // Batched projections: one weight read serves all T rows.
4892                    let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?;
4893                    let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?;
4894                    let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?;
4895                    let gated =
4896                        geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4897                    let (mut q, gate) = if gated {
4898                        let mut qs = e.uninit(t * n_head * head_dim)?;
4899                        let mut gs = e.uninit(t * n_head * head_dim)?;
4900                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
4901                        (qs, Some(gs))
4902                    } else {
4903                        (qf, None)
4904                    };
4905                    let mut qn = e.uninit(t * n_head * head_dim)?;
4906                    e.rms_norm(
4907                        &q,
4908                        fa.q_norm.float_data(),
4909                        &mut qn,
4910                        head_dim,
4911                        t * n_head,
4912                        eps,
4913                    )?;
4914                    q = qn;
4915                    let mut kn = e.uninit(t * n_head_kv * head_dim)?;
4916                    e.rms_norm(
4917                        &k,
4918                        fa.k_norm.float_data(),
4919                        &mut kn,
4920                        head_dim,
4921                        t * n_head_kv,
4922                        eps,
4923                    )?;
4924                    k = kn;
4925                    e.rope_neox(
4926                        &mut q, &pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
4927                    )?;
4928                    e.rope_neox(
4929                        &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4930                    )?;
4931
4932                    // Per-row append + attend: row r sees rows 0..r in KV (causal within the
4933                    // draft), each through the b_n=1 serving kernels at its own t_kv.
4934                    let q_dim = n_head * head_dim;
4935                    let kv_dim = n_head_kv * head_dim;
4936                    let mut attn = e.uninit(t * q_dim)?;
4937                    let (kdk, kdv, ktb, vtb, kv_view) = {
4938                        let kvl = cache.kv[il].as_ref().unwrap();
4939                        let s = &e.gpu.stream();
4940                        let (pk, _g) = kvl.k.device_ptr(s);
4941                        let (pv, _g2) = kvl.v.device_ptr(s);
4942                        (
4943                            kvl.kv_dim_k,
4944                            kvl.kv_dim_v,
4945                            kvl.k_tok_bytes,
4946                            kvl.v_tok_bytes,
4947                            e.htod_u64(&[pk as u64, pv as u64])?,
4948                        )
4949                    };
4950                    if let Some((_, ctr)) = stream {
4951                        // STREAM ARM (2b): one batched dc append + the multi-row dc attention
4952                        // — the generic stream arm's exact shape (rows kernels are pinned
4953                        // byte-identical to the per-row programs by kernel-check). Host len
4954                        // stays a stale lower bound; the burst drain reconciles it.
4955                        let kvl = cache.kv[il].as_mut().unwrap();
4956                        e.append_kv_quantized_rows_dc(
4957                            &k,
4958                            &v,
4959                            &mut kvl.k,
4960                            &mut kvl.v,
4961                            ctr,
4962                            t,
4963                            kdk,
4964                            kdv,
4965                            ktb,
4966                            vtb,
4967                            Engine::kv_fp8_on(),
4968                        )?;
4969                        let upper = (kvl.len + t + 64).min(cache.max_ctx);
4970                        let k_view = e.view_u8(&kvl.k, upper * ktb);
4971                        let v_view = e.view_u8(&kvl.v, upper * vtb);
4972                        e.fa_decode_rows_dc(
4973                            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr,
4974                            upper, t, scale, ktb, vtb, 0, false,
4975                        )?;
4976                    } else {
4977                        for r in 0..t {
4978                            // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
4979                            // whose row 0 is this row (arithmetic-free materialization copies,
4980                            // same as decode's per-seq fallback arm).
4981                            let mut k_row = e.uninit(kv_dim)?;
4982                            e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
4983                            let mut v_row = e.uninit(kv_dim)?;
4984                            e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
4985                            let pos_row = &pos_rows[r];
4986                            let kvl = cache.kv[il].as_mut().unwrap();
4987                            if seqs_append {
4988                                e.append_kv_quantized_seqs(
4989                                    &k_row,
4990                                    &v_row,
4991                                    &kv_view.slice(0..2),
4992                                    pos_row,
4993                                    1,
4994                                    kdk,
4995                                    kdv,
4996                                    ktb,
4997                                    vtb,
4998                                )?;
4999                                kvl.len += 1;
5000                            } else {
5001                                e.append_kv_quantized_view(
5002                                    &k_row.slice(0..kv_dim),
5003                                    &v_row.slice(0..kv_dim),
5004                                    &mut kvl.k,
5005                                    &mut kvl.v,
5006                                    kvl.len,
5007                                    kvl.kv_dim_k,
5008                                    kvl.kv_dim_v,
5009                                    kvl.k_tok_bytes,
5010                                    kvl.v_tok_bytes,
5011                                    Engine::kv_fp8_on(),
5012                                )?;
5013                                kvl.len += 1;
5014                            }
5015                            let t_kv = kvl.len;
5016                            let mut q_row = e.uninit(q_dim)?;
5017                            e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
5018                            let mut a_row = e.uninit(q_dim)?;
5019                            if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
5020                                let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
5021                                e.fa_decode_batch_seqs_v4(
5022                                    &q_row,
5023                                    &kv_view.slice(0..2),
5024                                    pos_row,
5025                                    &mut a_row,
5026                                    head_dim,
5027                                    n_head,
5028                                    n_head_kv,
5029                                    1,
5030                                    t_kv,
5031                                    scale,
5032                                    sp0_r,
5033                                    ktb,
5034                                    vtb,
5035                                )?;
5036                            } else {
5037                                let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
5038                                let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
5039                                let mut a_view = a_row.slice_mut(0..q_dim);
5040                                e.fa_decode_kvmod_view(
5041                                    &q_row.slice(0..q_dim),
5042                                    &k_view,
5043                                    &v_view,
5044                                    &mut a_view,
5045                                    head_dim,
5046                                    n_head,
5047                                    n_head_kv,
5048                                    t_kv,
5049                                    scale,
5050                                    kvl.k_tok_bytes,
5051                                    kvl.v_tok_bytes,
5052                                    Engine::kv_fp8_on(),
5053                                )?;
5054                            }
5055                            e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
5056                        }
5057                    }
5058
5059                    // Output gate (element-wise) + o-proj at m=T.
5060                    let attn_g = match &gate {
5061                        Some(g) => {
5062                            let n = t * q_dim;
5063                            let mut gsig = e.uninit(n)?;
5064                            e.sigmoid(g, &mut gsig, n)?;
5065                            let mut ag = e.uninit(n)?;
5066                            e.mul(&attn, &gsig, &mut ag, n)?;
5067                            ag
5068                        }
5069                        None => attn,
5070                    };
5071                    e.matmul(&fa.wo, &attn_g, t)?
5072                }
5073                // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
5074                // per-row serving-kernel chain cannot run (host state swaps keyed on host
5075                // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
5076                // rebuild — the per-row chain only produces per-column clones). GDN rides
5077                // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
5078                // and its one-scan recurrence is pinned bit-identical to T chained T=1
5079                // steps (its header + kernel-check). Position-independent, so no counter
5080                // plumbing is needed. Guards mirror the generic call site exactly.
5081                Mixer::Linear(la) if stream.is_some() => {
5082                    if !(t >= 3 || (t == 2 && spec_m2()))
5083                        || !self.mixer_in_q8_1_fast(e, &layer.mixer)
5084                        || !e.uses_q8_1_fast(&la.ssm_out)
5085                    {
5086                        return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
5087                                    (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
5088                            .into());
5089                    }
5090                    let want = ckpt.is_some();
5091                    let (out, stash) = self.linear_attn_verify_t(
5092                        e,
5093                        la,
5094                        &xn,
5095                        Some((&hq, &hd)),
5096                        t,
5097                        cache,
5098                        il,
5099                        want,
5100                    )?;
5101                    if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
5102                        ck.gdn[il] = Some(st);
5103                    }
5104                    out
5105                }
5106                Mixer::Linear(_) => {
5107                    unreachable!("linear layers ride qwen35_tparallel_linear_layer")
5108                }
5109            };
5110
5111            // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
5112            let pnorm = layer.post_attn_norm.float_data();
5113            let mut x1 = e.uninit(t * n_embd)?;
5114            let mut zn = e.uninit(t * n_embd)?;
5115            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
5116            let ffn_out = match &layer.ffn {
5117                crate::hybrid::Ffn::Dense {
5118                    ffn_gate,
5119                    ffn_up,
5120                    ffn_down,
5121                } => {
5122                    assert!(
5123                        self.cfg.m3.is_none(),
5124                        "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
5125                    );
5126                    let n_ff = ffn_gate.out_features();
5127                    let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
5128                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
5129                    let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
5130                    let mut act = e.uninit(t * n_ff)?;
5131                    e.silu_mul(&g, &u, &mut act, t * n_ff)?;
5132                    let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5133                    e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5134                }
5135                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
5136            };
5137            let mut x2 = e.uninit(t * n_embd)?;
5138            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5139            // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
5140            self.dflash_tap(e, cache, il, &x2, t)?;
5141            x = x2;
5142            il += 1;
5143        }
5144        Ok(x)
5145    }
5146
5147    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
5148    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
5149    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
5150    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
5151    /// bit-identical by construction:
5152    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
5153    ///   the device sequence is driven entirely by the 6-entry pointer table, which
5154    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
5155    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
5156    ///   legacy post-swap clone read.
5157    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
5158    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
5159    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
5160    /// None builds the per-verify table exactly as before.
5161    #[allow(clippy::too_many_arguments)]
5162    fn qwen35_tparallel_linear_layer(
5163        &self,
5164        e: &Engine,
5165        il: usize,
5166        x: &CudaSlice<f32>,
5167        t: usize,
5168        cache: &mut Cache,
5169        mut ckpt: Option<&mut VerifyCkpt>,
5170        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
5171        table_src: Option<(&CudaSlice<u64>, usize)>,
5172    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5173        use cudarc::driver::DevicePtr;
5174        let cfg = &self.cfg;
5175        let n_embd = cfg.n_embd as usize;
5176        let eps = cfg.rms_eps;
5177        let layer = &self.layers[il];
5178        let Mixer::Linear(la) = &layer.mixer else {
5179            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
5180        };
5181        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
5182        let anorm = layer.attn_norm.float_data();
5183        let mut xn = e.uninit(t * n_embd)?;
5184        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
5185        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
5186
5187        let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
5188        let d_state = ssm.state_size as usize;
5189        let num_k = ssm.group_count as usize;
5190        let num_v = ssm.time_step_rank as usize;
5191        let d_conv = ssm.conv_kernel as usize;
5192        let key_dim = d_state * num_k;
5193        let value_dim = d_state * num_v;
5194        let conv_dim = key_dim * 2 + value_dim;
5195        let gdn_scale = 1.0 / (d_state as f32).sqrt();
5196
5197        // ---- batched projections: one weight read for all T rows ----
5198        let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?;
5199        let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?;
5200        let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?;
5201        let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?;
5202        let beta_w = la.ssm_beta.out_features();
5203        let alpha_w = la.ssm_alpha.out_features();
5204        let qkv_w = la.wqkv.out_features();
5205
5206        // ---- per-row state chain through the b_n=1 serving kernels ----
5207        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
5208        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
5209        let table_local: Option<CudaSlice<u64>> = match table_src {
5210            Some(_) => None,
5211            None => {
5212                let rl = cache.recur[il].as_ref().unwrap();
5213                let s = &e.gpu.stream();
5214                let (pc, _g0) = rl.conv_state.device_ptr(s);
5215                let (p0, _g1) = rl.ssm_state.device_ptr(s);
5216                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
5217                Some(e.htod_u64(&[
5218                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
5219                ])?)
5220            }
5221        };
5222        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
5223            Some((tb, off)) => (tb, off),
5224            None => (table_local.as_ref().unwrap(), 0),
5225        };
5226        let mut o_all = e.uninit(t * value_dim)?;
5227        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5228            if ckpt.is_some() && stash.is_none() && t >= 2 {
5229                Some(Vec::with_capacity(t - 1))
5230            } else {
5231                None
5232            };
5233        let mut stash = stash;
5234        // Per-row scratch reused across rows (uninit is cheap but not free at
5235        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
5236        // [T, ...] buffers — zero arithmetic-free copies in this loop.
5237        let mut conv_out = e.uninit(conv_dim)?;
5238        let mut q_l2 = e.uninit(value_dim)?;
5239        let mut k_l2 = e.uninit(value_dim)?;
5240        let mut v_gd = e.uninit(value_dim)?;
5241        let mut beta_b = e.uninit(num_v)?;
5242        let mut g_log = e.uninit(num_v)?;
5243        for r in 0..t {
5244            let base = toff + if r % 2 == 0 { 0 } else { 3 };
5245            let conv_view = table.slice(base..base + 1);
5246            let in_view = table.slice(base + 1..base + 2);
5247            let out_view = table.slice(base + 2..base + 3);
5248            e.ssm_conv1d_fused_decode_b_view(
5249                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
5250                &conv_view,
5251                la.ssm_conv1d.float_data(),
5252                &mut conv_out,
5253                conv_dim,
5254                d_conv,
5255                1,
5256            )?;
5257            e.gdn_prep_decode_b_view(
5258                &conv_out,
5259                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
5260                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
5261                la.ssm_dt.float_data(),
5262                la.ssm_a.float_data(),
5263                &mut q_l2,
5264                &mut k_l2,
5265                &mut v_gd,
5266                &mut beta_b,
5267                &mut g_log,
5268                d_state,
5269                num_v,
5270                num_k,
5271                key_dim,
5272                eps,
5273                conv_dim,
5274                1,
5275            )?;
5276            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
5277            e.gdn_scan_s128_batched_view(
5278                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
5279                gdn_scale,
5280            )?;
5281            if r + 1 < t {
5282                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
5283                // odd rows write s0 — the same physical state the legacy post-swap
5284                // canonical clone read.
5285                let rl = cache.recur[il]
5286                    .as_ref()
5287                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
5288                let ssm_src = if r % 2 == 0 {
5289                    &rl.ssm_state_alt
5290                } else {
5291                    &rl.ssm_state
5292                };
5293                match stash.as_mut() {
5294                    Some((conv_slab, ssm_slab)) => {
5295                        // BOTH stash reads go through the pointer table at run time: the
5296                        // ssm handles ping-pong between rounds, and the ctx (with its
5297                        // captured graphs) outlives the Cache — a fresh generation's
5298                        // conv/ssm buffers land at new addresses that only the per-round
5299                        // table refresh knows. A baked direct copy would read freed
5300                        // memory (parity was the slice-3 smoke divergence; cache
5301                        // lifetime is the cross-generation twin).
5302                        e.copy_indirect_src_f32(
5303                            &conv_view,
5304                            conv_slab,
5305                            r * conv_dim * (d_conv - 1),
5306                            conv_dim * (d_conv - 1),
5307                        )?;
5308                        // The ssm handles PING-PONG between rounds: a captured direct
5309                        // copy would bake the capture-time physical buffer and read the
5310                        // wrong parity after any odd-vt round (the slice-3 smoke
5311                        // divergence). Read the src address from row r's OUT table
5312                        // entry at run time — the same entry the scan just wrote.
5313                        e.copy_indirect_src_f32(
5314                            &out_view,
5315                            ssm_slab,
5316                            r * d_state * d_state * num_v,
5317                            d_state * d_state * num_v,
5318                        )?;
5319                    }
5320                    None => {
5321                        if let Some(states) = col_states.as_mut() {
5322                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
5323                        }
5324                    }
5325                }
5326            }
5327        }
5328        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
5329        // handle motion is identical and the device sequence never read the handles.
5330        if t % 2 == 1 {
5331            let rl = cache.recur[il].as_mut().unwrap();
5332            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5333        }
5334        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
5335            checkpoint.cols[il] = Some(states);
5336        }
5337
5338        // ---- batched gated norm + out-projection at m=T ----
5339        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
5340            let (gq, gd) = e.gated_rmsnorm_q8_1(
5341                &o_all,
5342                la.ssm_norm.float_data(),
5343                &z,
5344                d_state,
5345                t * num_v,
5346                eps,
5347            )?;
5348            let g0 = e.zeros(0)?;
5349            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
5350        } else {
5351            let mut gn = e.uninit(t * value_dim)?;
5352            e.gated_rmsnorm(
5353                &o_all,
5354                la.ssm_norm.float_data(),
5355                &z,
5356                &mut gn,
5357                d_state,
5358                t * num_v,
5359                eps,
5360            )?;
5361            e.matmul(&la.ssm_out, &gn, t)?
5362        };
5363
5364        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
5365        let pnorm = layer.post_attn_norm.float_data();
5366        let mut x1 = e.uninit(t * n_embd)?;
5367        let mut zn = e.uninit(t * n_embd)?;
5368        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
5369        let ffn_out = match &layer.ffn {
5370            crate::hybrid::Ffn::Dense {
5371                ffn_gate,
5372                ffn_up,
5373                ffn_down,
5374            } => {
5375                assert!(
5376                    self.cfg.m3.is_none(),
5377                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
5378                );
5379                let n_ff = ffn_gate.out_features();
5380                let (zq, zd) = e.quantize_q8_1(&zn, t, n_embd)?;
5381                let g = e.matmul_pre(ffn_gate, &zq, &zd, &zn, t)?;
5382                let u = e.matmul_pre(ffn_up, &zq, &zd, &zn, t)?;
5383                let mut act = e.uninit(t * n_ff)?;
5384                e.silu_mul(&g, &u, &mut act, t * n_ff)?;
5385                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5386                e.matmul_pre(ffn_down, &aq, &ad, &act, t)?
5387            }
5388            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
5389        };
5390        let mut x2 = e.uninit(t * n_embd)?;
5391        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5392        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
5393        self.dflash_tap(e, cache, il, &x2, t)?;
5394        Ok(x2)
5395    }
5396
5397    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
5398    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
5399    /// carried in from outside the range) and exits with the range's final residual materialized
5400    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
5401    /// instead of one.
5402    ///
5403    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
5404    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
5405    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
5406    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
5407    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
5408    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
5409    /// code — there is no "split version" of the verify math.
5410    ///
5411    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
5412    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
5413    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
5414    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
5415    #[allow(clippy::too_many_arguments)]
5416    fn verify_layers(
5417        &self,
5418        e: &Engine,
5419        mut x: CudaSlice<f32>,
5420        lo: usize,
5421        hi: usize,
5422        pos_d: &CudaSlice<i32>,
5423        pos0: usize,
5424        t: usize,
5425        cache: &mut Cache,
5426        mut ckpt: Option<&mut VerifyCkpt>,
5427        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5428        graphs: Option<&mut DsparkVerifyGraphs>,
5429    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5430        if self.cfg.step35.is_some() {
5431            if stream.is_some() {
5432                return Err(
5433                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
5434                            cannot express the SWA offset KV view)"
5435                        .into(),
5436                );
5437            }
5438            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
5439        }
5440        if self.qwen35_serving_class() {
5441            return self.qwen35_verify_batch_layers(
5442                e,
5443                x,
5444                lo,
5445                hi,
5446                pos0,
5447                t,
5448                cache,
5449                ckpt.take(),
5450                stream,
5451                graphs,
5452            );
5453        }
5454        let n_embd = self.cfg.n_embd as usize;
5455        let eps = self.cfg.rms_eps;
5456        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
5457        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
5458        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
5459        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
5460        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
5461        // residual the next layer needs) as its `res` output. Falls back to the separate add
5462        // when the next layer is off the fused-q8 path.
5463        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
5464        for il in lo..hi {
5465            let layer = &self.layers[il];
5466            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
5467            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
5468            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
5469            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
5470            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
5471            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
5472            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
5473            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
5474            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
5475            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
5476            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
5477            // projections only; Linear mixer: the batched arm — the per-column fallback needs
5478            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
5479            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
5480            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
5481            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
5482            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
5483            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
5484            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
5485            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
5486            let lin_q8_only = match &layer.mixer {
5487                Mixer::Linear(la) => {
5488                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
5489                }
5490                Mixer::Full(_) if self.cfg.step35.is_some() => false,
5491                _ => true,
5492            };
5493            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
5494            // a non-fused layer still performs the residual add.
5495            let taken = pending.take();
5496            let (h, h_q8) = if norm_fused && lin_q8_only {
5497                let pair = match taken {
5498                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
5499                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
5500                    Some((x1p, f1p)) => {
5501                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
5502                        let p = e.add_rms_norm_q8_1(
5503                            &x1p,
5504                            &f1p,
5505                            layer.attn_norm.float_data(),
5506                            &mut x2,
5507                            n_embd,
5508                            t,
5509                            eps,
5510                        )?;
5511                        x = x2;
5512                        p
5513                    }
5514                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
5515                };
5516                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
5517            } else {
5518                if let Some((x1p, f1p)) = taken {
5519                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5520                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5521                    x = x2;
5522                }
5523                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
5524                if norm_fused {
5525                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5526                } else {
5527                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
5528                }
5529                (h, None)
5530            };
5531            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
5532
5533            let mixed = match &layer.mixer {
5534                Mixer::Full(fa) => self.full_attn_verify(
5535                    e,
5536                    fa,
5537                    &h,
5538                    h_q8_ref,
5539                    pos_d,
5540                    t,
5541                    cache,
5542                    il,
5543                    stream.map(|(_, c)| c),
5544                )?,
5545                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
5546                Mixer::Linear(la) => {
5547                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
5548                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
5549                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
5550                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
5551                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
5552                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
5553                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
5554                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
5555                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
5556                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
5557                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
5558                    if (t >= 3 || (t == 2 && spec_m2()))
5559                        && mixer_fast
5560                        && e.uses_q8_1_fast(&la.ssm_out)
5561                    {
5562                        let want = ckpt.is_some();
5563                        let (out, stash) =
5564                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
5565                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
5566                            ck.gdn[il] = Some(st);
5567                        }
5568                        out
5569                    } else {
5570                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
5571                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5572                            if ckpt.is_some() && t >= 2 {
5573                                Some(Vec::with_capacity(t - 1))
5574                            } else {
5575                                None
5576                            };
5577                        for col in 0..t {
5578                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
5579                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
5580                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
5581                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
5582                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
5583                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
5584                            // (pure dtod — cannot change any computed value). Last column skipped:
5585                            // rebuild targets are j <= t-1 columns.
5586                            if let Some(cs) = col_states.as_mut() {
5587                                if col + 1 < t {
5588                                    let rl = cache.recur[il].as_ref().unwrap();
5589                                    cs.push((
5590                                        e.clone_dtod(&rl.conv_state)?,
5591                                        e.clone_dtod(&rl.ssm_state)?,
5592                                    ));
5593                                }
5594                            }
5595                        }
5596                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
5597                            // ReplaySSM-assessment instrumentation (2026-07-30): the
5598                            // per-column clones are the only true state snapshots left in
5599                            // the verify (the batched path stashes INPUTS and replays).
5600                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
5601                                static ONCE: std::sync::Once = std::sync::Once::new();
5602                                let bytes: usize =
5603                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
5604                                ONCE.call_once(|| eprintln!(
5605                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
5606                                    cs.len(), bytes as f64 / 1e6));
5607                            }
5608                            ck.cols[il] = Some(cs);
5609                        }
5610                        out
5611                    }
5612                }
5613            };
5614
5615            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
5616            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
5617            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
5618            let ffn_fuse = match &layer.ffn {
5619                crate::hybrid::Ffn::Dense {
5620                    ffn_gate, ffn_up, ..
5621                } => {
5622                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
5623                        && e.uses_q8_1_fast(ffn_gate)
5624                        && e.uses_q8_1_fast(ffn_up)
5625                }
5626                crate::hybrid::Ffn::Moe(_) => false,
5627            };
5628            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
5629            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
5630            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
5631            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
5632            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
5633            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
5634            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
5635            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
5636            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
5637            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
5638            // mirror decode's dispatch or spec self-consistency fails.
5639            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
5640            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
5641            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
5642            let mut z = e.zeros(0)?; // replaced below on the unfused arms
5643            let z_q8 = if fuse_q8 {
5644                Some(e.add_rms_norm_q8_1(
5645                    &x,
5646                    &mixed,
5647                    layer.post_attn_norm.float_data(),
5648                    &mut x1,
5649                    n_embd,
5650                    t,
5651                    eps,
5652                )?)
5653            } else {
5654                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
5655                if ffn_fuse {
5656                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
5657                    e.rms_norm_decode(
5658                        &x1,
5659                        layer.post_attn_norm.float_data(),
5660                        &mut zf,
5661                        n_embd,
5662                        t,
5663                        eps,
5664                    )?;
5665                } else {
5666                    e.add_rms_norm(
5667                        &x,
5668                        &mixed,
5669                        layer.post_attn_norm.float_data(),
5670                        &mut x1,
5671                        &mut zf,
5672                        n_embd,
5673                        t,
5674                        eps,
5675                    )?;
5676                }
5677                z = zf;
5678                None
5679            };
5680            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
5681            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
5682            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
5683            let ffn_out = match &layer.ffn {
5684                crate::hybrid::Ffn::Dense {
5685                    ffn_gate,
5686                    ffn_up,
5687                    ffn_down,
5688                } => {
5689                    let n_ff = ffn_gate.out_features();
5690                    if let Some((zq, zd)) = z_q8.as_ref() {
5691                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
5692                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
5693                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
5694                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
5695                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
5696                        // structure at nrows=t.
5697                        let pair =
5698                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
5699                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
5700                                None => None,
5701                            };
5702                        let (gate, gs, up, us) = match pair {
5703                            Some(x4) => x4,
5704                            None => (
5705                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
5706                                1.0, // scale already applied inside _pre
5707                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
5708                                1.0,
5709                            ),
5710                        };
5711                        if e.uses_q8_1_fast(ffn_down) {
5712                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
5713                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
5714                        } else {
5715                            let mut act = vbuf(e, t * n_ff)?;
5716                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
5717                            e.matmul_decode_exact(ffn_down, &act, t)?
5718                        }
5719                    } else {
5720                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
5721                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
5722                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
5723                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
5724                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
5725                        let (gate, up) =
5726                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
5727                                Some(pair) => pair,
5728                                None => (
5729                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
5730                                    e.matmul_decode_exact(ffn_up, &z, t)?,
5731                                ),
5732                            };
5733                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
5734                        Self::ffn_act_lim(
5735                            e,
5736                            &self.cfg,
5737                            &gate,
5738                            &up,
5739                            1.0,
5740                            1.0,
5741                            dense_lim,
5742                            &mut act,
5743                            t * n_ff,
5744                        )?;
5745                        e.matmul_decode_exact(ffn_down, &act, t)?
5746                    }
5747                }
5748                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
5749            };
5750            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
5751            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
5752            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
5753            pending = Some((x1, ffn_out));
5754        }
5755        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
5756        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
5757        if let Some((x1p, f1p)) = pending.take() {
5758            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
5759            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
5760            x = x2;
5761        }
5762        Ok(x)
5763    }
5764    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
5765    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
5766    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
5767    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
5768    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
5769    /// ssm state exactly like T sequential decode steps.
5770    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
5771    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
5772    #[allow(clippy::too_many_arguments)]
5773    fn linear_attn_verify_t(
5774        &self,
5775        e: &Engine,
5776        la: &LinearAttnLayer,
5777        h: &CudaSlice<f32>,
5778        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
5779        t: usize,
5780        cache: &mut Cache,
5781        il: usize,
5782        want_stash: bool,
5783    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
5784        let cfg = &self.cfg;
5785        let ssm = cfg.ssm.as_ref().unwrap();
5786        let d_state = ssm.state_size as usize;
5787        let num_k = ssm.group_count as usize;
5788        let num_v = ssm.time_step_rank as usize;
5789        let d_conv = ssm.conv_kernel as usize;
5790        let key_dim = d_state * num_k;
5791        let conv_dim = key_dim * 2 + d_state * num_v;
5792        let eps = cfg.rms_eps;
5793        let scale = 1.0 / (d_state as f32).sqrt();
5794
5795        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
5796        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
5797        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
5798        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
5799        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
5800        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
5801        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
5802        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
5803        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
5804        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
5805        // Bit-identical per (tensor,token,row) — see spec_fused_t().
5806        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
5807        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
5808        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
5809        // and feeds every projection; the caller guaranteed all four input projections are
5810        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
5811        let h_q8_t = if h_q8.is_none()
5812            && spec_fused_t()
5813            && (2..=4).contains(&t)
5814            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
5815                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
5816        {
5817            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
5818        } else {
5819            None
5820        };
5821        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
5822        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
5823            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
5824        let (qkv_mixed, z) = {
5825            let mut fused = None;
5826            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
5827                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5828                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
5829            } else if let Some((hq, hd)) = hq8_any {
5830                if spec_fused_t() && (2..=4).contains(&t) {
5831                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
5832                }
5833            }
5834            match (fused, hq8_any) {
5835                (Some(pair), _) => pair,
5836                (None, Some((hq, hd))) if h_q8.is_some() => (
5837                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
5838                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
5839                ),
5840                (None, _) => (
5841                    e.matmul_decode_exact(&la.wqkv, h, t)?,
5842                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
5843                ),
5844            }
5845        };
5846        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
5847        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
5848        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
5849        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
5850        let (beta_raw, alpha) = if t == 1 {
5851            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
5852            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
5853                Some(((mut b, bs), (mut a, as_))) => {
5854                    if bs != 1.0 {
5855                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
5856                    }
5857                    if as_ != 1.0 {
5858                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
5859                    }
5860                    (b, a)
5861                }
5862                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
5863                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
5864                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
5865                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
5866                    Some((b, a)) => (b, a),
5867                    None => (
5868                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
5869                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
5870                    ),
5871                },
5872            }
5873        } else {
5874            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
5875            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
5876            let mut nvfp4_fused = None;
5877            let mut q8_fused = None;
5878            if let Some((hq, hd)) = hq8_any {
5879                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
5880                    nvfp4_fused =
5881                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5882                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
5883                        static ONCE: std::sync::Once = std::sync::Once::new();
5884                        ONCE.call_once(|| {
5885                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
5886                        });
5887                    }
5888                }
5889                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
5890                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
5891                }
5892            }
5893            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
5894                if bs != 1.0 {
5895                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
5896                }
5897                if as_ != 1.0 {
5898                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
5899                }
5900                (b, a)
5901            } else if let Some(pair) = q8_fused {
5902                pair
5903            } else {
5904                match hq8_any {
5905                    Some((hq, hd)) if h_q8.is_some() => (
5906                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
5907                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
5908                    ),
5909                    _ => (
5910                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
5911                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
5912                    ),
5913                }
5914            }
5915        };
5916
5917        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
5918        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
5919        let rl = cache.recur[il].as_mut().unwrap();
5920        let mut conv_out = e.uninit(conv_dim * t)?;
5921        e.ssm_conv1d_tm_state(
5922            &qkv_mixed,
5923            &mut rl.conv_state,
5924            la.ssm_conv1d.float_data(),
5925            &mut conv_out,
5926            conv_dim,
5927            t,
5928            d_conv,
5929        )?;
5930
5931        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
5932        let mut q_g = e.uninit(d_state * num_v * t)?;
5933        let mut k_g = e.uninit(d_state * num_v * t)?;
5934        let mut v_g = e.uninit(d_state * num_v * t)?;
5935        e.qkv_to_gdn_repack(
5936            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
5937        )?;
5938        let mut q_l2 = e.uninit(d_state * num_v * t)?;
5939        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
5940        let mut k_l2 = e.uninit(d_state * num_v * t)?;
5941        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
5942        let mut beta = e.uninit(t * num_v)?;
5943        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
5944        let mut g_log = e.uninit(t * num_v)?;
5945        e.gdn_glog(
5946            &alpha,
5947            la.ssm_dt.float_data(),
5948            la.ssm_a.float_data(),
5949            &mut g_log,
5950            num_v,
5951            t,
5952        )?;
5953
5954        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
5955        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
5956        let mut o = e.uninit(d_state * num_v * t)?;
5957        {
5958            let crate::cache::RecurLayer {
5959                ssm_state,
5960                ssm_state_alt,
5961                ..
5962            } = rl;
5963            e.gdn_scan_s128(
5964                &q_l2,
5965                &k_l2,
5966                &v_g,
5967                &g_log,
5968                &beta,
5969                ssm_state,
5970                ssm_state_alt,
5971                &mut o,
5972                num_v,
5973                t,
5974                scale,
5975            )?;
5976        }
5977        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
5978
5979        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
5980        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
5981        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
5982        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
5983        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
5984        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
5985        let out = if e.uses_q8_1_fast(&la.ssm_out) {
5986            let (gq, gd) =
5987                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
5988            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
5989        } else {
5990            let mut gn = e.uninit(d_state * num_v * t)?;
5991            e.gated_rmsnorm(
5992                &o,
5993                la.ssm_norm.float_data(),
5994                &z,
5995                &mut gn,
5996                d_state,
5997                num_v * t,
5998                eps,
5999            )?;
6000            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
6001            // would fall to dp4a with a different FP reduction order — same class of bug as
6002            // the input projs).
6003            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
6004        };
6005        let stash = if want_stash {
6006            Some(GdnStash {
6007                qkv_mixed,
6008                q_l2,
6009                k_l2,
6010                v_g,
6011                g_log,
6012                beta,
6013            })
6014        } else {
6015            None
6016        };
6017        Ok((out, stash))
6018    }
6019
6020    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
6021    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
6022    /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
6023    ///   are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
6024    ///   verify-probe gates), so keeping them == replaying them.
6025    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
6026    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
6027    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
6028    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
6029    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
6030    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
6031    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
6032    fn commit_verified_prefix(
6033        &self,
6034        e: &Engine,
6035        cache: &mut Cache,
6036        snap: &crate::cache::CacheSnapshot,
6037        ckpt: &VerifyCkpt,
6038        j: usize,
6039        kv_lens_done: bool,
6040        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
6041    ) -> Result<(), Box<dyn std::error::Error>> {
6042        let cfg = &self.cfg;
6043        let ssm = cfg.ssm.as_ref().unwrap();
6044        let d_state = ssm.state_size as usize;
6045        let num_k = ssm.group_count as usize;
6046        let num_v = ssm.time_step_rank as usize;
6047        let d_conv = ssm.conv_kernel as usize;
6048        let conv_dim = d_state * num_k * 2 + d_state * num_v;
6049        let scale = 1.0 / (d_state as f32).sqrt();
6050        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
6051        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
6052        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
6053        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
6054        // buffers and stream order are identical to the per-layer memcpy sequence; the
6055        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
6056        let mut batched_cols = false;
6057        if state_copy_batch_on() && dev_j.is_none() {
6058            use cudarc::driver::DevicePtr;
6059            let s = &e.gpu.stream();
6060            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
6061            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
6062            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
6063            let mut uniform = true;
6064            for il in 0..self.layers.len() {
6065                let Some(rl) = cache.recur[il].as_ref() else {
6066                    continue;
6067                };
6068                if ckpt.gdn[il].is_some() {
6069                    continue; // kernel-rebuild arm restores below, per layer
6070                }
6071                let Some(cols) = &ckpt.cols[il] else {
6072                    continue; // missing-ckpt error surfaces in the main loop
6073                };
6074                let (c, st) = &cols[j - 1];
6075                if conv_pairs.is_empty() {
6076                    conv_words = c.len();
6077                    ssm_words = st.len();
6078                } else if c.len() != conv_words || st.len() != ssm_words {
6079                    uniform = false;
6080                    break;
6081                }
6082                let (pc, _g0) = c.device_ptr(s);
6083                let (dc, _g1) = rl.conv_state.device_ptr(s);
6084                let (ps, _g2) = st.device_ptr(s);
6085                let (ds, _g3) = rl.ssm_state.device_ptr(s);
6086                conv_pairs.push((pc as u64, dc as u64));
6087                ssm_pairs.push((ps as u64, ds as u64));
6088            }
6089            if uniform && !conv_pairs.is_empty() {
6090                let n = conv_pairs.len();
6091                let mut t = vec![0u64; 2 * n];
6092                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
6093                    t[k] = src;
6094                    t[n + k] = dst;
6095                }
6096                let conv_t = e.htod_u64(&t)?;
6097                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
6098                    t[k] = src;
6099                    t[n + k] = dst;
6100                }
6101                let ssm_t = e.htod_u64(&t)?;
6102                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
6103                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
6104                batched_cols = true;
6105            }
6106        }
6107        for il in 0..self.layers.len() {
6108            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6109                kvl.len = saved + j;
6110                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
6111                if !kv_lens_done {
6112                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6113                }
6114            }
6115            if let Some(rl) = cache.recur[il].as_mut() {
6116                if let Some(st) = &ckpt.gdn[il] {
6117                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
6118                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
6119                    if let Some((acc, base, t_v)) = dev_j {
6120                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
6121                        e.ssm_conv_ring_rebuild_dc(
6122                            &st.qkv_mixed,
6123                            ring_old,
6124                            &mut rl.conv_state,
6125                            conv_dim,
6126                            acc,
6127                            base,
6128                            t_v,
6129                            d_conv,
6130                        )?;
6131                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
6132                        e.gdn_scan_s128_dc(
6133                            &st.q_l2,
6134                            &st.k_l2,
6135                            &st.v_g,
6136                            &st.g_log,
6137                            &st.beta,
6138                            state_in,
6139                            &mut rl.ssm_state,
6140                            &mut o,
6141                            num_v,
6142                            acc,
6143                            base,
6144                            t_v,
6145                            scale,
6146                        )?;
6147                    } else {
6148                        e.ssm_conv_ring_rebuild(
6149                            &st.qkv_mixed,
6150                            ring_old,
6151                            &mut rl.conv_state,
6152                            conv_dim,
6153                            j,
6154                            d_conv,
6155                        )?;
6156                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
6157                        e.gdn_scan_s128(
6158                            &st.q_l2,
6159                            &st.k_l2,
6160                            &st.v_g,
6161                            &st.g_log,
6162                            &st.beta,
6163                            state_in,
6164                            &mut rl.ssm_state,
6165                            &mut o,
6166                            num_v,
6167                            j,
6168                            scale,
6169                        )?;
6170                    }
6171                } else if let Some(cols) = &ckpt.cols[il] {
6172                    if !batched_cols {
6173                        let (c, s) = &cols[j - 1];
6174                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
6175                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
6176                    }
6177                } else {
6178                    return Err(
6179                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
6180                    );
6181                }
6182            }
6183        }
6184        cache.pos = snap.pos + j;
6185        Ok(())
6186    }
6187
6188    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
6189    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
6190    fn commit_verified_prefix_stream(
6191        &self,
6192        e: &Engine,
6193        cache: &mut Cache,
6194        snap: &crate::cache::CacheSnapshot,
6195        ckpt: &VerifyCkpt,
6196        acc: &CudaSlice<u32>,
6197        base: usize,
6198        t_v: usize,
6199    ) -> Result<(), Box<dyn std::error::Error>> {
6200        let cfg = &self.cfg;
6201        let ssm = cfg.ssm.as_ref().unwrap();
6202        let d_state = ssm.state_size as usize;
6203        let num_k = ssm.group_count as usize;
6204        let num_v = ssm.time_step_rank as usize;
6205        let d_conv = ssm.conv_kernel as usize;
6206        let conv_dim = d_state * num_k * 2 + d_state * num_v;
6207        let scale = 1.0 / (d_state as f32).sqrt();
6208        for il in 0..self.layers.len() {
6209            if let Some(rl) = cache.recur[il].as_mut() {
6210                let st = ckpt.gdn[il]
6211                    .as_ref()
6212                    .ok_or("stream restore: batched-linear stash missing")?;
6213                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
6214                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
6215                e.ssm_conv_ring_rebuild_dc(
6216                    &st.qkv_mixed,
6217                    ring_old,
6218                    &mut rl.conv_state,
6219                    conv_dim,
6220                    acc,
6221                    base,
6222                    t_v,
6223                    d_conv,
6224                )?;
6225                let mut o = e.uninit(d_state * num_v * t_v)?;
6226                e.gdn_scan_s128_dc(
6227                    &st.q_l2,
6228                    &st.k_l2,
6229                    &st.v_g,
6230                    &st.g_log,
6231                    &st.beta,
6232                    state_in,
6233                    &mut rl.ssm_state,
6234                    &mut o,
6235                    num_v,
6236                    acc,
6237                    base,
6238                    t_v,
6239                    scale,
6240                )?;
6241            }
6242        }
6243        Ok(())
6244    }
6245
6246    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
6247    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
6248    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
6249    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
6250    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
6251    pub fn decode_step_t_aux2(
6252        &self,
6253        e: &Engine,
6254        tokens: &[u32],
6255        pos0: usize,
6256        cache: &mut Cache,
6257        aux_layers: &[usize],
6258        pred_col: Option<usize>,
6259    ) -> Result<
6260        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
6261        Box<dyn std::error::Error>,
6262    > {
6263        let cfg = &self.cfg;
6264        let n_embd = cfg.n_embd as usize;
6265        let eps = cfg.rms_eps;
6266        let t = tokens.len();
6267        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6268        let pos_d = e.htod_i32(&pos_vec)?;
6269        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
6270        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
6271        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
6272        let want_pred = pred_col.is_some();
6273
6274        for (il, layer) in self.layers.iter().enumerate() {
6275            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
6276            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
6277            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
6278            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
6279            if norm_fused {
6280                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6281            } else {
6282                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6283            }
6284            let mixed = match &layer.mixer {
6285                Mixer::Full(fa) => {
6286                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
6287                }
6288                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6289                Mixer::Linear(la) => {
6290                    let mut out = e.zeros(t * n_embd)?;
6291                    for col in 0..t {
6292                        let mut h_col = e.zeros(n_embd)?;
6293                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
6294                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
6295                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
6296                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
6297                    }
6298                    out
6299                }
6300            };
6301            let ffn_fuse = match &layer.ffn {
6302                crate::hybrid::Ffn::Dense {
6303                    ffn_gate, ffn_up, ..
6304                } => {
6305                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
6306                        && e.uses_q8_1_fast(ffn_gate)
6307                        && e.uses_q8_1_fast(ffn_up)
6308                }
6309                crate::hybrid::Ffn::Moe(_) => false,
6310            };
6311            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
6312            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
6313            if ffn_fuse {
6314                e.add(&x, &mixed, &mut x1, t * n_embd)?;
6315                e.rms_norm_decode(
6316                    &x1,
6317                    layer.post_attn_norm.float_data(),
6318                    &mut z,
6319                    n_embd,
6320                    t,
6321                    eps,
6322                )?;
6323            } else {
6324                e.add_rms_norm(
6325                    &x,
6326                    &mixed,
6327                    layer.post_attn_norm.float_data(),
6328                    &mut x1,
6329                    &mut z,
6330                    n_embd,
6331                    t,
6332                    eps,
6333                )?;
6334            }
6335            let ffn_out = match &layer.ffn {
6336                crate::hybrid::Ffn::Dense {
6337                    ffn_gate,
6338                    ffn_up,
6339                    ffn_down,
6340                } => {
6341                    let n_ff = ffn_gate.out_features();
6342                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
6343                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
6344                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
6345                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
6346                    Self::ffn_act_lim(
6347                        e,
6348                        &self.cfg,
6349                        &gate,
6350                        &up,
6351                        1.0,
6352                        1.0,
6353                        self.cfg.clamp_shexp_at(il as u32),
6354                        &mut act,
6355                        t * n_ff,
6356                    )?;
6357                    e.matmul_decode_exact(ffn_down, &act, t)?
6358                }
6359                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
6360            };
6361            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6362            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6363            if aux_layers.contains(&il) {
6364                let mut a = e.zeros(n_embd)?;
6365                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
6366                aux_last.push(a);
6367                if let Some(pc) = pred_col {
6368                    let mut ap = e.zeros(n_embd)?;
6369                    e.copy_view_into(
6370                        &mut ap,
6371                        0,
6372                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
6373                        n_embd,
6374                    )?;
6375                    aux_pred.push(ap);
6376                }
6377            }
6378            x = x2;
6379        }
6380        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
6381        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6382        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
6383        let host = e.dtoh(&logits)?;
6384        cache.pos += t;
6385        Ok((
6386            host,
6387            aux_last,
6388            if want_pred { Some(aux_pred) } else { None },
6389        ))
6390    }
6391
6392    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
6393    /// `step35_decode_attn`.
6394    ///
6395    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
6396    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
6397    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
6398    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
6399    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
6400    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
6401    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
6402    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
6403    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
6404    /// position of each query row. A batched twin would have to reproduce all of that AND the
6405    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
6406    /// take one `base_len`, not a per-row offset).
6407    ///
6408    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
6409    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
6410    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
6411    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
6412    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
6413    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
6414    /// step35 twin is a perf lane's job and must be gated against this arm.
6415    ///
6416    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
6417    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
6418    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
6419    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
6420    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
6421    #[allow(clippy::too_many_arguments)]
6422    fn step35_verify(
6423        &self,
6424        e: &Engine,
6425        fa: &FullAttnLayer,
6426        h: &CudaSlice<f32>,
6427        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
6428        t: usize,
6429        cache: &mut Cache,
6430        il: usize,
6431    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6432        let n_embd = self.cfg.n_embd as usize;
6433        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
6434        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
6435        // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
6436        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
6437        // cannot regress it into silently reading an empty buffer.
6438        assert_eq!(
6439            h.len(),
6440            t * n_embd,
6441            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
6442             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
6443            h_q8.is_some()
6444        );
6445        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
6446        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
6447        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
6448        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
6449        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
6450        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
6451        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
6452        for r in 0..t {
6453            // Absolute position of this query row. `cache.pos` is the committed length at round
6454            // start and every row before r has already been appended by this loop, so the r-th
6455            // verify token sits at cache.pos + r — the same position eager decode would give it.
6456            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
6457            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
6458            e.copy_view_into(
6459                &mut h_row,
6460                0,
6461                &h.slice(r * n_embd..(r + 1) * n_embd),
6462                n_embd,
6463            )?;
6464            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
6465            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
6466            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
6467            debug_assert_eq!(
6468                o.len(),
6469                n_embd,
6470                "step35_decode_attn returns post-wo [n_embd]"
6471            );
6472            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
6473        }
6474        Ok(out)
6475    }
6476
6477    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
6478    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
6479    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
6480    #[allow(clippy::too_many_arguments)]
6481    fn full_attn_verify(
6482        &self,
6483        e: &Engine,
6484        fa: &FullAttnLayer,
6485        h: &CudaSlice<f32>,
6486        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
6487        pos_d: &CudaSlice<i32>,
6488        t: usize,
6489        cache: &mut Cache,
6490        il: usize,
6491        stream_ctr: Option<&CudaSlice<i32>>,
6492    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6493        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
6494        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
6495        // its own arm. A verify that silently computes different attention than decode defeats the
6496        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
6497        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
6498        // shape and not laziness.
6499        if self.cfg.step35.is_some() {
6500            if stream_ctr.is_some() {
6501                return Err(
6502                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
6503                            cannot express the SWA offset KV view; same root cause as the dc \
6504                            decode refusal) — run spec without the stream arm"
6505                        .into(),
6506                );
6507            }
6508            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
6509        }
6510        let cfg = &self.cfg;
6511        let geometry = cfg.full_attention_geometry_at(il as u32);
6512        let n_head = geometry.n_head as usize;
6513        let n_head_kv = geometry.n_head_kv as usize;
6514        let head_dim = geometry.head_dim_k as usize;
6515        let eps = cfg.rms_eps;
6516        let scale = geometry.attention_scale();
6517        let n_embd = cfg.n_embd as usize;
6518
6519        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
6520        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
6521        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
6522        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
6523        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
6524        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
6525        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
6526        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
6527        let (qf, mut k, v) = {
6528            let mut fused = None;
6529            let qkv_fast =
6530                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
6531            if t == 1 && qkv_fast {
6532                let (hq_o, hd_o);
6533                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
6534                    Some(p) => p,
6535                    None => {
6536                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
6537                        (&hq_o, &hd_o)
6538                    }
6539                };
6540                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
6541            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
6542                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
6543                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
6544                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
6545                let (hq_o, hd_o);
6546                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
6547                    Some(p) => p,
6548                    None => {
6549                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
6550                        (&hq_o, &hd_o)
6551                    }
6552                };
6553                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
6554            }
6555            match (fused, h_q8) {
6556                (Some(triple), _) => triple,
6557                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
6558                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
6559                (None, Some((hq, hd))) if qkv_fast => (
6560                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
6561                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
6562                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
6563                ),
6564                (None, _) => (
6565                    e.matmul_decode_exact(&fa.wq, h, t)?,
6566                    e.matmul_decode_exact(&fa.wk, h, t)?,
6567                    e.matmul_decode_exact(&fa.wv, h, t)?,
6568                ),
6569            }
6570        };
6571        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
6572        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6573        let (mut q, gate) = if gated {
6574            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
6575            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
6576            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
6577            (q, Some(gate))
6578        } else {
6579            (qf, None)
6580        };
6581
6582        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
6583        e.rms_norm(
6584            &q,
6585            fa.q_norm.float_data(),
6586            &mut qn,
6587            head_dim,
6588            n_head * t,
6589            eps,
6590        )?;
6591        q = qn;
6592        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
6593        e.rms_norm(
6594            &k,
6595            fa.k_norm.float_data(),
6596            &mut kn,
6597            head_dim,
6598            n_head_kv * t,
6599            eps,
6600        )?;
6601        k = kn;
6602        let rope_dims = geometry.n_rot as usize;
6603        e.rope_neox(
6604            &mut q,
6605            pos_d,
6606            head_dim,
6607            rope_dims,
6608            n_head,
6609            t,
6610            geometry.rope_base,
6611            1.0,
6612        )?;
6613        e.rope_neox(
6614            &mut k,
6615            pos_d,
6616            head_dim,
6617            rope_dims,
6618            n_head_kv,
6619            t,
6620            geometry.rope_base,
6621            1.0,
6622        )?;
6623
6624        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
6625        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
6626        let kvl = cache.kv[il].as_mut().unwrap();
6627        let (kv_dim_k, kv_dim_v, ktb, vtb) =
6628            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
6629        if let Some(ctr) = stream_ctr {
6630            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
6631            // math on a (block, token) grid, documented byte-identical); host len is a stale
6632            // LOWER BOUND under pre-issue (drain reconciles it).
6633            e.append_kv_quantized_rows_dc(
6634                &k,
6635                &v,
6636                &mut kvl.k,
6637                &mut kvl.v,
6638                ctr,
6639                t,
6640                kv_dim_k,
6641                kv_dim_v,
6642                ktb,
6643                vtb,
6644                crate::Engine::kv_fp8_on(),
6645            )?;
6646        } else {
6647            for i in 0..t {
6648                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
6649                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
6650                e.append_kv_quantized_view(
6651                    &k_row,
6652                    &v_row,
6653                    &mut kvl.k,
6654                    &mut kvl.v,
6655                    kvl.len + i,
6656                    kv_dim_k,
6657                    kv_dim_v,
6658                    ktb,
6659                    vtb,
6660                    crate::Engine::kv_fp8_on(),
6661                )?;
6662            }
6663            kvl.len += t;
6664        }
6665
6666        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
6667        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
6668        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
6669        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
6670        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
6671        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
6672        // keys. The verify appends all T tokens first but bounds the key range per row.
6673        //
6674        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
6675        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
6676        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
6677        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
6678        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
6679        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
6680        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
6681        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
6682        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
6683        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
6684        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
6685        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
6686        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
6687        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
6688        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
6689        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
6690        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
6691        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
6692        if let Some(ctr) = stream_ctr {
6693            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
6694            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
6695            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
6696            let upper = kvl.len + t + 64;
6697            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
6698            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
6699            e.fa_decode_rows_dc(
6700                &q,
6701                &k_view,
6702                &v_view,
6703                &mut attn,
6704                head_dim,
6705                n_head,
6706                n_head_kv,
6707                ctr,
6708                upper.min(cache.max_ctx),
6709                t,
6710                scale,
6711                ktb,
6712                vtb,
6713                0,
6714                false,
6715            )?;
6716        } else if spec_lean() && t == 1 {
6717            let t_kv = base_len + 1;
6718            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
6719            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
6720            e.fa_decode_kvmod(
6721                &q,
6722                &k_view,
6723                &v_view,
6724                &mut attn,
6725                head_dim,
6726                n_head,
6727                n_head_kv,
6728                t_kv,
6729                scale,
6730                ktb,
6731                vtb,
6732                crate::Engine::kv_fp8_on(),
6733            )?;
6734        } else if e.fa_rows_eligible(base_len, head_dim) {
6735            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
6736            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
6737            e.fa_decode_rows(
6738                &q,
6739                &k_view,
6740                &v_view,
6741                &mut attn,
6742                head_dim,
6743                n_head,
6744                n_head_kv,
6745                base_len,
6746                t,
6747                scale,
6748                ktb,
6749                vtb,
6750                None,
6751                false,
6752                crate::Engine::kv_fp8_on(),
6753                None,
6754            )?;
6755        } else {
6756            for r in 0..t {
6757                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
6758                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
6759                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
6760                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
6761                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
6762                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
6763                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
6764                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
6765                e.fa_decode_kvmod(
6766                    &q_row,
6767                    &k_view_r,
6768                    &v_view_r,
6769                    &mut attn_row,
6770                    head_dim,
6771                    n_head,
6772                    n_head_kv,
6773                    t_kv_r,
6774                    scale,
6775                    ktb,
6776                    vtb,
6777                    crate::Engine::kv_fp8_on(),
6778                )?;
6779                e.copy_into(
6780                    &mut attn,
6781                    r * n_head * head_dim,
6782                    &attn_row,
6783                    n_head * head_dim,
6784                )?;
6785            }
6786        }
6787
6788        let attn_g = match &gate {
6789            Some(gate) => {
6790                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
6791                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
6792                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
6793                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
6794                ag
6795            }
6796            None => attn,
6797        };
6798        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
6799        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
6800        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
6801    }
6802
6803    /// Context-linear bytes for a plain serving session's trunk cache.
6804    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
6805        crate::cache::cache_bytes_per_token(&self.cfg)
6806    }
6807
6808    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
6809    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
6810        (
6811            self.plain_session_kv_bytes_per_token(),
6812            crate::cache::cache_ring_bytes_per_token(&self.cfg),
6813            crate::cache::cache_ring_row_cap(&self.cfg),
6814        )
6815    }
6816
6817    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
6818    /// scratch. With no MTP head this equals the plain coefficient.
6819    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
6820        let scratch = self
6821            .mtp
6822            .as_ref()
6823            .map(|mtp| {
6824                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6825                k + v
6826            })
6827            .unwrap_or(0);
6828        self.plain_session_kv_bytes_per_token()
6829            .saturating_add(scratch)
6830    }
6831
6832    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
6833    /// capped by the same SWA ring rows as the trunk.
6834    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
6835        let total = self.spec_session_kv_bytes_per_token();
6836        let (_, mut ring, rows) = self.plain_session_kv_shape();
6837        if rows > 0 {
6838            ring = ring.saturating_add(
6839                self.mtp
6840                    .as_ref()
6841                    .map(|mtp| {
6842                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
6843                        k + v
6844                    })
6845                    .unwrap_or(0),
6846            );
6847        }
6848        (total, ring, rows)
6849    }
6850
6851    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
6852    /// the NextN head to draft K tokens then verifies them in one batched target forward.
6853    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
6854    /// acceptance rate. `k` = draft length per round.
6855    ///
6856    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
6857    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
6858    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
6859    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
6860    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
6861    /// captured graph references is event-free; the spec loop is strictly single-stream.
6862    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
6863    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
6864    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
6865    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
6866    /// generate_spec_inner2.
6867    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
6868    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
6869    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
6870    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
6871    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
6872    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
6873    pub fn new_session(
6874        &self,
6875        e: &Engine,
6876        max_ctx: usize,
6877    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
6878        Ok(SpecSession {
6879            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
6880            // is the SERVING spec-session path, and with the ppN door open across two cards a
6881            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
6882            // round — the wrong-card class already fixed on the two batched serving paths
6883            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
6884            // branch, same allocations), so single-device behavior is byte-unchanged.
6885            cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
6886            scratch: MtpScratch::new(
6887                e,
6888                &self.cfg,
6889                max_ctx,
6890                self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
6891            )?,
6892            committed: Vec::new(),
6893            last_h: None,
6894            next_pred: None,
6895            sctr: 0,
6896            uctr: 0,
6897            draft_ctx: None,
6898            pending_tok: None,
6899            turn_ckpt: None,
6900            telem: SpecTelemetryCounters::default(),
6901            capture_at: None,
6902            boundary_capture: None,
6903        })
6904    }
6905
6906    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
6907    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
6908    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
6909    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
6910    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
6911    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
6912    /// worker always receives a fully-warm continuation session (committed = whole
6913    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
6914    /// boundary logits on the empty-suffix shape).
6915    ///
6916    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
6917    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
6918    /// request, and plain feeds a carried suffix via eager `decode_step` below
6919    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
6920    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
6921    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
6922    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
6923    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
6924    /// burst prime.
6925    ///
6926    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
6927    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
6928    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
6929    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
6930    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
6931    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
6932    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
6933    /// cold session draws from the identical row at counter 0 and then runs its rounds from
6934    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
6935    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
6936    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
6937    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
6938    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
6939    ///
6940    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
6941    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
6942    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
6943    /// and are never routed here.
6944    ///
6945    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
6946    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
6947    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
6948    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
6949    /// entry stays published for the next request.
6950    #[allow(clippy::too_many_arguments)]
6951    pub fn spec_session_from_restored(
6952        &self,
6953        e: &Engine,
6954        mut cache: Cache,
6955        prefix: Vec<u32>,
6956        suffix: &[u32],
6957        draft_k: &CudaSlice<u8>,
6958        draft_v: &CudaSlice<u8>,
6959        draft_k_tok_bytes: usize,
6960        draft_v_tok_bytes: usize,
6961        draft_len: usize,
6962        last_h: &[f32],
6963        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
6964        // when a suffix follows — the feed's own logits are the boundary then.
6965        boundary_logits: &[f32],
6966        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
6967        // ONE place instead of being half-applied by the worker.
6968        sampling: Option<SpecSampling>,
6969        require_anchor: bool,
6970        max_ctx: usize,
6971    ) -> Result<SpecSession, (Option<Cache>, String)> {
6972        let pos = prefix.len();
6973        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
6974            Err((Some(cache), msg))
6975        };
6976        if self.mtp.is_none() {
6977            return fail(cache, "no MTP head attached (nothing to draft with)".into());
6978        }
6979        if pos == 0 {
6980            return fail(cache, "empty committed prefix".into());
6981        }
6982        if cache.pos != pos {
6983            let msg = format!(
6984                "restored cache pos {} != restored prefix len {pos}",
6985                cache.pos
6986            );
6987            return fail(cache, msg);
6988        }
6989        if draft_len != pos {
6990            return fail(
6991                cache,
6992                format!("draft plane len {draft_len} != restored prefix len {pos}"),
6993            );
6994        }
6995        if pos + suffix.len() >= max_ctx {
6996            return fail(
6997                cache,
6998                format!(
6999                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
7000                    pos + suffix.len(),
7001                ),
7002            );
7003        }
7004        let mut scratch = match MtpScratch::new(
7005            e,
7006            &self.cfg,
7007            max_ctx,
7008            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7009        ) {
7010            Ok(s) => s,
7011            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
7012        };
7013        if scratch.kv.ring.is_some() {
7014            return fail(
7015                cache,
7016                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
7017            );
7018        }
7019        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
7020            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
7021        {
7022            return fail(
7023                cache,
7024                format!(
7025                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
7026                     {}/{} bytes/token (stale entry across a format change)",
7027                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
7028                ),
7029            );
7030        }
7031        if pos > scratch.cap {
7032            return fail(
7033                cache,
7034                format!(
7035                    "draft plane rows {pos} exceed scratch capacity {}",
7036                    scratch.cap
7037                ),
7038            );
7039        }
7040        let kb = pos * draft_k_tok_bytes;
7041        let vb = pos * draft_v_tok_bytes;
7042        if draft_k.len() < kb || draft_v.len() < vb {
7043            return fail(
7044                cache,
7045                format!(
7046                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
7047                    draft_k.len(),
7048                    draft_v.len(),
7049                ),
7050            );
7051        }
7052        if kb > 0 {
7053            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
7054                return fail(cache, format!("draft K restore copy failed: {err}"));
7055            }
7056        }
7057        if vb > 0 {
7058            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
7059                return fail(cache, format!("draft V restore copy failed: {err}"));
7060            }
7061        }
7062        if let Err(err) = scratch.set_len(e, pos) {
7063            return fail(cache, format!("draft scratch len set failed: {err}"));
7064        }
7065        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
7066            // anchor upload failure is acceptance-only when a suffix feed follows (fill
7067            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
7068            // burst entry asserts committed + last_h + next_pred) — the caller says which.
7069            e.htod(last_h).ok()
7070        } else {
7071            None
7072        };
7073        if require_anchor && last_h_dev.is_none() {
7074            return fail(
7075                cache,
7076                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
7077            );
7078        }
7079        let mut committed = prefix;
7080        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
7081        // what the empty-suffix continuation assert in the burst entry requires.
7082        let next_pred;
7083        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
7084        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
7085        // drawing its own first token from the same row.
7086        let mut sctr = 0u32;
7087        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
7088        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
7089        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
7090        // after the suffix joins `committed` below.
7091        let mut boundary_capture: Option<SpecBoundaryCapture> = None;
7092        if !suffix.is_empty() {
7093            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
7094            // From here on the trunk cache mutates: failures return Err((None, _)) and
7095            // the worker serves the request cold-plain instead of reusing the carrier.
7096            let dirty =
7097                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
7098            let n_embd = self.cfg.n_embd as usize;
7099            let t = suffix.len();
7100            let mut h_rows = match e.uninit(t * n_embd) {
7101                Ok(b) => b,
7102                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
7103            };
7104            let mut feed_logits = Vec::new();
7105            let batched = t >= crate::hybrid_forward::PRIME_MIN_T
7106                && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
7107                && !e.frozen_cpu_experts_prefer_tokenwise_prime();
7108            if batched {
7109                // prefill_tick's prime arm: one request-level prime_cache call.
7110                match self.prime_cache(e, suffix, &mut cache, 0) {
7111                    Ok((l, _h_seed, hiddens)) => {
7112                        if let Err(err) = e.copy_into(&mut h_rows, 0, &hiddens, t * n_embd) {
7113                            return dirty(format!("suffix hidden copy: {err}"));
7114                        }
7115                        feed_logits = l;
7116                    }
7117                    Err(err) => return dirty(format!("suffix prime failed: {err}")),
7118                }
7119            } else {
7120                // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
7121                for (i, &tok) in suffix.iter().enumerate() {
7122                    match self.decode_step_h(e, tok, &mut cache) {
7123                        Ok((l, h)) => {
7124                            if let Err(err) = e.copy_into(&mut h_rows, i * n_embd, &h, n_embd) {
7125                                return dirty(format!("suffix hidden copy: {err}"));
7126                            }
7127                            feed_logits = l;
7128                        }
7129                        Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
7130                    }
7131                }
7132            }
7133            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
7134            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
7135            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
7136            // with T). Fill failures are acceptance-only — truncate to the restored rows
7137            // and continue; the burst's own set_len keeps the invariant.
7138            let mtp = self.mtp.as_ref().expect("mtp checked above");
7139            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7140            let embd_gpu = if spec_host_embd() {
7141                None
7142            } else {
7143                Some(
7144                    self.embd_gpu
7145                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7146                )
7147            };
7148            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7149            let fill_chunk = 4096usize;
7150            let mut filled = true;
7151            let mut start = 0usize;
7152            'fill: while start < t {
7153                let end = (start + fill_chunk).min(t);
7154                let tc = end - start;
7155                let Ok(mut phs) = e.zeros(tc * n_embd) else {
7156                    filled = false;
7157                    break 'fill;
7158                };
7159                let (src_lo, dst_off, n_copy) = if start == 0 {
7160                    (0, n_embd, (tc - 1) * n_embd)
7161                } else {
7162                    ((start - 1) * n_embd, 0, tc * n_embd)
7163                };
7164                if start == 0 {
7165                    if let Some(lh) = last_h_dev.as_ref() {
7166                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
7167                            filled = false;
7168                            break 'fill;
7169                        }
7170                    }
7171                }
7172                if n_copy > 0
7173                    && e.copy_view_into(
7174                        &mut phs,
7175                        dst_off,
7176                        &h_rows.slice(src_lo..src_lo + n_copy),
7177                        n_copy,
7178                    )
7179                    .is_err()
7180                {
7181                    filled = false;
7182                    break 'fill;
7183                }
7184                if self
7185                    .mtp_kv_fill(
7186                        e,
7187                        mtp,
7188                        &suffix[start..end],
7189                        &phs,
7190                        pos + start,
7191                        &mut scratch,
7192                        embd_dev,
7193                    )
7194                    .is_err()
7195                {
7196                    filled = false;
7197                    break 'fill;
7198                }
7199                start = end;
7200            }
7201            if !filled {
7202                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
7203                // so keep only the restored rows resident and let verify arbitrate.
7204                if let Err(err) = scratch.set_len(e, pos) {
7205                    return dirty(format!("scratch truncation after failed fill: {err}"));
7206                }
7207            }
7208            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
7209            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
7210            // finding (d)). Pre-lane, publication was armed only for COLD sessions
7211            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
7212            // non-continuation burst — but a converted hit's first burst IS a continuation,
7213            // so a growing conversation learned exactly ONE boundary and turn 3 could never
7214            // hit a longer prefix than turn 2 did.
7215            //
7216            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
7217            // line — the trunk is primed over the whole prompt, nothing is generated, and the
7218            // draft plane rows [0..prompt) are filled just above. That is a complete
7219            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
7220            // publishes; the worker's existing publication sweep picks it up because it is
7221            // keyed on `boundary_capture.is_some()` and is sampler- and resume-independent.
7222            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
7223            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
7224            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
7225            // publication is an optimization, never a correctness dependency.
7226            if spec_restore_republish_on() {
7227                debug_assert_eq!(
7228                    cache.pos,
7229                    pos + t,
7230                    "extended-entry capture must sit at the restored session's prompt end",
7231                );
7232                if let Ok(snap) = cache.snapshot(e) {
7233                    boundary_capture = Some(SpecBoundaryCapture {
7234                        snap,
7235                        pos: pos + t,
7236                        logits: feed_logits.clone(),
7237                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
7238                    });
7239                }
7240            }
7241            // continuation seed: the feed's boundary logits ARE the plain path's boundary
7242            // logits (same program), so greedy's argmax here is plain's first emitted token,
7243            // and the sampled draw is the cold sampled session's own first token.
7244            next_pred = Some(if sampled {
7245                let sp = sampling.expect("sampled implies a sampler");
7246                // `committed` is still the restored prefix here; the suffix joins it below —
7247                // so this is the last-N window over the WHOLE prompt, exactly the cold
7248                // session's own window at its first token.
7249                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
7250                match sample_boundary_token(
7251                    e,
7252                    &feed_logits,
7253                    &sp,
7254                    &hist,
7255                    &mut sctr,
7256                    "restore-suffix-feed",
7257                ) {
7258                    Ok(t) => t,
7259                    // the trunk is already fed: hand nothing back, the worker serves the
7260                    // request cold-plain. Never fall back to an argmax — that would put a
7261                    // greedy token in a sampled stream to save a slow path.
7262                    Err(err) => {
7263                        return dirty(format!("boundary token draw failed: {err}"));
7264                    }
7265                }
7266            } else {
7267                argmax(&feed_logits) as u32
7268            });
7269            let mut lh = match e.uninit(n_embd) {
7270                Ok(b) => b,
7271                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
7272            };
7273            if let Err(err) = e.copy_view_into(
7274                &mut lh,
7275                0,
7276                &h_rows.slice((t - 1) * n_embd..t * n_embd),
7277                n_embd,
7278            ) {
7279                return dirty(format!("boundary hidden copy: {err}"));
7280            }
7281            last_h_dev = Some(lh);
7282            committed.extend_from_slice(suffix);
7283        } else {
7284            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
7285            // ENTRY's boundary logits are the boundary row, and this is the token the cold
7286            // session emits from that same row. Owned here rather than in the worker so the
7287            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
7288            if boundary_logits.is_empty() {
7289                return fail(
7290                    cache,
7291                    "full-cover restore without the entry's boundary logits".into(),
7292                );
7293            }
7294            next_pred = Some(if sampled {
7295                let sp = sampling.expect("sampled implies a sampler");
7296                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
7297                match sample_boundary_token(
7298                    e,
7299                    boundary_logits,
7300                    &sp,
7301                    &hist,
7302                    &mut sctr,
7303                    "restore-full-cover",
7304                ) {
7305                    Ok(t) => t,
7306                    // nothing has been mutated on this shape — hand the carrier back and let
7307                    // the hit serve PLAIN (the banked pre-lane path).
7308                    Err(err) => {
7309                        return fail(cache, format!("boundary token draw failed: {err}"));
7310                    }
7311                }
7312            } else {
7313                argmax(boundary_logits) as u32
7314            });
7315        }
7316        Ok(SpecSession {
7317            cache,
7318            scratch,
7319            committed,
7320            last_h: last_h_dev,
7321            next_pred,
7322            sctr,
7323            uctr: 0,
7324            draft_ctx: None,
7325            pending_tok: None,
7326            turn_ckpt: None,
7327            telem: SpecTelemetryCounters::default(),
7328            capture_at: None,
7329            boundary_capture,
7330        })
7331    }
7332
7333    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
7334    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
7335    /// snapshot, or draft-KV row that only corrupts the following round.
7336    pub fn optipipe_compare_session_state(
7337        &self,
7338        e: &Engine,
7339        reference: &SpecSession,
7340        candidate: &SpecSession,
7341    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
7342        fn fail(what: &str) -> Box<dyn std::error::Error> {
7343            format!("optipipe state mismatch: {what}").into()
7344        }
7345        fn same_f32(a: &[f32], b: &[f32]) -> bool {
7346            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
7347        }
7348        fn compare_layers(
7349            es: &Engine,
7350            range: std::ops::Range<usize>,
7351            reference: &SpecSession,
7352            candidate: &SpecSession,
7353            report: &mut OptiForkStateIdentity,
7354        ) -> Result<(), Box<dyn std::error::Error>> {
7355            for il in range {
7356                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
7357                    (Some(a), Some(b)) => {
7358                        if a.len != b.len {
7359                            return Err(fail(&format!(
7360                                "layer {il} host KV len {} != {}",
7361                                a.len, b.len
7362                            )));
7363                        }
7364                        let ad = es.dtoh_i32(&a.len_d)?;
7365                        let bd = es.dtoh_i32(&b.len_d)?;
7366                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
7367                            return Err(fail(&format!(
7368                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
7369                                a.len,
7370                            )));
7371                        }
7372                        let kb = a.len * a.k_tok_bytes;
7373                        let vb = a.len * a.v_tok_bytes;
7374                        if kb > 0 {
7375                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
7376                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
7377                            if ak != bk {
7378                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
7379                                return Err(fail(&format!(
7380                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
7381                                    at / a.k_tok_bytes,
7382                                    at % a.k_tok_bytes,
7383                                    ak[at],
7384                                    bk[at],
7385                                )));
7386                            }
7387                        }
7388                        if vb > 0 {
7389                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
7390                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
7391                            if av != bv {
7392                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
7393                                return Err(fail(&format!(
7394                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
7395                                    at / a.v_tok_bytes,
7396                                    at % a.v_tok_bytes,
7397                                    av[at],
7398                                    bv[at],
7399                                )));
7400                            }
7401                        }
7402                        report.trunk_kv_bytes += kb + vb;
7403                    }
7404                    (None, None) => {}
7405                    _ => return Err(fail(&format!("layer {il} KV presence"))),
7406                }
7407                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
7408                    (Some(a), Some(b)) => {
7409                        let ac = es.dtoh(&a.conv_state)?;
7410                        let bc = es.dtoh(&b.conv_state)?;
7411                        if !same_f32(&ac, &bc) {
7412                            return Err(fail(&format!("layer {il} conv state")));
7413                        }
7414                        let as_ = es.dtoh(&a.ssm_state)?;
7415                        let bs = es.dtoh(&b.ssm_state)?;
7416                        if !same_f32(&as_, &bs) {
7417                            return Err(fail(&format!("layer {il} SSM state")));
7418                        }
7419                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
7420                    }
7421                    (None, None) => {}
7422                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
7423                }
7424            }
7425            Ok(())
7426        }
7427
7428        if reference.committed != candidate.committed {
7429            return Err(fail("committed token ids"));
7430        }
7431        if reference.cache.pos != candidate.cache.pos
7432            || reference.cache.max_ctx != candidate.cache.max_ctx
7433        {
7434            return Err(fail("cache pos/capacity"));
7435        }
7436        if reference.pending_tok != candidate.pending_tok
7437            || reference.next_pred != candidate.next_pred
7438            || reference.sctr != candidate.sctr
7439            || reference.uctr != candidate.uctr
7440        {
7441            return Err(fail("pending/prediction/counter tail"));
7442        }
7443
7444        let mut report = OptiForkStateIdentity::default();
7445        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
7446            let rt = crate::pp::PpNRt::get(e)?;
7447            for stage in 0..rt.n_stages() {
7448                let _scope = rt.enter(stage);
7449                compare_layers(
7450                    rt.engine(stage, e),
7451                    fence[stage]..fence[stage + 1],
7452                    reference,
7453                    candidate,
7454                    &mut report,
7455                )?;
7456            }
7457        } else {
7458            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
7459        }
7460
7461        let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
7462        if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
7463            return Err(fail("draft scratch length"));
7464        }
7465        let kb = a.len * a.k_tok_bytes;
7466        let vb = a.len * a.v_tok_bytes;
7467        if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
7468            return Err(fail("draft scratch K bytes"));
7469        }
7470        if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
7471            return Err(fail("draft scratch V bytes"));
7472        }
7473        report.scratch_kv_bytes = kb + vb;
7474
7475        match (&reference.last_h, &candidate.last_h) {
7476            (Some(a), Some(b)) => {
7477                let ah = e.dtoh(a)?;
7478                let bh = e.dtoh(b)?;
7479                if !same_f32(&ah, &bh) {
7480                    return Err(fail("last hidden/seed bytes"));
7481                }
7482                report.hidden_bytes = ah.len() * 4;
7483            }
7484            (None, None) => {}
7485            _ => return Err(fail("last hidden/seed presence")),
7486        }
7487        Ok(report)
7488    }
7489
7490    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
7491    /// retained prompt-end checkpoint, so a request whose prompt matches
7492    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
7493    ///
7494    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
7495    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
7496    /// restored from the device copy taken there, draft scratch length reset, `committed`
7497    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
7498    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
7499    /// every burst after it are identical to a cold run of the same token stream — the
7500    /// committed-tokens-authoritative contract.
7501    ///
7502    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
7503    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
7504    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
7505    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
7506    /// (the scratch KV, the resident embedding), none of which the rewind moves.
7507    ///
7508    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
7509    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
7510    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
7511    pub fn spec_rewind_to_checkpoint(
7512        &self,
7513        e: &Engine,
7514        sess: &mut SpecSession,
7515    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
7516        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
7517            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
7518        }) {
7519            return Err(
7520                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
7521            );
7522        }
7523        let Some(ckpt) = sess.turn_ckpt.take() else {
7524            return Ok(None);
7525        };
7526        assert!(
7527            ckpt.pos <= sess.committed.len(),
7528            "checkpoint past committed ({} > {})",
7529            ckpt.pos,
7530            sess.committed.len()
7531        );
7532        // Restore through each layer's owning engine. A single primary-engine rollback is not
7533        // sufficient when the serving cache is stage-owned under cross-device PP.
7534        crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
7535        debug_assert_eq!(
7536            sess.cache.pos, ckpt.pos,
7537            "rollback landed off the checkpoint"
7538        );
7539        sess.scratch.set_len(e, ckpt.pos)?;
7540        sess.committed.truncate(ckpt.pos);
7541        sess.last_h = Some(ckpt.last_h);
7542        sess.next_pred = None;
7543        sess.pending_tok = None;
7544        Ok(Some(ckpt.pos))
7545    }
7546
7547    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
7548    /// checkpoint without re-priming the checkpoint prefix.
7549    ///
7550    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
7551    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
7552    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
7553    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
7554    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
7555    ///
7556    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
7557    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
7558    pub fn spec_grow_and_rewind_to_checkpoint(
7559        &self,
7560        e: &Engine,
7561        sess: &mut SpecSession,
7562        target_cap: usize,
7563    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
7564        if target_cap <= sess.cache.max_ctx {
7565            return self.spec_rewind_to_checkpoint(e, sess);
7566        }
7567        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
7568            return Ok(None);
7569        };
7570        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
7571            return Err(format!(
7572                "checkpoint pos {} outside committed length {}",
7573                ckpt.pos,
7574                sess.committed.len(),
7575            )
7576            .into());
7577        }
7578        if ckpt.pos > target_cap {
7579            return Err(format!(
7580                "checkpoint pos {} exceeds grown capacity {target_cap}",
7581                ckpt.pos,
7582            )
7583            .into());
7584        }
7585
7586        let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
7587        let mut grown_scratch = MtpScratch::new(
7588            e,
7589            &self.cfg,
7590            target_cap,
7591            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7592        )?;
7593        crate::pp::restore_cache_checkpoint(
7594            e,
7595            &self.cfg,
7596            Some(&sess.cache),
7597            &mut grown_cache,
7598            &ckpt.snap,
7599        )?;
7600
7601        let src = &sess.scratch.kv;
7602        let dst = &mut grown_scratch.kv;
7603        if ckpt.pos > src.len
7604            || src.kv_dim_k != dst.kv_dim_k
7605            || src.kv_dim_v != dst.kv_dim_v
7606            || src.k_tok_bytes != dst.k_tok_bytes
7607            || src.v_tok_bytes != dst.v_tok_bytes
7608        {
7609            return Err(format!(
7610                "checkpoint draft layout mismatch (pos {}, source len {})",
7611                ckpt.pos, src.len,
7612            )
7613            .into());
7614        }
7615        let kb = ckpt.pos * src.k_tok_bytes;
7616        let vb = ckpt.pos * src.v_tok_bytes;
7617        if kb > 0 {
7618            e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
7619        }
7620        if vb > 0 {
7621            e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
7622        }
7623        grown_scratch.set_len(e, ckpt.pos)?;
7624        // The old scratch is dropped immediately after publication below. Bound its D2D reads
7625        // first; growth happens once per rewritten turn, outside the decode hot loop.
7626        e.stream().synchronize()?;
7627
7628        let ckpt = sess
7629            .turn_ckpt
7630            .take()
7631            .expect("checkpoint remained present through transactional grow");
7632        let pos = ckpt.pos;
7633        sess.cache = grown_cache;
7634        sess.scratch = grown_scratch;
7635        sess.committed.truncate(pos);
7636        sess.last_h = Some(ckpt.last_h);
7637        sess.next_pred = None;
7638        sess.pending_tok = None;
7639        sess.draft_ctx = None;
7640        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
7641        debug_assert_eq!(
7642            sess.scratch.kv.len, pos,
7643            "grown draft rewind landed off checkpoint"
7644        );
7645        Ok(Some(pos))
7646    }
7647
7648    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
7649    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
7650    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
7651    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
7652    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
7653    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
7654    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
7655    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
7656    /// park-time flush is a future request whose sampler is not knowable here (residual
7657    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
7658    pub fn spec_flush_pending(
7659        &self,
7660        e: &Engine,
7661        sess: &mut SpecSession,
7662        sampling: Option<SpecSampling>,
7663    ) -> Result<(), Box<dyn std::error::Error>> {
7664        let Some(b) = sess.pending_tok.take() else {
7665            return Ok(());
7666        };
7667        let mtp = self
7668            .mtp
7669            .as_ref()
7670            .expect("pending carry requires an MTP head");
7671        let n_embd = self.cfg.n_embd as usize;
7672        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
7673        let embd_gpu = if spec_host_embd() {
7674            None
7675        } else {
7676            Some(
7677                self.embd_gpu
7678                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
7679            )
7680        };
7681        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
7682        let pos_b = sess.cache.pos;
7683        sess.scratch.set_len(e, pos_b)?;
7684        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
7685        sess.next_pred = Some(match sampling {
7686            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
7687                // window includes `b` itself: it is committed by this pass, and the pre-lane
7688                // code never counted a boundary token in the penalty history at all.
7689                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
7690                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
7691            }
7692            _ => argmax(&lg_b) as u32,
7693        });
7694        let anchor = sess
7695            .last_h
7696            .as_ref()
7697            .expect("pending carry requires last_h (the predecessor-row anchor)");
7698        self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
7699        sess.last_h = Some(hb);
7700        sess.committed.push(b);
7701        Ok(())
7702    }
7703
7704    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
7705    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
7706    /// rounds through that same graph. Other model families keep their eager T=1 contract.
7707    fn spec_target_step_h(
7708        &self,
7709        e: &Engine,
7710        token: u32,
7711        cache: &mut Cache,
7712    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
7713        if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
7714            return self.decode_step_h(e, token, cache);
7715        }
7716        let pos0 = cache.pos;
7717        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
7718        Ok((e.dtoh(&logits)?, hidden))
7719    }
7720
7721    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
7722    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
7723    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
7724    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
7725    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
7726    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
7727    /// dispatch sites cannot drift apart again.
7728    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
7729    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
7730    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
7731    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
7732    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
7733    /// eligibility sites so they cannot drift (the qwen35_serving_class lesson).
7734    fn mtp_graph_capturable(&self) -> bool {
7735        self.mtp
7736            .as_ref()
7737            .map(|m| match &m.ffn {
7738                crate::hybrid::Ffn::Dense { .. } => true,
7739                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
7740            })
7741            .unwrap_or(false)
7742    }
7743
7744    fn qwen35_serving_class(&self) -> bool {
7745        matches!(
7746            self.cfg.arch,
7747            memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
7748        )
7749    }
7750
7751    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
7752    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
7753    /// session already exist.
7754    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
7755        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
7756            || !spec_devacc()
7757            || spec_replay_env_enabled()
7758            || spec_stream()
7759            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
7760            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
7761            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
7762            || std::env::var("MEMRA_SPEC_PMIN")
7763                .ok()
7764                .and_then(|v| v.parse::<f32>().ok())
7765                .unwrap_or(0.0)
7766                > 0.0
7767            || self.is_gemma4_e4b()
7768            || self.cfg.gemma4.is_some()
7769            || self.mtp.is_none()
7770        {
7771            return false;
7772        }
7773        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
7774            return false;
7775        };
7776        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
7777            return false;
7778        }
7779        crate::pp::PpNRt::get(e)
7780            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
7781            .unwrap_or(false)
7782    }
7783
7784    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
7785    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
7786    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
7787    #[allow(clippy::too_many_arguments)]
7788    pub fn generate_spec_session_pair(
7789        &self,
7790        e: &Engine,
7791        sess_a: &mut SpecSession,
7792        max_new_a: usize,
7793        k_a: usize,
7794        sess_b: &mut SpecSession,
7795        max_new_b: usize,
7796        k_b: usize,
7797    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
7798    {
7799        if !self.spec_pipe_available(e) {
7800            return Err("two-session speculative pipeline is outside its reduced matrix".into());
7801        }
7802        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
7803            return Err(
7804                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
7805            );
7806        }
7807        for sess in [&*sess_a, &*sess_b] {
7808            if sess.committed.is_empty()
7809                || sess.last_h.is_none()
7810                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
7811            {
7812                return Err("two-session speculative pipeline requires warm continuations".into());
7813            }
7814        }
7815
7816        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
7817            && !spec_host_embd()
7818            && self.mtp_graph_capturable()
7819            && !crate::model::full_prec_enabled();
7820        let graph_a = graph_ok && k_a + 2 < 96;
7821        let graph_b = graph_ok && k_b + 2 < 96;
7822        let was_tracking = e.ctx().is_event_tracking();
7823        if (graph_a || graph_b) && was_tracking {
7824            unsafe {
7825                e.ctx().disable_event_tracking();
7826            }
7827        }
7828
7829        static LOGGED: std::sync::Once = std::sync::Once::new();
7830        LOGGED.call_once(|| {
7831            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
7832        });
7833        let sync = std::sync::Arc::new(SpecPipeSync::new());
7834        let lane_a = SpecPipeLane {
7835            sync: sync.clone(),
7836            lane: 0,
7837        };
7838        let lane_b = SpecPipeLane { sync, lane: 1 };
7839        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
7840        let (result_a, result_b) = std::thread::scope(|scope| {
7841            let b = scope.spawn(move || {
7842                let mut finish = SpecPipeFinish::new(&lane_b);
7843                let sess_b = unsafe { sess_b_ptr.get_mut() };
7844                let result = e
7845                    .ctx()
7846                    .bind_to_thread()
7847                    .map_err(|err| err.to_string())
7848                    .and_then(|_| {
7849                        self.generate_spec_inner2(
7850                            e,
7851                            &[],
7852                            max_new_b,
7853                            k_b,
7854                            graph_b,
7855                            Some(sess_b),
7856                            None,
7857                            None,
7858                            None,
7859                            None,
7860                            Some(&lane_b),
7861                        )
7862                        .map_err(|err| err.to_string())
7863                    });
7864                finish.close(result.is_err());
7865                result
7866            });
7867            let mut finish = SpecPipeFinish::new(&lane_a);
7868            let result_a = self.generate_spec_inner2(
7869                e,
7870                &[],
7871                max_new_a,
7872                k_a,
7873                graph_a,
7874                Some(sess_a),
7875                None,
7876                None,
7877                None,
7878                None,
7879                Some(&lane_a),
7880            );
7881            finish.close(result_a.is_err());
7882            let result_b = b
7883                .join()
7884                .map_err(|_| "paired speculative session B panicked".to_string())
7885                .and_then(|r| r);
7886            (result_a, result_b)
7887        });
7888
7889        if (graph_a || graph_b) && was_tracking {
7890            unsafe {
7891                e.ctx().enable_event_tracking();
7892            }
7893        }
7894        let result_a = result_a?;
7895        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
7896        Ok((result_a, result_b))
7897    }
7898
7899    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
7900    /// message rendered through the chat template continuation). Returns (new tokens emitted,
7901    /// drafted, accepted); session.committed grows by suffix + emitted.
7902    pub fn generate_spec_session(
7903        &self,
7904        e: &Engine,
7905        sess: &mut SpecSession,
7906        suffix: &[u32],
7907        max_new: usize,
7908        k: usize,
7909    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7910        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
7911    }
7912
7913    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
7914    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
7915    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
7916    /// for the filtered target (feat/filtered-spec).
7917    ///
7918    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
7919    /// output — once right after the prime's first token, then once per round commit — so a
7920    /// streaming caller can flush text at round cadence instead of once per burst. The slices
7921    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
7922    /// timing only: token bytes, session state, and exactness are untouched.
7923    ///
7924    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
7925    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
7926    /// the caller's scheduler regains control without waiting the burst out. Burst size is
7927    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
7928    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
7929    /// drains and the defensive tail flush can land with nothing new committed).
7930    #[allow(clippy::too_many_arguments)]
7931    pub fn generate_spec_session_sampled(
7932        &self,
7933        e: &Engine,
7934        sess: &mut SpecSession,
7935        suffix: &[u32],
7936        max_new: usize,
7937        k: usize,
7938        sampling: Option<SpecSampling>,
7939        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7940    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7941        self.generate_spec_session_sampled_prime_split(
7942            e, sess, suffix, max_new, k, sampling, None, on_commit,
7943        )
7944    }
7945
7946    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
7947    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
7948    /// pass `None` and stay on the existing zero-prime path.
7949    #[allow(clippy::too_many_arguments)]
7950    pub fn generate_spec_session_sampled_prime_split(
7951        &self,
7952        e: &Engine,
7953        sess: &mut SpecSession,
7954        suffix: &[u32],
7955        max_new: usize,
7956        k: usize,
7957        sampling: Option<SpecSampling>,
7958        prime_split: Option<usize>,
7959        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7960    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7961        self.generate_spec_session_constrained_prime_split(
7962            e,
7963            sess,
7964            suffix,
7965            max_new,
7966            k,
7967            sampling,
7968            None,
7969            prime_split,
7970            on_commit,
7971        )
7972    }
7973
7974    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
7975    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
7976    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
7977    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
7978    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
7979    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
7980    /// may drop (drafter is unconstrained); that is measured, not hidden.
7981    #[allow(clippy::too_many_arguments)]
7982    pub fn generate_spec_session_constrained(
7983        &self,
7984        e: &Engine,
7985        sess: &mut SpecSession,
7986        suffix: &[u32],
7987        max_new: usize,
7988        k: usize,
7989        sampling: Option<SpecSampling>,
7990        constraint: Option<&mut dyn SpecConstraint>,
7991        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
7992    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
7993        self.generate_spec_session_constrained_prime_split(
7994            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
7995        )
7996    }
7997
7998    #[allow(clippy::too_many_arguments)]
7999    pub fn generate_spec_session_constrained_prime_split(
8000        &self,
8001        e: &Engine,
8002        sess: &mut SpecSession,
8003        suffix: &[u32],
8004        max_new: usize,
8005        k: usize,
8006        sampling: Option<SpecSampling>,
8007        constraint: Option<&mut dyn SpecConstraint>,
8008        prime_split: Option<usize>,
8009        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8010    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8011        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
8012            return Err(
8013                "constrained spec decode is greedy-only (worker routes sampled \
8014                        constrained to plain decode)"
8015                    .into(),
8016            );
8017        }
8018        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
8019        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
8020        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
8021        // serve continuation case — consume the carry in-loop with zero solo passes.
8022        if sess.pending_tok.is_some()
8023            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
8024        {
8025            self.spec_flush_pending(e, sess, sampling)?;
8026        }
8027
8028        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
8029        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
8030        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
8031        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8032            && !spec_host_embd()
8033            && self.mtp_graph_capturable()
8034            && k + 2 < 96
8035            && !crate::model::full_prec_enabled();
8036        let was_tracking = e.ctx().is_event_tracking();
8037        if graph_draft && was_tracking {
8038            unsafe {
8039                e.ctx().disable_event_tracking();
8040            }
8041        }
8042        let r = self.generate_spec_inner2(
8043            e,
8044            suffix,
8045            max_new,
8046            k,
8047            graph_draft,
8048            Some(sess),
8049            sampling,
8050            constraint,
8051            on_commit,
8052            prime_split,
8053            None,
8054        );
8055        if graph_draft && was_tracking {
8056            unsafe {
8057                e.ctx().enable_event_tracking();
8058            }
8059        }
8060        let (out, d, a) = r?;
8061        Ok((out, d, a))
8062    }
8063
8064    pub fn generate_spec(
8065        &self,
8066        e: &Engine,
8067        prompt: &[u32],
8068        max_new: usize,
8069        k: usize,
8070    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8071        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
8072        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
8073        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8074            && !spec_host_embd()
8075            && self.mtp_graph_capturable()
8076            && k + 2 < 96
8077            && !crate::model::full_prec_enabled();
8078        if !graph_draft {
8079            return self.generate_spec_inner2(
8080                e, prompt, max_new, k, false, None, None, None, None, None, None,
8081            );
8082        }
8083        let was_tracking = e.ctx().is_event_tracking();
8084        if was_tracking {
8085            unsafe {
8086                e.ctx().disable_event_tracking();
8087            }
8088        }
8089        let r = self.generate_spec_inner2(
8090            e, prompt, max_new, k, true, None, None, None, None, None, None,
8091        );
8092        if was_tracking {
8093            unsafe {
8094                e.ctx().enable_event_tracking();
8095            }
8096        }
8097        r
8098    }
8099
8100    fn generate_spec_inner2(
8101        &self,
8102        e: &Engine,
8103        prompt: &[u32],
8104        max_new: usize,
8105        k: usize,
8106        graph_draft: bool,
8107        mut sess: Option<&mut SpecSession>,
8108        sampling: Option<SpecSampling>,
8109        mut constraint: Option<&mut dyn SpecConstraint>,
8110        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8111        prime_split: Option<usize>,
8112        pipe: Option<&SpecPipeLane>,
8113    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8114        assert!(k >= 1, "k must be >= 1");
8115        if let Some(p) = pipe {
8116            p.setup_begin()?;
8117        }
8118        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
8119        let mut flushed = 0usize;
8120        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
8121        // at the next round boundary (same exit as max_new reached — the session tail runs).
8122        // Initialized by the unconditional post-prime flush below.
8123        let mut keep_going;
8124        let mtp = self
8125            .mtp
8126            .as_ref()
8127            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
8128        let n_vocab = self.output.out_features();
8129        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
8130        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
8131        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
8132        let d_vocab = mtp
8133            .shared_head_head
8134            .as_ref()
8135            .unwrap_or(&self.output)
8136            .out_features();
8137        let n_embd = self.cfg.n_embd as usize;
8138        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
8139        // already committed (their state is in the caches); 0 = fresh single-shot call.
8140        let session_mode = sess.is_some();
8141        let max_ctx = match sess.as_ref() {
8142            Some(s) => s.cache.max_ctx,
8143            None => prompt.len() + max_new + k + 8,
8144        };
8145        let mut own_cache;
8146        let mut own_scratch;
8147        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
8148        // (requested split, destination slot). Single-shot per burst; fresh calls have none.
8149        let mut sess_capture: Option<(Option<usize>, &mut Option<SpecBoundaryCapture>)> = None;
8150        let (
8151            cache,
8152            scratch,
8153            mut sess_tail,
8154            mut sess_draft_slot,
8155            mut sess_pending_slot,
8156            sess_ckpt_slot,
8157            sess_telem,
8158        ): (
8159            &mut Cache,
8160            &mut MtpScratch,
8161            Option<(
8162                &mut Vec<u32>,
8163                &mut Option<CudaSlice<f32>>,
8164                &mut Option<u32>,
8165                &mut u32,
8166                &mut u32,
8167            )>,
8168            Option<&mut Option<DraftGraphCtx>>,
8169            Option<&mut Option<u32>>,
8170            Option<&mut Option<SpecCheckpoint>>,
8171            Option<&SpecTelemetryCounters>,
8172        ) = match sess.take() {
8173            Some(sr) => {
8174                let SpecSession {
8175                    cache,
8176                    scratch,
8177                    committed,
8178                    last_h,
8179                    next_pred,
8180                    sctr: s_sctr,
8181                    uctr: s_uctr,
8182                    draft_ctx,
8183                    pending_tok,
8184                    turn_ckpt,
8185                    telem,
8186                    capture_at,
8187                    boundary_capture,
8188                } = sr;
8189                sess_capture = Some((capture_at.take(), boundary_capture));
8190                (
8191                    cache,
8192                    scratch,
8193                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
8194                    Some(draft_ctx),
8195                    Some(pending_tok),
8196                    Some(turn_ckpt),
8197                    Some(telem),
8198                )
8199            }
8200            None => {
8201                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
8202                // `Cache::new` verbatim.
8203                own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
8204                // Persistent scratch = max_ctx rows (~2KB/token quantized).
8205                own_scratch = MtpScratch::new(
8206                    e,
8207                    &self.cfg,
8208                    max_ctx,
8209                    self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8210                )?;
8211                (
8212                    &mut own_cache,
8213                    &mut own_scratch,
8214                    None,
8215                    None,
8216                    None,
8217                    None,
8218                    None,
8219                )
8220            }
8221        };
8222        let base = cache.pos;
8223        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
8224        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
8225        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
8226        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
8227        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
8228        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
8229        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
8230        // acceptance-only — exactness is verify's job either way).
8231        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
8232        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
8233        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
8234        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
8235        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
8236        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
8237        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
8238        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
8239        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
8240        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
8241        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
8242        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
8243        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
8244        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
8245        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
8246        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
8247        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
8248        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
8249        // + fallback seam).
8250        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
8251        // bar — the retained verify-state commit proven equivalent to sequential serving —
8252        // was waiting on this arch running the serving batched verify class, which the
8253        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
8254        // replay-free commit consumes is now produced by the SAME serving-class verify that
8255        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
8256        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
8257        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
8258        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
8259        // rollback + A/B seam.
8260        let spec_replay = spec_replay_env_enabled();
8261        if constraint.is_some() && spec_replay {
8262            return Err(
8263                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
8264                        (legacy replay commits an unmasked bonus)"
8265                    .into(),
8266            );
8267        }
8268        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
8269        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
8270        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
8271        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
8272
8273        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
8274        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
8275        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
8276        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
8277        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
8278        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
8279        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
8280        // generation exactly where the last turn stopped — no prime at all. The stashed
8281        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
8282        // committed.last() by the same rule this entry applies to a cold prime's last row —
8283        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
8284        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
8285        // where the sampler and the session's Philox counters were live). `last_h` seeds the
8286        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
8287        let continuation = prompt.is_empty();
8288        if continuation {
8289            assert!(session_mode, "empty prompt requires a session");
8290            assert!(
8291                sess_tail
8292                    .as_ref()
8293                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
8294                        && lh.is_some()
8295                        && (np.is_some() || carried_pending.is_some())),
8296                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
8297            );
8298        }
8299        let mut prime_logits;
8300        let mut prompt_h: Option<CudaSlice<f32>> = None;
8301        let t_prime = std::time::Instant::now();
8302        let batched_prime = !continuation
8303            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
8304            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
8305            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
8306        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
8307        if prime_split.is_some() && (continuation || base != 0) {
8308            return Err("spec prime split is cold-session-only".into());
8309        }
8310        if continuation {
8311            prime_logits = Vec::new();
8312        } else if let Some(split) = prime_split {
8313            if split < crate::hybrid_forward::PRIME_MIN_T {
8314                return Err(format!(
8315                    "spec prime split {split} is below PRIME_MIN_T {}",
8316                    crate::hybrid_forward::PRIME_MIN_T,
8317                )
8318                .into());
8319            }
8320            // Mirror the plain worker's affinity boundary exactly. The prefix is a request-level
8321            // prime (`queued_after` keeps Step35 arm selection independent of this stop); a tail
8322            // below PRIME_MIN_T then takes the same eager tokenwise continuation as prefill_tick.
8323            // Retain every hidden row so the draft scratch fill remains one coherent prompt.
8324            let mut h_all = e.uninit(prompt.len() * n_embd)?;
8325            let (l, _, h_prefix) =
8326                self.prime_cache(e, &prompt[..split], &mut *cache, prompt.len() - split)?;
8327            e.copy_into(&mut h_all, 0, &h_prefix, split * n_embd)?;
8328            prime_logits = l;
8329            // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm states
8330            // are about to be advanced in place by the tail prime, so this is the ONLY moment
8331            // the boundary's recurrent state exists. Capture iff the worker requested exactly
8332            // this split. cache.pos == split here (the prefix prime just finished). A failed
8333            // snapshot is silent (turn_ckpt convention) — publication is an optimization,
8334            // never a correctness dependency.
8335            if let Some((requested, slot)) = sess_capture.as_mut() {
8336                if *requested == Some(split) {
8337                    debug_assert_eq!(cache.pos, split, "boundary capture off the prime split");
8338                    if let Ok(snap) = cache.snapshot(e) {
8339                        **slot = Some(SpecBoundaryCapture {
8340                            snap,
8341                            pos: split,
8342                            logits: prime_logits.clone(),
8343                            // rows [0..split) of h_all are the prefix prime's hiddens — copied
8344                            // just above, before the tail prime overwrites nothing (append-only).
8345                            last_h: capture_boundary_hidden(e, &h_all, split, n_embd),
8346                        });
8347                    }
8348                }
8349            }
8350            let tail = &prompt[split..];
8351            if tail.len() >= crate::hybrid_forward::PRIME_MIN_T
8352                && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
8353                && !e.frozen_cpu_experts_prefer_tokenwise_prime()
8354            {
8355                let (l, _, h_tail) = self.prime_cache(e, tail, &mut *cache, 0)?;
8356                e.copy_into(&mut h_all, split * n_embd, &h_tail, tail.len() * n_embd)?;
8357                prime_logits = l;
8358            } else {
8359                for (i, &tok) in tail.iter().enumerate() {
8360                    let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
8361                    e.copy_into(&mut h_all, (split + i) * n_embd, &h, n_embd)?;
8362                    prime_logits = l;
8363                }
8364            }
8365            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8366                eprintln!("[spec-prime] affinity split={split} tail={}", tail.len());
8367            }
8368            prompt_h = Some(h_all);
8369        } else if batched_prime {
8370            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
8371            prime_logits = l;
8372            prompt_h = Some(hiddens);
8373        } else {
8374            prime_logits = Vec::new();
8375            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
8376            for (i, &tok) in prompt.iter().enumerate() {
8377                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
8378                if let Some(ph) = prompt_h.as_mut() {
8379                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
8380                }
8381                prime_logits = l;
8382            }
8383        }
8384        e.stream().synchronize()?;
8385        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
8386        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
8387        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
8388        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
8389        // prime_split. The mid-prompt capture above already consumed the request if it matched.
8390        if !continuation && base == 0 {
8391            if let Some((requested, slot)) = sess_capture.as_mut() {
8392                if *requested == Some(prompt.len()) && slot.is_none() {
8393                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
8394                    if let Ok(snap) = cache.snapshot(e) {
8395                        **slot = Some(SpecBoundaryCapture {
8396                            snap,
8397                            pos: prompt.len(),
8398                            logits: prime_logits.clone(),
8399                            last_h: prompt_h
8400                                .as_ref()
8401                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
8402                                .unwrap_or_default(),
8403                        });
8404                    }
8405                }
8406            }
8407        }
8408        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
8409        // prime-subtraction hack.
8410        crate::PRIME_NANOS.store(
8411            t_prime.elapsed().as_nanos() as u64,
8412            std::sync::atomic::Ordering::Relaxed,
8413        );
8414
8415        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8416        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
8417        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
8418        let host_embd = spec_host_embd();
8419        let embd_gpu = if host_embd {
8420            None
8421        } else {
8422            Some(
8423                self.embd_gpu
8424                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8425            )
8426        };
8427        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8428        if host_embd {
8429            eprintln!(
8430                "[spec] host-row embedding: {} bytes kept off HBM",
8431                self.embd.raw.len()
8432            );
8433        }
8434        let mut out: Vec<u32> = Vec::with_capacity(max_new);
8435        let mut total_drafted = 0usize;
8436        let mut total_accepted = 0usize;
8437
8438        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
8439        // The sampler config, the session's Philox counters and the penalty window are parsed
8440        // HERE, above the boundary-token selection, because the boundary token must be drawn
8441        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
8442        // selection, which is the whole mechanical reason the boundary token was an argmax:
8443        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
8444        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
8445        // below takes the argmax path it always took).
8446        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
8447        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
8448        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
8449        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
8450        let sp = sampling.unwrap_or_else(|| SpecSampling {
8451            temp: std::env::var("MEMRA_SPEC_TEMP")
8452                .ok()
8453                .and_then(|v| v.parse().ok())
8454                .unwrap_or(0.0),
8455            seed: std::env::var("MEMRA_SEED")
8456                .ok()
8457                .and_then(|v| v.parse().ok())
8458                .unwrap_or(42),
8459            top_k: std::env::var("MEMRA_TOP_K")
8460                .ok()
8461                .and_then(|v| v.parse().ok())
8462                .unwrap_or(0),
8463            top_p: std::env::var("MEMRA_TOP_P")
8464                .ok()
8465                .and_then(|v| v.parse().ok())
8466                .unwrap_or(1.0),
8467            min_p: std::env::var("MEMRA_MIN_P")
8468                .ok()
8469                .and_then(|v| v.parse().ok())
8470                .unwrap_or(0.0),
8471            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
8472                .ok()
8473                .and_then(|v| v.parse().ok())
8474                .unwrap_or(0),
8475            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
8476                .ok()
8477                .and_then(|v| v.parse().ok())
8478                .unwrap_or(1.0),
8479            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
8480                .ok()
8481                .and_then(|v| v.parse().ok())
8482                .unwrap_or(0.0),
8483            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
8484                .ok()
8485                .and_then(|v| v.parse().ok())
8486                .unwrap_or(0.0),
8487        });
8488        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
8489        let sampled = sp_temp > 0.0;
8490        // Counters resume from the session (burst continuity: randomness must never repeat
8491        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
8492        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
8493        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
8494        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
8495        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
8496        // for the penalized+filtered target). History = generated tokens, host-tracked window.
8497        let pen_on = sampled
8498            && sp.penalty_last_n > 0
8499            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
8500        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
8501        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
8502        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
8503        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
8504        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
8505        // which is what the API contract says and what the plain sampler's own `history` does.
8506        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
8507        let mut pen_hist: Vec<u32> = if pen_on {
8508            let sess_hist: &[u32] = if spec_pen_session_on() {
8509                sess_tail
8510                    .as_ref()
8511                    .map(|(c, ..)| c.as_slice())
8512                    .unwrap_or(&[])
8513            } else {
8514                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
8515            };
8516            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
8517        } else {
8518            Vec::new()
8519        };
8520        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
8521        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
8522        // request's own filtered/penalized target through the session's Philox stream
8523        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
8524        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
8525        // Emit it, then FEED it to establish the loop invariant below.
8526        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
8527        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
8528        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
8529        // prompt's last logits (plain constrained-greedy identity); a continuation without
8530        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
8531        // worker never resumes constrained sessions from the pool, so this cannot fire).
8532        if let Some(c) = constraint.as_deref_mut() {
8533            if continuation && carried_pending.is_none() {
8534                return Err("constrained spec continuation requires a carried pending \
8535                            (pool resume is unconstrained-only)"
8536                    .into());
8537            }
8538            if !continuation {
8539                c.mask_logits(&mut prime_logits)
8540                    .map_err(|e2| format!("constraint: {e2}"))?;
8541            }
8542        }
8543        let mut last_token = if let Some(b) = carried_pending {
8544            b
8545        } else if continuation {
8546            // A continuation's boundary token was DRAWN by the burst that stashed it (the
8547            // session tail below), or by `spec_session_from_restored` for a converted
8548            // prefix-cache hit — in both cases from the correct logits row with this same
8549            // session's Philox stream, which is why it can be consumed here as-is.
8550            sess_tail.as_ref().unwrap().2.unwrap()
8551        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
8552            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
8553        } else {
8554            // greedy (byte contract), the rollback door, or constrained (masked-argmax
8555            // identity — the worker routes sampled+constrained to the plain path, and this
8556            // function refuses the combination outright above).
8557            argmax(&prime_logits) as u32
8558        };
8559        if pen_on {
8560            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
8561            // emitted token into its penalty history, and pre-lane the burst's first token
8562            // was invisible to penalties forever (never pushed, and never in `committed`
8563            // until this burst's tail). Covers the carry/continuation seeds too — neither is
8564            // in `committed` yet.
8565            pen_hist.push(last_token);
8566        }
8567        if carried_pending.is_none() {
8568            out.push(last_token);
8569            // grammar advances with every emitted token (carried pendings were consumed
8570            // by the burst that emitted them).
8571            if let Some(c) = constraint.as_deref_mut() {
8572                c.consume(last_token)
8573                    .map_err(|e2| format!("constraint: {e2}"))?;
8574            }
8575        }
8576        if continuation {
8577            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
8578            // overhang so the chain's first append lands at slot base (== committed.len()).
8579            scratch.set_len(e, base)?;
8580        }
8581        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
8582        // concatenating to the full `out`). Called after the prime's first token and after each
8583        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
8584        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
8585        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
8586        fn flush_commit(
8587            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
8588            out: &[u32],
8589            flushed: &mut usize,
8590        ) -> bool {
8591            if let Some(f) = cb.as_mut() {
8592                let keep = f(&out[*flushed..]);
8593                *flushed = out.len();
8594                keep
8595            } else {
8596                true
8597            }
8598        }
8599        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
8600        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
8601        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
8602        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
8603        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
8604        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
8605        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
8606        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
8607        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
8608        // those, so their residual mass is p(x), correct by construction).
8609        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
8610            match &mtp.d2t {
8611                Some(map) => Some(e.htod_u32_v(map)?),
8612                None => None,
8613            }
8614        } else {
8615            None
8616        };
8617        let mut q_full_buf: Option<CudaSlice<f32>> = None;
8618        // host Philox4x32-10 (mirrors spec_sample.cu; independent stream via ctr_lo tag)
8619        let host_u01 = |seed: u64, ctr: u32| -> f32 {
8620            let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
8621            let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
8622            let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
8623            for _ in 0..10 {
8624                let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
8625                let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
8626                let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
8627                c0 = n0;
8628                c1 = n1;
8629                c2 = n2;
8630                c3 = n3;
8631                k0 = k0.wrapping_add(0x9E3779B9);
8632                k1 = k1.wrapping_add(0xBB67AE85);
8633            }
8634            (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
8635        };
8636        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
8637        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
8638        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
8639        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
8640        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
8641        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
8642        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
8643        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
8644        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
8645        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
8646        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
8647        let t_ent = std::time::Instant::now();
8648
8649        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
8650        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
8651        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
8652        // the one that matters (a history-rewriting client mutates what the session GENERATED,
8653        // so the next turn's prompt agrees with this one up to exactly here).
8654        //
8655        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
8656        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
8657        // hold exactly `base + prompt.len()` rows and nothing generated.
8658        //
8659        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
8660        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
8661        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
8662        // `<think>` block the client strips, so every later turn's diff diverged exactly one
8663        // token below the checkpoint and affinity declined 100% of the time. Measured on the
8664        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
8665        // whole mechanism inert while looking, from the outside, like a working
8666        // correctness-declines-safely path — hence the decline log carries the offsets.
8667        //
8668        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
8669        // state (the reason a spec session could not rewind before). The draft scratch needs no
8670        // copy: rows below the boundary are rewritten by the next turn's own fill.
8671        //
8672        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
8673        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
8674        // checkpoint rather than replacing it with a strictly worse one.
8675        //
8676        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
8677        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
8678        // fail the burst that is already running — so the error is swallowed, loud only under
8679        // MEMRA_DEBUG_SPEC.
8680        if let Some(slot) = sess_ckpt_slot {
8681            if !continuation {
8682                let pos = cache.pos;
8683                debug_assert_eq!(
8684                    pos,
8685                    base + prompt.len(),
8686                    "turn checkpoint must sit at the prompt end, before the init feed"
8687                );
8688                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8689                    if let Some(ph) = &prompt_h {
8690                        // hidden of the LAST primed row = the predecessor anchor at this
8691                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
8692                        // last_h, and what the next prime's fill reads for its first row).
8693                        let np = prompt.len();
8694                        e.uninit(n_embd).and_then(|mut a| {
8695                            e.copy_view_into(
8696                                &mut a,
8697                                0,
8698                                &ph.slice((np - 1) * n_embd..np * n_embd),
8699                                n_embd,
8700                            )?;
8701                            Ok(a)
8702                        })
8703                    } else {
8704                        Err("no prompt hiddens".into())
8705                    };
8706                match (cache.snapshot(e), anchor) {
8707                    (Ok(snap), Ok(last_h)) => {
8708                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
8709                    }
8710                    (s, a) => {
8711                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
8712                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
8713                            let err = s
8714                                .err()
8715                                .map(|e| e.to_string())
8716                                .or_else(|| a.err().map(|e| e.to_string()))
8717                                .unwrap_or_default();
8718                            eprintln!(
8719                                "[spec] turn checkpoint skipped ({err}); \
8720                                       next turn re-primes in full"
8721                            );
8722                        }
8723                    }
8724                }
8725            }
8726        }
8727        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
8728        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
8729        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
8730        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
8731        let mut last_pred = 0u32;
8732        let mut last_col_logits: Option<CudaSlice<f32>> = None;
8733        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
8734        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
8735        let mut init_logits_host: Option<Vec<f32>> = None;
8736        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
8737            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
8738            last_pred = argmax(&init_logits) as u32;
8739            if constraint.is_some() {
8740                init_logits_host = Some(init_logits.clone());
8741            }
8742            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
8743            if sampled {
8744                last_col_logits = Some(e.htod(&init_logits)?);
8745            }
8746            h
8747        } else {
8748            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
8749            let lh = sess_tail
8750                .as_ref()
8751                .unwrap()
8752                .1
8753                .as_ref()
8754                .expect("pending carry requires last_h");
8755            e.clone_dtod(lh)?
8756        };
8757        let t_init = t_ent.elapsed();
8758        let mut last_col_stats: Option<(f32, f32, f32)> = None;
8759        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
8760        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
8761        // stable pointer for the graph-draft round-start copy.
8762        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
8763        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
8764        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
8765        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
8766        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
8767        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
8768        // overwritten below).
8769        let mut fill_prev = e.clone_dtod(&h_seed0)?;
8770        {
8771            if let Some(ph) = &prompt_h {
8772                let np = prompt.len();
8773                e.copy_view_into(
8774                    &mut h_seed_buf,
8775                    0,
8776                    &ph.slice((np - 1) * n_embd..np * n_embd),
8777                    n_embd,
8778                )?;
8779            } else if continuation {
8780                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
8781                    if let Some(lh) = lh.as_ref() {
8782                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
8783                    }
8784                }
8785            }
8786        }
8787        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
8788        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
8789
8790        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
8791        let fork_mode = OptiForkGateMode::configured();
8792        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
8793        // the end. Metric normalization vs the reference engine: BOTH engines count
8794        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
8795        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
8796        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
8797        let mut st_drafted = vec![0usize; k];
8798        let mut st_accepted = vec![0usize; k];
8799        let mut st_len_hist = vec![0usize; k + 1];
8800        let mut st_full = 0usize;
8801        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
8802        // stop the draft chain early when the head's softmax confidence in its own pick drops
8803        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
8804        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
8805        let p_min = *PMIN.get_or_init(|| {
8806            std::env::var("MEMRA_SPEC_PMIN")
8807                .ok()
8808                .and_then(|v| v.parse().ok())
8809                .unwrap_or(0.0)
8810        });
8811        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
8812        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
8813        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
8814        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
8815        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
8816        // verify batch is not); the j==0 exemption stays for pending-less rounds.
8817        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
8818            .map(|v| v == "1")
8819            .unwrap_or(false);
8820
8821        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
8822        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
8823        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
8824        // cuBLAS path in an exotic head) falls back to the eager draft chain.
8825        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
8826        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
8827        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
8828        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
8829        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
8830        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
8831        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
8832        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
8833        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
8834            Some(c) => c,
8835            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
8836        };
8837        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
8838        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
8839        if sampled && dctx.g_q.len() < d_vocab {
8840            dctx.g_q = e.zeros(d_vocab)?;
8841            dctx.g_perturb = e.zeros(d_vocab)?;
8842        }
8843        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
8844        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
8845        // truncation (the correctness backstop) stops cutting every tight-schema round.
8846        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
8847        // shape, so a parked graph of the other shape is dropped and recaptured.
8848        let dmask_on = constraint
8849            .as_deref()
8850            .is_some_and(|c| c.draft_mask_enabled());
8851        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
8852        if dmask_on && dctx.g_dmask.len() < dmask_words {
8853            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
8854            dctx.graph = None; // the old capture baked the old (or no) mask pointer
8855            dctx.failed.clear_greedy();
8856            dctx.keeper.clear();
8857        }
8858        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
8859            dctx.graph = None;
8860            dctx.failed.clear_greedy();
8861            dctx.keeper.clear();
8862        }
8863        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
8864            let DraftGraphCtx {
8865                g_tok,
8866                g_pos,
8867                g_seed,
8868                g_p,
8869                g_dmask,
8870                ..
8871            } = &mut dctx;
8872            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
8873            // host uploads the position's real words, so the warmups stay grammar-free.
8874            if dmask_on {
8875                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
8876            }
8877            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
8878            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
8879            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
8880            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
8881            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
8882            // passes (and, in serve, other sessions) recycle those addresses and the replay then
8883            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
8884            let cap_res = e.capture_graph_retained(|e| {
8885                self.mtp_head_forward_cap(
8886                    e,
8887                    mtp,
8888                    g_tok,
8889                    g_pos,
8890                    g_seed,
8891                    g_p,
8892                    &mut *scratch,
8893                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
8894                    true,
8895                    embd_gpu.expect("graph draft requires resident embedding"),
8896                    embd_qt,
8897                    embd_rb,
8898                    d_vocab,
8899                    None,
8900                    None,
8901                    if dmask_on {
8902                        Some((g_dmask_ro, dmask_words))
8903                    } else {
8904                        None
8905                    },
8906                )
8907            });
8908            match cap_res {
8909                Ok((g, keep)) => {
8910                    scratch.set_len(e, base)?;
8911                    dctx.graph = Some(g);
8912                    dctx.graph_masked = dmask_on;
8913                    dctx.keeper = keep;
8914                }
8915                Err(err) => {
8916                    scratch.set_len(e, base)?;
8917                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
8918                    // silent. Once per flip — mark returns None on an already-failed ctx.
8919                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
8920                        eprintln!("{line}");
8921                    }
8922                }
8923            }
8924        }
8925        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
8926        // graph object, built only when sampled && graph-eligible — the greedy capture above is
8927        // untouched (and skipped when sampled: its graph would never be launched). Same head
8928        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
8929        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
8930        // once per round); the raw head logits land in the persistent g_q for the host's
8931        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
8932        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
8933        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
8934        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
8935        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
8936        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
8937        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
8938        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
8939        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
8940        // this compare misses at most ONCE per resumed request — the first burst recaptures
8941        // and every later burst in that request replays. A client that wants the parked graph
8942        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
8943        // stable across its whole conversation.
8944        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
8945        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
8946        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
8947        // force the eager draft (which computes stats/penalties per row).
8948        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
8949        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
8950        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
8951        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
8952        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
8953        // the request shape the vendor-default flip makes the majority).
8954        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
8955        let pure_temp = s_key.pure_temp();
8956        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
8957            dctx.graph_s = None;
8958            dctx.failed.clear_sampled();
8959            dctx.s_key = None;
8960            dctx.q_slots.clear();
8961            dctx.keeper_s.clear();
8962        }
8963        if graph_draft
8964            && sampled
8965            && pure_temp
8966            && dctx.graph_s.is_none()
8967            && !dctx.failed.sampled_failed()
8968        {
8969            let DraftGraphCtx {
8970                g_tok,
8971                g_pos,
8972                g_seed,
8973                g_p,
8974                g_ctr,
8975                g_perturb,
8976                g_q,
8977                ..
8978            } = &mut dctx;
8979            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
8980            let cap_res = e.capture_graph_retained(|e| {
8981                self.mtp_head_forward_cap(
8982                    e,
8983                    mtp,
8984                    g_tok,
8985                    g_pos,
8986                    g_seed,
8987                    g_p,
8988                    &mut *scratch,
8989                    p_min > 0.0,
8990                    true,
8991                    embd_gpu.expect("graph draft requires resident embedding"),
8992                    embd_qt,
8993                    embd_rb,
8994                    d_vocab,
8995                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
8996                    None,
8997                    None, // constrained spec is greedy-only — sampled never carries a hook
8998                )
8999            });
9000            match cap_res {
9001                Ok((g, keep)) => {
9002                    scratch.set_len(e, base)?;
9003                    for _ in 0..k {
9004                        dctx.q_slots.push(e.zeros(d_vocab)?);
9005                    }
9006                    dctx.graph_s = Some(g);
9007                    dctx.s_key = Some(s_key);
9008                    dctx.keeper_s = keep;
9009                }
9010                Err(err) => {
9011                    scratch.set_len(e, base)?;
9012                    // LOUD flip (audit Q2): same contract as the greedy capture above.
9013                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
9014                        eprintln!("{line}");
9015                    }
9016                }
9017            }
9018        }
9019        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
9020        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
9021        // captured under this request's exact regime, and capture requires `pure_temp` — so a
9022        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
9023        // the graph arm, so it is asserted here rather than assumed: a future change that widens
9024        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
9025        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
9026        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
9027        // rather than launching it; the launch site re-tests `pure_temp` independently.
9028        if sampled && !pure_temp && dctx.graph_s.is_some() {
9029            debug_assert!(
9030                false,
9031                "sampled draft graph parked under {:?} survived into a FILTERED request \
9032                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
9033                 softmax, so the verify's filtered q would test a distribution the draft was \
9034                 never sampled from",
9035                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
9036            );
9037            eprintln!(
9038                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
9039                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
9040                 EAGER — the key must carry every field that shapes q",
9041                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
9042            );
9043            dctx.graph_s = None;
9044            dctx.s_key = None;
9045            dctx.q_slots.clear();
9046            dctx.keeper_s.clear();
9047        }
9048        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
9049        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
9050        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
9051        // arms below print which chain actually ran, so the probe never restates the condition.
9052        if skey_probe() {
9053            eprintln!(
9054                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
9055                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
9056                sampled as u8,
9057                pure_temp as u8,
9058                sp_temp,
9059                sp.top_k,
9060                sp.top_p,
9061                sp.min_p,
9062                pen_on as u8,
9063                k,
9064                graph_draft as u8,
9065                dctx.graph_s.is_some() as u8,
9066                dctx.s_key,
9067            );
9068        }
9069        let t_cap = t_ent.elapsed();
9070        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
9071        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
9072        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
9073        // fill: the first chain step processes it and appends its entry at slot prompt.len().
9074        if let Some(ph) = &prompt_h {
9075            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
9076            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
9077            // global positions [base..base+tp). Fresh call: base==0, identical to before.
9078            scratch.set_len(e, base)?;
9079            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
9080            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
9081            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
9082            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
9083            let tp = prompt.len();
9084            let fill_chunk: usize = if crate::cache::swa_ring_on() {
9085                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
9086            } else {
9087                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
9088                // meaning one monolithic fill.
9089                std::env::var("MEMRA_PRIME_CHUNK")
9090                    .ok()
9091                    .and_then(|v| v.parse().ok())
9092                    .unwrap_or(4096)
9093            };
9094            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
9095            let mut start = 0usize;
9096            while start < tp {
9097                let end = (start + fill_chunk).min(tp);
9098                let tc = end - start;
9099                {
9100                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
9101                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
9102                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
9103                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
9104                    let mut phs = e.zeros(tc * n_embd)?;
9105                    let (src_lo, dst_off) = if start == 0 {
9106                        (0, n_embd)
9107                    } else {
9108                        ((start - 1) * n_embd, 0)
9109                    };
9110                    let n_copy = if start == 0 {
9111                        (tc - 1) * n_embd
9112                    } else {
9113                        tc * n_embd
9114                    };
9115                    if start == 0 {
9116                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
9117                            if let Some(lh) = lh.as_ref() {
9118                                e.copy_into(&mut phs, 0, lh, n_embd)?;
9119                            }
9120                        }
9121                    }
9122                    if n_copy > 0 {
9123                        e.copy_view_into(
9124                            &mut phs,
9125                            dst_off,
9126                            &ph.slice(src_lo..src_lo + n_copy),
9127                            n_copy,
9128                        )?;
9129                    }
9130                    self.mtp_kv_fill(
9131                        e,
9132                        mtp,
9133                        &prompt[start..end],
9134                        &phs,
9135                        base + start,
9136                        &mut *scratch,
9137                        embd_dev,
9138                    )?;
9139                }
9140                start = end;
9141            }
9142        }
9143        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
9144        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
9145        // (=1 brackets the whole call in run_spec.rs, prime included.)
9146        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
9147            unsafe extern "C" {
9148                fn cudaProfilerStart() -> i32;
9149            }
9150            unsafe {
9151                cudaProfilerStart();
9152            }
9153        }
9154        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
9155        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
9156        // consume each other's device outputs; the host drains the ring every M rounds. v1
9157        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
9158        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
9159        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
9160        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
9161        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
9162        let stream_on = crate::spec::spec_stream()
9163            && !sampled
9164            && !spec_replay
9165            && constraint.is_none()
9166            && !session_mode
9167            && embd_gpu.is_some()
9168            && !crate::model::full_prec_enabled()
9169            && k + 2 < 96;
9170        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
9171        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
9172        if stream_on {
9173            let cap = e.capture_graph(|e| {
9174                for j in 0..k.max(1) {
9175                    self.mtp_head_forward_cap(
9176                        e,
9177                        mtp,
9178                        &mut dctx.g_tok,
9179                        &mut dctx.g_pos,
9180                        &mut dctx.g_seed,
9181                        &mut dctx.g_p,
9182                        &mut *scratch,
9183                        true,
9184                        true,
9185                        embd_gpu.expect("round stream requires resident embedding"),
9186                        embd_qt,
9187                        embd_rb,
9188                        d_vocab,
9189                        None,
9190                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
9191                        None, // round-stream requires constraint.is_none() (see stream_on)
9192                    )?;
9193                }
9194                Ok(())
9195            });
9196            match cap {
9197                Ok(g) => {
9198                    scratch.set_len(e, 0)?;
9199                    stream_graph = Some(g);
9200                }
9201                Err(err) => {
9202                    scratch.set_len(e, 0)?;
9203                    if debug_spec {
9204                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
9205                    }
9206                }
9207            }
9208        }
9209        let stream_active = stream_on && stream_graph.is_some();
9210        if debug_spec {
9211            eprintln!(
9212                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
9213                crate::spec::spec_stream(),
9214                dctx.graph.is_some(),
9215                stream_graph.is_some()
9216            );
9217        }
9218        let t_v_s = k + 1;
9219        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
9220        // module (extracted 2026-07-12; the gemma burst reuses them).
9221        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
9222        let crate::round_stream::StreamBufs {
9223            mut vtok_d,
9224            mut brk_d,
9225            mut pend_d,
9226            last_pred_d,
9227            mut pos_ctr,
9228            mut pos_start_d,
9229            mut ring_d,
9230            acc_d: mut stream_acc,
9231            m_rounds,
9232            k: _,
9233        } = sb;
9234        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
9235            Some(crate::round_stream::kv_len_ptr_table(
9236                e,
9237                cache,
9238                Some(&pos_ctr),
9239            )?)
9240        } else {
9241            None
9242        };
9243
9244        let t_fill = t_ent.elapsed();
9245        let mut round = 0usize;
9246        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
9247        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
9248        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
9249        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
9250        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
9251        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
9252        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
9253        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
9254        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
9255        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
9256        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
9257        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
9258        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
9259        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
9260        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
9261        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
9262        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
9263        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
9264        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
9265        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
9266        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
9267        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
9268        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
9269        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
9270        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
9271        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
9272        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
9273        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
9274        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
9275        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
9276            .ok()
9277            .and_then(|v| v.parse().ok());
9278        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
9279            4
9280        } else if self.cfg.n_embd as usize >= 2500 {
9281            2
9282        } else {
9283            1
9284        };
9285        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
9286        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
9287        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
9288        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
9289        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
9290            .ok()
9291            .and_then(|v| v.parse().ok())
9292            .unwrap_or(1024);
9293        let floor_at = |pos: usize| -> usize {
9294            if adapt_floor_env.is_some() || pos < floor_ctx {
9295                adapt_floor
9296            } else if adapt_floor >= 4 {
9297                1
9298            } else {
9299                adapt_floor
9300            }
9301        };
9302        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
9303        // fixed-K default path is untouched by this whole block.
9304        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
9305            .ok()
9306            .and_then(|v| v.parse().ok())
9307            .unwrap_or(7);
9308        let k_cap = k.min(cap_max).max(1);
9309        let mut kc = k_cap;
9310        let mut opti_fork: Option<OptiForkState> = None;
9311        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
9312        if fork_mode != OptiForkGateMode::Disabled {
9313            let fence = crate::pp::pp_cuts(self.layers.len());
9314            let refusal = if !session_mode {
9315                Some("not-session")
9316            } else if k != 1 || adapt {
9317                Some("requires-fixed-k1")
9318            } else if sampled || constraint.is_some() || spec_replay {
9319                Some("sampled-constrained-or-replay")
9320            } else if pipe.is_some() {
9321                Some("two-session-pipeline")
9322            } else if !spec_devacc() {
9323                Some("requires-device-accept")
9324            } else if stream_active || crate::spec::spec_stream() {
9325                Some("round-stream")
9326            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
9327                Some("swa-ring")
9328            } else if crate::pp::pp_host_bounce_active() {
9329                Some("host-bounce")
9330            } else if fork_mode == OptiForkGateMode::Controller
9331                && cache.recur.iter().any(Option::is_some)
9332            {
9333                Some("controller-requires-zero-recurrent-state")
9334            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
9335                Some("requires-pp2")
9336            } else {
9337                None
9338            };
9339            if let Some(reason) = refusal {
9340                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9341                eprintln!("[opti-fork] refused reason={reason}");
9342            } else {
9343                let fence = fence.expect("validated PP-2 fence");
9344                let rt = crate::pp::PpNRt::get(e)?;
9345                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
9346                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
9347                let primary_supported =
9348                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
9349                if !rt.cross_device() || !primary_supported {
9350                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9351                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
9352                } else {
9353                    // Both recurrent snapshots and both seed generations are allocated before
9354                    // the first fork, each through its owning PP stage. Allocation failure
9355                    // therefore happens before any optimistic state mutation can occur.
9356                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
9357                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
9358                    let fork = OptiForkState::new(
9359                        e,
9360                        cache,
9361                        fork_mode,
9362                        alternate_snapshot,
9363                        &h_seed_buf,
9364                        &fill_prev,
9365                        rt,
9366                        fence[1],
9367                        self.layers.len(),
9368                    )?;
9369                    eprintln!(
9370                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
9371                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
9372                        fence[1],
9373                        fork.logical_payload_bytes[0],
9374                        fork.logical_payload_bytes[1],
9375                        fork.controller.map_or(0.0, |policy| policy.threshold),
9376                    );
9377                    fork_snapshot = Some(current_snapshot);
9378                    opti_fork = Some(fork);
9379                }
9380            }
9381        }
9382        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
9383        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
9384        let mut snap = match fork_snapshot {
9385            Some(snapshot) => snapshot,
9386            None => cache.snapshot(e)?,
9387        };
9388        let mut carried_opti: Option<OptiControllerTicket> = None;
9389        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
9390        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
9391        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
9392            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
9393        } else {
9394            None
9395        };
9396        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
9397        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
9398        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
9399        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
9400        // pass of any kind). Verify still
9401        // checks every emitted token against the target -> exactness holds by construction; only
9402        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
9403        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
9404        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
9405        let mut pending: Option<u32> = carried_pending;
9406        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
9407        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
9408        // the verify accept readback). Printed once at loop end via spec-stats.
9409        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
9410        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
9411        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
9412        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
9413        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
9414        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
9415        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
9416        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
9417        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
9418        let mut ph_wait = 0f64;
9419        let mut ph_commit = 0f64;
9420        let mut ph_t = std::time::Instant::now();
9421        let mut ph_mark = |acc: &mut f64, on: bool| {
9422            if on {
9423                let now = std::time::Instant::now();
9424                *acc += (now - ph_t).as_secs_f64();
9425                ph_t = now;
9426            }
9427        };
9428        if let Some(p) = pipe {
9429            p.setup_end();
9430        }
9431        while keep_going && out.len() < max_new {
9432            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
9433            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
9434            if let (true, Some(sg), Some(ptrs)) = (
9435                stream_active && round >= 1 && pending.is_some(),
9436                &stream_graph,
9437                &stream_ptrs,
9438            ) {
9439                if debug_spec {
9440                    static ONCE: std::sync::Once = std::sync::Once::new();
9441                    ONCE.call_once(|| {
9442                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
9443                    });
9444                }
9445                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
9446                e.set_u32_one(&mut pend_d, pending.unwrap())?;
9447                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
9448                for _mi in 0..m_rounds {
9449                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
9450                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
9451                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
9452                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
9453                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
9454                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9455                    sg.launch()?;
9456                    e.spec_assemble_verify(
9457                        &g_tokp2k,
9458                        &pend_d,
9459                        d2t_dev.as_ref(),
9460                        &mut vtok_d,
9461                        &mut brk_d,
9462                        p_min,
9463                        k,
9464                        pmin0,
9465                    )?;
9466                    let mut ck = VerifyCkpt::new(self.layers.len());
9467                    let dummy = vec![0u32; t_v_s];
9468                    let (tl_d, vx) = self.decode_step_t_core_stream(
9469                        e,
9470                        &dummy,
9471                        0,
9472                        &mut *cache,
9473                        embd_dev,
9474                        Some(&mut ck),
9475                        Some((&vtok_d, &pos_ctr)),
9476                        None,
9477                        None,
9478                        None,
9479                    )?;
9480                    for j in 0..t_v_s {
9481                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
9482                    }
9483                    e.spec_accept_greedy_dc(
9484                        &preds_d,
9485                        &vtok_d,
9486                        &last_pred_d,
9487                        &brk_d,
9488                        &mut stream_acc,
9489                    )?;
9490                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
9491                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
9492                    self.commit_verified_prefix_stream(
9493                        e,
9494                        &mut *cache,
9495                        &snap,
9496                        &ck,
9497                        &stream_acc,
9498                        1,
9499                        t_v_s,
9500                    )?;
9501                    e.spec_rollback_stream(
9502                        ptrs,
9503                        &pos_start_d,
9504                        &stream_acc,
9505                        1,
9506                        self.layers.len() + 1,
9507                    )?;
9508                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
9509                }
9510                e.stream().synchronize()?;
9511                let ring_h = e.dtoh_u32(&ring_d)?;
9512                let cnt = ring_h[0] as usize;
9513                for i in 0..cnt {
9514                    if out.len() < max_new {
9515                        out.push(ring_h[1 + i]);
9516                    }
9517                }
9518                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
9519                for il in 0..self.layers.len() {
9520                    if let Some(kvl) = cache.kv[il].as_mut() {
9521                        kvl.len = pos_h;
9522                    }
9523                }
9524                cache.pos = pos_h;
9525                scratch.kv.len = pos_h;
9526                pending = Some(ring_h[cnt]); // last drained token = the live bonus
9527                last_token = ring_h[cnt];
9528                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
9529                total_accepted += cnt.saturating_sub(m_rounds);
9530                if let Some(t) = sess_telem {
9531                    // totals only — the burst's per-round accept counts stayed on device
9532                    // (that is the point of the round-stream arm). pos_* untouched.
9533                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
9534                }
9535                round += m_rounds;
9536                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
9537                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
9538                continue;
9539            }
9540            let pipe_draft = match pipe {
9541                Some(p) => Some(p.draft_begin(round)?),
9542                None => None,
9543            };
9544            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
9545            let mut current_opti = carried_opti.take();
9546            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
9547                match opti_fork.as_mut() {
9548                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
9549                    None => None,
9550                    Some(_) => None,
9551                }
9552            } else {
9553                None
9554            };
9555            if current_opti.is_none() {
9556                if let Some(fork) = opti_fork.as_ref() {
9557                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
9558                } else {
9559                    cache.snapshot_into(e, &mut snap)?;
9560                }
9561            } else if snap.pos != pos {
9562                return Err(format!(
9563                    "optipipe carried snapshot pos {} != current pos {pos}",
9564                    snap.pos
9565                )
9566                .into());
9567            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
9568            ph_mark(&mut ph_rest, phase_on);
9569
9570            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
9571            // p-min semantics (both paths): stop the chain early when the head's confidence in
9572            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
9573            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
9574            let base0 = if pending.is_some() { 1usize } else { 0usize };
9575            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
9576            // accepted run + 1 (the gemma law — see the setup block above the loop).
9577            let k_this = if adapt { kc } else { k };
9578            let mut draft: Vec<u32> = Vec::with_capacity(k);
9579            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
9580            let mut controller_draft_prob: Option<f32> = None;
9581            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
9582            if let Some(ticket) = current_opti.as_mut() {
9583                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
9584                if ticket.verify_tokens[0] != carried_pending {
9585                    return Err(format!(
9586                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
9587                        ticket.verify_tokens[0],
9588                    )
9589                    .into());
9590                }
9591                draft.push(ticket.verify_tokens[1]);
9592                controller_draft_prob = Some(ticket.draft_prob);
9593                controller_eager_state = ticket
9594                    .take_eager_seed()
9595                    .map(|seed| (ticket.verify_tokens[1], seed));
9596            } else {
9597                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
9598                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
9599                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
9600                // rejected drafts and p-min extras via the len mechanism).
9601                scratch.set_len(e, pos + base0 - 1)?;
9602                if pen_on {
9603                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
9604                    // device dedup: `penalty_last_n` is usize::MAX for any serve request with
9605                    // a penalty, so without the cap this grew with the whole session.
9606                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
9607                    let w0 = pen_hist.len().saturating_sub(win);
9608                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
9609                }
9610                if sampled {
9611                    draft_logits.clear();
9612                    draft_stats.clear();
9613                }
9614                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
9615                // position's mask is computed on that clone and advanced by the PROPOSED token. The
9616                // real state moves only on emission (verify's job), so the emitted stream is
9617                // unchanged — the mask only removes tokens the verify would have truncated anyway.
9618                let mut dmask_live = dmask_on;
9619                if dmask_live {
9620                    let t_c = std::time::Instant::now();
9621                    constraint
9622                        .as_deref_mut()
9623                        .unwrap()
9624                        .draft_begin()
9625                        .map_err(|e2| format!("constraint: {e2}"))?;
9626                    dm_clone_ns += t_c.elapsed().as_nanos();
9627                    dm_rounds += 1;
9628                }
9629                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
9630                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
9631                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
9632                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
9633                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
9634                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
9635                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9636                    for j in 0..k_this {
9637                        // per-position mask upload (contents only — the graph's baked pointer is
9638                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
9639                        // mask node degrades to a no-op ban instead of needing a second graph.
9640                        if dmask_live
9641                            && !upload_draft_mask(
9642                                e,
9643                                constraint.as_deref_mut().unwrap(),
9644                                &mut dctx.g_dmask,
9645                                mtp.d2t.as_ref(),
9646                                d_vocab,
9647                                dmask_words,
9648                            )?
9649                        {
9650                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
9651                            // genuinely miss the legal set): neutralize the captured mask node and
9652                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
9653                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
9654                            dmask_live = false;
9655                        }
9656                        gr.launch()?;
9657                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
9658                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
9659                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
9660                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
9661                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
9662                        // replay's embed node, and the MMU fault kills the CUDA context for the
9663                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
9664                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
9665                        // buffer (g_seed = the verify-side handoff vs head-side compute).
9666                        if (idx as usize) >= d_vocab {
9667                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
9668                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
9669                            // seed, untouched since the round-start copy — the pair discriminates
9670                            // "seed arrived poisoned" from "head forward produced NaN".
9671                            let seed_h = e.dtoh(&dctx.g_seed)?;
9672                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9673                            let in_h = e.dtoh(&h_seed_buf)?;
9674                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
9675                            return Err(format!(
9676                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9677                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
9678                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
9679                             the embed row (#87 trap)"
9680                            )
9681                            .into());
9682                        }
9683                        // trimmed draft vocab -> target token id (identity when no d2t map)
9684                        let d = match &mtp.d2t {
9685                            Some(map) => map[idx as usize],
9686                            None => idx,
9687                        };
9688                        let draft_p = if p_min > 0.0
9689                            || opti_fork
9690                                .as_ref()
9691                                .is_some_and(|fork| fork.controller.is_some())
9692                        {
9693                            Some(e.dtoh(&dctx.g_p)?[0])
9694                        } else {
9695                            None
9696                        };
9697                        if j == 0 {
9698                            controller_draft_prob = draft_p;
9699                        }
9700                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9701                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9702                                break;
9703                            }
9704                        }
9705                        draft.push(d);
9706                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
9707                        // index the argmax wrote — patch the persistent token buffer (4B htod).
9708                        if d != idx {
9709                            e.set_u32_one(&mut dctx.g_tok, d)?;
9710                        }
9711                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
9712                        // unmasked drafting for the remaining positions (verify still arbitrates).
9713                        // speculative advance; a chain the grammar can no longer follow (EOS
9714                        // proposed) ends here. The captured mask node always runs, so a dead chain
9715                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
9716                        if dmask_live
9717                            && !constraint
9718                                .as_deref_mut()
9719                                .unwrap()
9720                                .draft_advance(d)
9721                                .map_err(|e2| format!("constraint: {e2}"))?
9722                        {
9723                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
9724                            break;
9725                        }
9726                    }
9727                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
9728                // legal ONLY in the regime it was captured in. The condition used to read
9729                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
9730                // which it could not, because the key omitted the filters. Both halves are now
9731                // enforced: the key drops a stale graph, and this site refuses to launch one.
9732                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
9733                    if skey_probe() {
9734                        eprintln!(
9735                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
9736                             top_p={} min_p={} s_key_parked={:?}",
9737                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9738                        );
9739                    }
9740                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
9741                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
9742                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
9743                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
9744                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
9745                    // stream. Host sctr advances in lockstep (computed, no readback needed).
9746                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
9747                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
9748                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
9749                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
9750                    for j in 0..k_this {
9751                        gr.launch()?;
9752                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
9753                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
9754                        // counts the p-min-discarded token too)
9755                        // q retention: ONE async D2D of the persistent head-logits buffer into this
9756                        // round's slot j (stream-ordered after the replay, before the next one).
9757                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
9758                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
9759                        // #87 SENTINEL TRAP (see the greedy graph arm above).
9760                        if (idx as usize) >= d_vocab {
9761                            let seed_h = e.dtoh(&dctx.g_seed)?;
9762                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9763                            return Err(format!(
9764                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
9765                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
9766                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
9767                             (#87 trap)"
9768                            )
9769                            .into());
9770                        }
9771                        let d = match &mtp.d2t {
9772                            Some(map) => map[idx as usize],
9773                            None => idx,
9774                        };
9775                        draft_idx.push(idx);
9776                        if p_min > 0.0 {
9777                            let p = e.dtoh(&dctx.g_p)?[0];
9778                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9779                                break;
9780                            }
9781                        }
9782                        draft.push(d);
9783                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
9784                        if d != idx {
9785                            e.set_u32_one(&mut dctx.g_tok, d)?;
9786                        }
9787                    }
9788                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
9789                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
9790                    for j in 0..draft.len().max(draft_idx.len()) {
9791                        let rows0 = e.htod_i32(&[0])?;
9792                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9793                        e.filter_stats(
9794                            &dctx.q_slots[j],
9795                            d_vocab,
9796                            &rows0,
9797                            &mut th_d,
9798                            &mut z_d,
9799                            &mut mx_d,
9800                            d_vocab,
9801                            1,
9802                            sp_temp,
9803                            sp.top_k,
9804                            sp.top_p,
9805                            sp.min_p,
9806                        )?;
9807                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
9808                    }
9809                } else {
9810                    if skey_probe() && sampled {
9811                        eprintln!(
9812                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
9813                             top_p={} min_p={} s_key_parked={:?}",
9814                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
9815                        );
9816                    }
9817                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
9818                    let mut e_tok = last_token;
9819                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
9820                    for j in 0..k_this {
9821                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
9822                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
9823                        let mtp_pos = pos + base0 + j;
9824                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
9825                        // A position with no legal draft-vocab row drops to unmasked drafting for
9826                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
9827                        if dmask_live {
9828                            dmask_live = upload_draft_mask(
9829                                e,
9830                                constraint.as_deref_mut().unwrap(),
9831                                &mut dctx.g_dmask,
9832                                mtp.d2t.as_ref(),
9833                                d_vocab,
9834                                dmask_words,
9835                            )?;
9836                        }
9837                        let (dl_d, h_nextn) = self.mtp_head_forward_dev(
9838                            e,
9839                            mtp,
9840                            e_tok,
9841                            &d_seed,
9842                            &mut *scratch,
9843                            mtp_pos,
9844                            embd_dev,
9845                            if dmask_live {
9846                                Some((&dctx.g_dmask, dmask_words))
9847                            } else {
9848                                None
9849                            },
9850                        )?;
9851                        let tok_d = if sampled {
9852                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
9853                            // the filtered softmax (filters off => th=0, exact v1 semantics).
9854                            if perturb_buf.is_none() {
9855                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
9856                            }
9857                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
9858                            if pen_on {
9859                                let h = pen_hist_d.as_ref().unwrap();
9860                                let nh = h.len();
9861                                e.penalize_logits(
9862                                    &mut q_row,
9863                                    h,
9864                                    nh,
9865                                    sp.penalty_repeat,
9866                                    sp.penalty_freq,
9867                                    sp.penalty_present,
9868                                    d_vocab,
9869                                )?;
9870                            }
9871                            let rows0 = e.htod_i32(&[0])?;
9872                            let (mut th_d, mut z_d, mut mx_d) =
9873                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
9874                            e.filter_stats(
9875                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
9876                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
9877                            )?;
9878                            let (th, z, mx) =
9879                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
9880                            let pb = perturb_buf.as_mut().unwrap();
9881                            e.gumbel_perturb_filtered(
9882                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
9883                            )?;
9884                            sctr += 1;
9885                            draft_logits.push(q_row);
9886                            draft_stats.push((mx, th, z));
9887                            e.argmax_token_device(pb, d_vocab)?
9888                        } else {
9889                            e.argmax_token_device(&dl_d, d_vocab)?
9890                        };
9891                        let idx = e.dtoh_u32_one(&tok_d)?;
9892                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
9893                        // here because the eager chain's operands are all readable: dl_d (the head
9894                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
9895                        if (idx as usize) >= d_vocab {
9896                            let dl_h = e.dtoh(&dl_d)?;
9897                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
9898                            let seed_h = e.dtoh(&d_seed)?;
9899                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
9900                            return Err(format!(
9901                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
9902                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
9903                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
9904                             embed row (#87 trap)"
9905                            )
9906                            .into());
9907                        }
9908                        let d = match &mtp.d2t {
9909                            Some(map) => map[idx as usize],
9910                            None => idx,
9911                        };
9912                        if sampled {
9913                            draft_idx.push(idx);
9914                        }
9915                        let draft_p = if p_min > 0.0
9916                            || opti_fork
9917                                .as_ref()
9918                                .is_some_and(|fork| fork.controller.is_some())
9919                        {
9920                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
9921                            Some(e.dtoh(&p_d)?[0])
9922                        } else {
9923                            None
9924                        };
9925                        if j == 0 {
9926                            controller_draft_prob = draft_p;
9927                        }
9928                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
9929                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
9930                                break;
9931                            }
9932                        }
9933                        draft.push(d);
9934                        e_tok = d;
9935                        d_seed = h_nextn;
9936                        // speculative advance; a chain the grammar can no longer follow (EOS
9937                        // proposed) ends here — the prefix already proposed still rides verify.
9938                        if dmask_live
9939                            && !constraint
9940                                .as_deref_mut()
9941                                .unwrap()
9942                                .draft_advance(d)
9943                                .map_err(|e2| format!("constraint: {e2}"))?
9944                        {
9945                            break;
9946                        }
9947                    }
9948                    if opti_fork
9949                        .as_ref()
9950                        .is_some_and(|fork| fork.controller.is_some())
9951                    {
9952                        controller_eager_state = Some((e_tok, d_seed));
9953                    }
9954                }
9955            }
9956            let k_round = draft.len();
9957            if let Some(p) = pipe {
9958                p.draft_end(round);
9959            }
9960            drop(pipe_draft);
9961
9962            ph_mark(&mut ph_draft, phase_on);
9963            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
9964            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
9965            let verify_tokens: Vec<u32> = match pending {
9966                Some(b) => {
9967                    let mut v = Vec::with_capacity(k_round + 1);
9968                    v.push(b);
9969                    v.extend_from_slice(&draft);
9970                    v
9971                }
9972                None => draft.clone(),
9973            };
9974            let base = if pending.is_some() { 1 } else { 0 };
9975            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
9976            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
9977            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
9978                Some(ticket.take_ckpt())
9979            } else if spec_replay {
9980                None
9981            } else {
9982                Some(VerifyCkpt::new(self.layers.len()))
9983            };
9984            let controller_can_probe = base == 1
9985                && k_round == 1
9986                && out.len().saturating_add(2) < max_new
9987                && controller_draft_prob.is_some()
9988                && opti_fork
9989                    .as_ref()
9990                    .and_then(|fork| fork.controller.as_ref())
9991                    .is_some_and(|policy| !policy.breaker_tripped);
9992            let mut successor_attempt: Option<OptiControllerTicket> = None;
9993            let mut rejected_probe: Option<(f32, u32)> = None;
9994            let mut controller_prepared: Option<OptiControllerPrepared> = None;
9995            if controller_can_probe {
9996                // Prepare d2/q and, on admission, d3 before either current verify half is
9997                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
9998                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
9999                // the primary stream after N stage 1 would serialize the supposed pipeline.
10000                let eager_pos = scratch.kv.len + 1;
10001                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
10002                    e,
10003                    mtp,
10004                    &mut dctx,
10005                    &mut *scratch,
10006                    d_vocab,
10007                    &mut controller_eager_state,
10008                    eager_pos,
10009                    embd_dev,
10010                )?;
10011                let first_probability = controller_draft_prob
10012                    .ok_or("optipipe controller probe lost first-token probability")?;
10013                let q_proxy = first_probability * pending_probability;
10014                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10015                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10016                let admitted = opti_fork
10017                    .as_ref()
10018                    .and_then(|fork| fork.controller.as_ref())
10019                    .ok_or("optipipe controller policy disappeared")?
10020                    .admit(q_proxy);
10021                if admitted {
10022                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10023                    let eager_pos = scratch.kv.len + 1;
10024                    let (optimistic_draft, optimistic_draft_probability) = self
10025                        .opti_controller_draft_step(
10026                            e,
10027                            mtp,
10028                            &mut dctx,
10029                            &mut *scratch,
10030                            d_vocab,
10031                            &mut controller_eager_state,
10032                            eager_pos,
10033                            embd_dev,
10034                        )?;
10035                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10036                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
10037                        debug_assert_eq!(token, optimistic_draft);
10038                        seed
10039                    });
10040                    controller_prepared = Some(OptiControllerPrepared {
10041                        verify_tokens: [optimistic_pending, optimistic_draft],
10042                        draft_prob: optimistic_draft_probability,
10043                        eager_seed,
10044                        q_proxy,
10045                        scratch_len: scratch.kv.len,
10046                    });
10047                } else {
10048                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10049                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10050                    rejected_probe = Some((q_proxy, optimistic_pending));
10051                    eprintln!(
10052                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
10053                        opti_fork
10054                            .as_ref()
10055                            .and_then(|fork| fork.controller.as_ref())
10056                            .expect("controller policy")
10057                            .threshold,
10058                    );
10059                }
10060            }
10061            let fork_attempt = match fork_generation.take() {
10062                Some(generation) if base == 1 && k_round == 1 => Some(generation),
10063                Some(generation) => {
10064                    opti_fork
10065                        .as_mut()
10066                        .expect("fork generation without fork state")
10067                        .retire(generation)?;
10068                    None
10069                }
10070                None => None,
10071            };
10072            let (tlogits_d, vx) = if let Some(p) = pipe {
10073                self.decode_step_t_core_pipelined(
10074                    e,
10075                    &verify_tokens,
10076                    pos,
10077                    &mut *cache,
10078                    embd_dev,
10079                    ckpt.as_mut(),
10080                    p,
10081                    round,
10082                )?
10083            } else if controller_can_probe {
10084                let fence = opti_fork
10085                    .as_ref()
10086                    .ok_or("optipipe controller probe lost fork state")?
10087                    .fence;
10088                let boundary = match current_opti.as_mut() {
10089                    Some(ticket) => ticket.take_boundary(),
10090                    None => self.verify_stage0_issue(
10091                        e,
10092                        &verify_tokens,
10093                        pos,
10094                        &mut *cache,
10095                        embd_dev,
10096                        ckpt.as_mut(),
10097                        None,
10098                        &fence,
10099                        Some(true),
10100                        None,
10101                    )?,
10102                };
10103                if let Some(prepared) = controller_prepared.take() {
10104                    let generation = {
10105                        let fork = opti_fork
10106                            .as_mut()
10107                            .ok_or("optipipe controller admission lost fork state")?;
10108                        let generation = fork.reserve_successor()?;
10109                        let rt = fork.rt;
10110                        let snapshot_fence = fork.fence;
10111                        opti_snapshot_one_stage_owned_into(
10112                            e,
10113                            cache,
10114                            rt,
10115                            &snapshot_fence,
10116                            0,
10117                            fork.successor_snapshot_mut(),
10118                        )?;
10119                        generation
10120                    };
10121                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
10122                    let successor_boundary = self.verify_stage0_issue(
10123                        e,
10124                        &prepared.verify_tokens,
10125                        pos + verify_tokens.len(),
10126                        &mut *cache,
10127                        embd_dev,
10128                        Some(&mut successor_ckpt),
10129                        None,
10130                        &fence,
10131                        Some(false),
10132                        None,
10133                    )?;
10134                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10135                    let fork = opti_fork
10136                        .as_ref()
10137                        .ok_or("optipipe controller ticket lost fork state")?;
10138                    successor_attempt = Some(fork.controller_ticket(
10139                        generation,
10140                        successor_boundary,
10141                        successor_ckpt,
10142                        prepared.verify_tokens,
10143                        prepared.draft_prob,
10144                        prepared.eager_seed,
10145                        prepared.q_proxy,
10146                        prepared.scratch_len,
10147                    ));
10148                    eprintln!(
10149                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
10150                         verify={:?}",
10151                        generation.id,
10152                        prepared.q_proxy,
10153                        fork.controller.expect("controller policy").threshold,
10154                        prepared.verify_tokens,
10155                    );
10156                }
10157                let result = self.verify_stage1_finish(
10158                    e,
10159                    boundary,
10160                    &mut *cache,
10161                    ckpt.as_mut(),
10162                    None,
10163                    &fence,
10164                    successor_attempt.is_none(),
10165                )?;
10166                if let Some(ticket) = current_opti.as_mut() {
10167                    ticket.settle();
10168                }
10169                if successor_attempt.is_some() {
10170                    let fork = opti_fork
10171                        .as_mut()
10172                        .ok_or("optipipe successor snapshot lost fork state")?;
10173                    let rt = fork.rt;
10174                    let snapshot_fence = fork.fence;
10175                    opti_snapshot_one_stage_owned_into(
10176                        e,
10177                        cache,
10178                        rt,
10179                        &snapshot_fence,
10180                        1,
10181                        fork.successor_snapshot_mut(),
10182                    )?;
10183                    // Publish N only after both independent successor-state queues are complete.
10184                    fork.rt.publish_to(1, &e.stream())?;
10185                }
10186                result
10187            } else if let Some(ticket) = current_opti.as_mut() {
10188                let fork = opti_fork
10189                    .as_mut()
10190                    .ok_or("optipipe carried controller ticket lost fork state")?;
10191                let boundary = ticket.take_boundary();
10192                let result = self.verify_stage1_finish(
10193                    e,
10194                    boundary,
10195                    &mut *cache,
10196                    ckpt.as_mut(),
10197                    None,
10198                    &fork.fence,
10199                    true,
10200                )?;
10201                ticket.settle();
10202                result
10203            } else if let Some(generation) = fork_attempt {
10204                let fork = opti_fork
10205                    .as_mut()
10206                    .expect("fork generation without fork state");
10207                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
10208                let action = fork.mode.action(generation.id);
10209                let boundary = self.verify_stage0_issue(
10210                    e,
10211                    &verify_tokens,
10212                    pos,
10213                    &mut *cache,
10214                    embd_dev,
10215                    ckpt.as_mut(),
10216                    None,
10217                    &fork.fence,
10218                    Some(true),
10219                    None,
10220                )?;
10221                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10222                let mut ticket = fork.ticket(generation, boundary);
10223                if action == OptiForkAction::Abort {
10224                    return Err(format!(
10225                        "optipipe forced abort with generation {} stage0 in flight",
10226                        generation.id,
10227                    )
10228                    .into());
10229                }
10230                fork.reconcile(
10231                    e,
10232                    &mut *cache,
10233                    &mut *scratch,
10234                    &snap,
10235                    &mut h_seed_buf,
10236                    &mut fill_prev,
10237                    generation,
10238                    action,
10239                    verify_tokens[0],
10240                )?;
10241                let result = if action == OptiForkAction::Hit {
10242                    let boundary = ticket.take_boundary();
10243                    self.verify_stage1_finish(
10244                        e,
10245                        boundary,
10246                        &mut *cache,
10247                        ckpt.as_mut(),
10248                        None,
10249                        &fork.fence,
10250                        true,
10251                    )?
10252                } else {
10253                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
10254                    // verify only after E_restart published the restored stage-0 state.
10255                    self.decode_step_t_core(
10256                        e,
10257                        &verify_tokens,
10258                        pos,
10259                        &mut *cache,
10260                        embd_dev,
10261                        ckpt.as_mut(),
10262                    )?
10263                };
10264                ticket.settle();
10265                debug_assert_eq!(ticket.generation, generation);
10266                fork.retire(generation)?;
10267                result
10268            } else {
10269                self.decode_step_t_core(
10270                    e,
10271                    &verify_tokens,
10272                    pos,
10273                    &mut *cache,
10274                    embd_dev,
10275                    ckpt.as_mut(),
10276                )?
10277            };
10278            let pipe_accept = match pipe {
10279                Some(p) => Some(p.accept_begin(round)?),
10280                None => None,
10281            };
10282
10283            ph_mark(&mut ph_verify, phase_on);
10284            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
10285            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
10286            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
10287            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
10288            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
10289            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
10290            // (== the bonus), so every index shifts by `base` and last_pred is unused.
10291            let t_v = verify_tokens.len();
10292            let mut preds: Vec<u32> = Vec::new();
10293            if !sampled {
10294                for j in 0..t_v {
10295                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
10296                }
10297                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
10298                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
10299                // next round's last_token = the next chain's embed lookup. Catch it at the
10300                // source with the column named — an all-NaN VERIFY column implicates the
10301                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
10302                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
10303                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
10304                    let mut probe = e.zeros(n_vocab)?;
10305                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
10306                    let col_h = e.dtoh(&probe)?;
10307                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
10308                    return Err(format!(
10309                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
10310                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
10311                         — the stage-split verify produced a poisoned column (#87 trap)",
10312                        preds[bad]
10313                    )
10314                    .into());
10315                }
10316            }
10317            ph_mark(&mut ph_wait, phase_on);
10318            let t_pred = |j: usize| -> u32 {
10319                if j == 0 && base == 0 {
10320                    last_pred
10321                } else {
10322                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
10323                    // used to call this from the sampled arm and panicked the worker; it now goes
10324                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
10325                    // out-of-range pred is a real bug, not something to paper over.
10326                    debug_assert!(
10327                        !sampled,
10328                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
10329                    );
10330                    preds[base + j - 1]
10331                }
10332            };
10333            let mut devacc_seeded = false;
10334            let mut devacc_acc: Option<CudaSlice<u32>> = None;
10335            let (n_acc, bonus) = if !sampled {
10336                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
10337                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
10338                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
10339                // gated on token identity vs the host walk (the arms below are bit-equal rules).
10340                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
10341                {
10342                    let draft_d = e.htod_u32_v(&draft)?;
10343                    let mut acc_out = e.alloc_u32_zeroed(2)?;
10344                    e.spec_accept_greedy(
10345                        &preds_d,
10346                        &draft_d,
10347                        last_pred,
10348                        base,
10349                        k_round,
10350                        &mut acc_out,
10351                    )?;
10352                    devacc_acc = Some(acc_out.clone());
10353                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
10354                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
10355                    // non-replay commit arms skip their host-offset seed copies (guarded below);
10356                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
10357                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
10358                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
10359                    // the update lands after the arms (devacc_seeded guard below).
10360                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
10361                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
10362                    // unified rule; full accept rewrites the verify-left value). Host mirrors
10363                    // update after the readback; commit_verified_prefix skips its len_d writes.
10364                    if let Some(successor) = successor_attempt.as_ref() {
10365                        opti_fork
10366                            .as_mut()
10367                            .ok_or("optipipe successor reconcile lost fork state")?
10368                            .queue_actual_reconcile(
10369                                e,
10370                                &snap,
10371                                &acc_out,
10372                                successor.verify_tokens[0],
10373                                base,
10374                            )?;
10375                    } else if let Some(ptrs) = &kv_len_ptrs {
10376                        let saved: Vec<i32> = (0..self.layers.len())
10377                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
10378                            .collect();
10379                        let saved_d = e.htod_i32(&saved)?;
10380                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
10381                    }
10382                    devacc_seeded = true;
10383                    let ab = e.dtoh_u32(&acc_out)?;
10384                    (ab[0] as usize, ab[1])
10385                } else {
10386                    let mut n_acc = 0usize;
10387                    for j in 0..k_round {
10388                        if t_pred(j) == draft[j] {
10389                            n_acc += 1;
10390                        } else {
10391                            break;
10392                        }
10393                    }
10394                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
10395                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
10396                    (n_acc, t_pred(n_acc))
10397                }
10398            } else {
10399                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
10400                if col_buf.is_none() {
10401                    col_buf = Some(e.zeros(n_vocab)?);
10402                }
10403                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
10404                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
10405                let mut pj = vec![0f32; k_round.max(1)];
10406                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
10407                if k_round > 0 {
10408                    let mut ids: Vec<u32> = Vec::new();
10409                    let mut rows: Vec<i32> = Vec::new();
10410                    for j in 0..k_round {
10411                        if j > 0 || base == 1 {
10412                            ids.push(draft[j]);
10413                            rows.push((base + j) as i32 - 1);
10414                        }
10415                    }
10416                    if !ids.is_empty() {
10417                        let nr = rows.len();
10418                        // penalties: materialize the used columns into one contiguous penalized
10419                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
10420                        // penalties: materialize used columns contiguously, penalize all rows in
10421                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
10422                        let p_rows: Vec<i32> = if pen_on {
10423                            (0..nr as i32).collect()
10424                        } else {
10425                            rows.clone()
10426                        };
10427                        if pen_on {
10428                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
10429                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
10430                            }
10431                            let pc = pcol_buf.as_mut().unwrap();
10432                            for (i2, &r) in rows.iter().enumerate() {
10433                                let c = r as usize;
10434                                e.copy_view_into(
10435                                    pc,
10436                                    i2 * n_vocab,
10437                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
10438                                    n_vocab,
10439                                )?;
10440                            }
10441                            let h = pen_hist_d.as_ref().unwrap();
10442                            let nh = h.len();
10443                            e.penalize_logits_rows(
10444                                pc,
10445                                h,
10446                                nh,
10447                                sp.penalty_repeat,
10448                                sp.penalty_freq,
10449                                sp.penalty_present,
10450                                n_vocab,
10451                                nr,
10452                            )?;
10453                        }
10454                        let p_src: &CudaSlice<f32> = if pen_on {
10455                            pcol_buf.as_ref().unwrap()
10456                        } else {
10457                            &tlogits_d
10458                        };
10459                        let rowsd = e.htod_i32(&p_rows)?;
10460                        let (mut th_d, mut z_d, mut mx_d) =
10461                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
10462                        e.filter_stats(
10463                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
10464                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
10465                        )?;
10466                        let idsd = e.htod_u32_v(&ids)?;
10467                        let mut outd = e.zeros(nr)?;
10468                        e.softmax_gather_filtered(
10469                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
10470                            sp_temp,
10471                        )?;
10472                        let outv = e.dtoh(&outd)?;
10473                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
10474                        let mut oi = 0usize;
10475                        for j in 0..k_round {
10476                            if j > 0 || base == 1 {
10477                                pj[j] = outv[oi];
10478                                oi += 1;
10479                            }
10480                        }
10481                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
10482                    }
10483                    if base == 0 {
10484                        let lc: &CudaSlice<f32> = if pen_on {
10485                            if col_buf.is_none() {
10486                                col_buf = Some(e.zeros(n_vocab)?);
10487                            }
10488                            let cb = col_buf.as_mut().unwrap();
10489                            e.copy_into(
10490                                cb,
10491                                0,
10492                                last_col_logits
10493                                    .as_ref()
10494                                    .expect("sampled: last_col_logits unset"),
10495                                n_vocab,
10496                            )?;
10497                            let h = pen_hist_d.as_ref().unwrap();
10498                            let nh = h.len();
10499                            e.penalize_logits(
10500                                cb,
10501                                h,
10502                                nh,
10503                                sp.penalty_repeat,
10504                                sp.penalty_freq,
10505                                sp.penalty_present,
10506                                n_vocab,
10507                            )?;
10508                            col_buf.as_ref().unwrap()
10509                        } else {
10510                            last_col_logits
10511                                .as_ref()
10512                                .expect("sampled: last_col_logits unset")
10513                        };
10514                        let rows0 = e.htod_i32(&[0])?;
10515                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10516                        e.filter_stats(
10517                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
10518                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
10519                        )?;
10520                        let idsd = e.htod_u32_v(&[draft[0]])?;
10521                        let mut outd = e.zeros(1)?;
10522                        e.softmax_gather_filtered(
10523                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
10524                        )?;
10525                        pj[0] = e.dtoh(&outd)?[0];
10526                        last_col_stats =
10527                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
10528                    }
10529                }
10530                // q source: the graph arm retained the head logits in the persistent q_slots;
10531                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
10532                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
10533                // computes them post-replay — graph engages only filter/penalty-free, so the
10534                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
10535                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
10536                    &dctx.q_slots
10537                } else {
10538                    &draft_logits
10539                };
10540                let mut n_acc = 0usize;
10541                for j in 0..k_round {
10542                    let (qmx, qth, qz) = draft_stats[j];
10543                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
10544                    let rowsd = e.htod_i32(&[0])?;
10545                    let thd = e.htod(&[qth])?;
10546                    let zd = e.htod(&[qz])?;
10547                    let _ = qmx;
10548                    let mut outd = e.zeros(1)?;
10549                    e.softmax_gather_filtered(
10550                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
10551                        sp_temp,
10552                    )?;
10553                    let qj = e.dtoh(&outd)?[0];
10554                    let u = host_u01(sp_seed, uctr);
10555                    uctr += 1;
10556                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
10557                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
10558                    // exactness signature (see `skey_probe`). Impossible when the draft was
10559                    // drawn from the same filtered distribution the verify reconstructs here;
10560                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
10561                    if skey_probe() && qj == 0.0 {
10562                        eprintln!(
10563                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
10564                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
10565                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
10566                        );
10567                    }
10568                    if accept {
10569                        n_acc += 1;
10570                    } else {
10571                        break;
10572                    }
10573                }
10574                let bonus = if n_acc == k_round {
10575                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
10576                    let col = base + k_round - 1;
10577                    let cb = col_buf.as_mut().unwrap();
10578                    e.copy_view_into(
10579                        cb,
10580                        0,
10581                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
10582                        n_vocab,
10583                    )?;
10584                    if pen_on {
10585                        let h = pen_hist_d.as_ref().unwrap();
10586                        let nh = h.len();
10587                        e.penalize_logits(
10588                            cb,
10589                            h,
10590                            nh,
10591                            sp.penalty_repeat,
10592                            sp.penalty_freq,
10593                            sp.penalty_present,
10594                            n_vocab,
10595                        )?;
10596                    }
10597                    if perturb_buf.is_none() {
10598                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
10599                    }
10600                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
10601                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
10602                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
10603                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
10604                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
10605                    // last gathered column, in both base arms. `th` is a threshold in e-units of
10606                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
10607                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
10608                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
10609                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
10610                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
10611                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
10612                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
10613                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
10614                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
10615                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
10616                    // and row_max is unused once nothing is masked), so this fix is a byte-level
10617                    // no-op for the untruncated serve default. One extra one-block filter_stats
10618                    // per full-accept round is the whole cost.
10619                    let (mx, th) = {
10620                        let rows0 = e.htod_i32(&[0])?;
10621                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10622                        let cb0 = col_buf.as_ref().unwrap();
10623                        e.filter_stats(
10624                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
10625                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
10626                        )?;
10627                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
10628                    };
10629                    let pb = perturb_buf.as_mut().unwrap();
10630                    let cb2 = col_buf.as_ref().unwrap();
10631                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
10632                    sctr += 1;
10633                    let td = e.argmax_token_device(pb, n_vocab)?;
10634                    e.dtoh_u32_one(&td)?
10635                } else {
10636                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
10637                    let cb = col_buf.as_mut().unwrap();
10638                    if n_acc > 0 || base == 1 {
10639                        let col = base + n_acc - 1;
10640                        e.copy_view_into(
10641                            cb,
10642                            0,
10643                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
10644                            n_vocab,
10645                        )?;
10646                    } else {
10647                        let lc = last_col_logits.as_ref().unwrap();
10648                        e.copy_into(cb, 0, lc, n_vocab)?;
10649                    }
10650                    if pen_on {
10651                        let h = pen_hist_d.as_ref().unwrap();
10652                        let nh = h.len();
10653                        e.penalize_logits(
10654                            cb,
10655                            h,
10656                            nh,
10657                            sp.penalty_repeat,
10658                            sp.penalty_freq,
10659                            sp.penalty_present,
10660                            n_vocab,
10661                        )?;
10662                    }
10663                    let cb2 = col_buf.as_ref().unwrap();
10664                    let sc = sctr;
10665                    sctr += 1;
10666                    // p-stats for the reject column: from col_stats when the col was gathered,
10667                    // else (j==0&&base==0) from last_col_stats.
10668                    let p_stats = if n_acc > 0 || base == 1 {
10669                        // col index within the gathered set == number of gathered cols before n_acc
10670                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
10671                        col_stats.get(gi).copied().unwrap_or_else(|| {
10672                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
10673                        })
10674                    } else {
10675                        last_col_stats.expect("sampled: last_col_stats unset at reject")
10676                    };
10677                    let q_stats = draft_stats[n_acc];
10678                    if let Some(map) = &d2t_dev {
10679                        if q_full_buf.is_none() {
10680                            q_full_buf = Some(e.zeros(n_vocab)?);
10681                        }
10682                        let qf = q_full_buf.as_mut().unwrap();
10683                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
10684                        let qf2 = q_full_buf.as_ref().unwrap();
10685                        e.residual_sample_filtered(
10686                            cb2,
10687                            Some(qf2),
10688                            n_vocab,
10689                            sp_temp,
10690                            sp_seed,
10691                            sc,
10692                            p_stats,
10693                            q_stats,
10694                            &mut sample_tok,
10695                        )?;
10696                    } else {
10697                        e.residual_sample_filtered(
10698                            cb2,
10699                            Some(&q_bufs[n_acc]),
10700                            n_vocab,
10701                            sp_temp,
10702                            sp_seed,
10703                            sc,
10704                            p_stats,
10705                            q_stats,
10706                            &mut sample_tok,
10707                        )?;
10708                    }
10709                    e.dtoh_u32(&sample_tok)?[0]
10710                };
10711                (n_acc, bonus)
10712            };
10713            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
10714            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
10715            // ordering). Walk the accepted drafts through the grammar in commit order; the
10716            // first illegal token truncates acceptance at its slot, and that slot's emission
10717            // is recomputed as the MASKED argmax of the target's own verify column — token-
10718            // identical to constrained plain greedy decode (an unmasked argmax that is
10719            // grammar-legal IS the masked argmax: masking only removes competitors). The
10720            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
10721            // measured in acceptance numbers, never hidden.
10722            let (n_acc, bonus) = match constraint.as_deref_mut() {
10723                None => (n_acc, bonus),
10724                Some(c) => {
10725                    fn ce(e2: String) -> Box<dyn std::error::Error> {
10726                        format!("constraint: {e2}").into()
10727                    }
10728                    let mut na = n_acc;
10729                    let mut cut = false;
10730                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
10731                        if c.is_allowed(d).map_err(ce)? {
10732                            c.consume(d).map_err(ce)?;
10733                        } else {
10734                            na = j;
10735                            cut = true;
10736                            dm_cut_tokens += n_acc - j;
10737                            break;
10738                        }
10739                    }
10740                    if cut {
10741                        dm_cuts += 1;
10742                    }
10743                    let mut bo = bonus;
10744                    if cut || !c.is_allowed(bo).map_err(ce)? {
10745                        let mut row = if na == 0 && base == 0 {
10746                            init_logits_host
10747                                .clone()
10748                                .ok_or("constraint: init logits missing (round-0 cut)")?
10749                        } else {
10750                            e.dtoh_view(
10751                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
10752                            )?
10753                        };
10754                        c.mask_logits(&mut row).map_err(ce)?;
10755                        bo = argmax(&row) as u32;
10756                    }
10757                    c.consume(bo).map_err(ce)?;
10758                    (na, bo)
10759                }
10760            };
10761            let mut successor_valid = false;
10762            if let Some((q_proxy, expected_d2)) = rejected_probe {
10763                let v_n = n_acc == 1 && bonus == expected_d2;
10764                eprintln!(
10765                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
10766                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
10767                );
10768            }
10769            if let Some(successor) = successor_attempt.as_ref() {
10770                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
10771                let generation = successor.generation;
10772                let q_proxy = successor.q_proxy;
10773                let expected_pending = successor.verify_tokens[0];
10774                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
10775                let fork = opti_fork
10776                    .as_mut()
10777                    .ok_or("optipipe successor resolution lost fork state")?;
10778                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
10779                if successor_valid {
10780                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10781                } else {
10782                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10783                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10784                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
10785                }
10786                let breaker_tripped = fork
10787                    .controller
10788                    .as_mut()
10789                    .expect("controller policy")
10790                    .resolve(successor_valid);
10791                if breaker_tripped {
10792                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10793                }
10794                eprintln!(
10795                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
10796                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
10797                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
10798                    generation.id, successor_valid, !successor_valid, breaker_tripped,
10799                );
10800                if !successor_valid {
10801                    let mut successor = successor_attempt
10802                        .take()
10803                        .expect("controller successor disappeared on miss");
10804                    successor.settle();
10805                    fork.retire(generation)?;
10806                }
10807            }
10808            total_drafted += k_round;
10809            total_accepted += n_acc;
10810            if let Some(t) = sess_telem {
10811                // Greedy, rejection-sampling, and grammar truncation all converge here after
10812                // the accept decision is already on host. Fixed-size relaxed atomics only.
10813                t.record_round(k_round, n_acc);
10814            }
10815            if spec_stats {
10816                st_len_hist[k_round] += 1;
10817                for j in 0..k_round {
10818                    st_drafted[j] += 1;
10819                }
10820                for j in 0..n_acc {
10821                    st_accepted[j] += 1;
10822                }
10823                if n_acc == k_round {
10824                    st_full += 1;
10825                }
10826            }
10827
10828            if debug_spec {
10829                eprintln!(
10830                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
10831                    out.len(),
10832                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
10833                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
10834                    // the GPU worker thread — a debug flag that killed the exact regime you would
10835                    // set it to investigate. See `debug_t_pred0`.
10836                    debug_t_pred0(sampled, base, last_pred, &preds)
10837                );
10838            }
10839
10840            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
10841            let commit_started = std::time::Instant::now();
10842            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
10843            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
10844            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
10845            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
10846            for j in 0..n_acc {
10847                if !session_mode && out.len() >= max_new {
10848                    break;
10849                }
10850                out.push(draft[j]);
10851            }
10852            if pen_on {
10853                pen_hist.extend_from_slice(&draft[0..n_acc]);
10854                pen_hist.push(bonus);
10855            }
10856            let bonus_emitted = session_mode || out.len() < max_new;
10857            if bonus_emitted {
10858                out.push(bonus);
10859            }
10860            last_token = bonus;
10861
10862            // --- 5. ROLLBACK + advance (§C) ---
10863            if n_acc == k_round && !spec_replay {
10864                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
10865                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
10866                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
10867                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
10868                // last_pred is dead in the pending path (t_pred reads verify col 0).
10869                //
10870                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
10871                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
10872                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
10873                // trunk hidden (the last verify column). set_len first: a p-min break may have
10874                // left one extra chain append at that slot. Partial accepts need NO fill (the
10875                // chain already covered every accepted position; round-start set_len truncates).
10876                let mut vh_seed = e.zeros(n_embd)?;
10877                e.copy_view_into(
10878                    &mut vh_seed,
10879                    0,
10880                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
10881                    n_embd,
10882                )?;
10883                if refresh {
10884                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
10885                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
10886                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
10887                    // the full stack (vx) is already resident from the verify. Replaces both the
10888                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
10889                    // (draft attention quality); exactness stays the verify's job.
10890                    scratch.set_len(e, pos)?;
10891                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
10892                    // (hidden of the last committed row before this verify batch).
10893                    let mut vxs = e.zeros(t_v * n_embd)?;
10894                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10895                    if t_v > 1 {
10896                        e.copy_view_into(
10897                            &mut vxs,
10898                            n_embd,
10899                            &vx.slice(0..(t_v - 1) * n_embd),
10900                            (t_v - 1) * n_embd,
10901                        )?;
10902                    }
10903                    self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
10904                } else {
10905                    scratch.set_len(e, pos + base + k_round - 1)?;
10906                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
10907                    let mut hp = e.zeros(n_embd)?;
10908                    if t_v >= 2 {
10909                        e.copy_view_into(
10910                            &mut hp,
10911                            0,
10912                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
10913                            n_embd,
10914                        )?;
10915                    } else {
10916                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
10917                    }
10918                    self.mtp_kv_fill(
10919                        e,
10920                        mtp,
10921                        &[draft[k_round - 1]],
10922                        &hp,
10923                        pos + base + k_round - 1,
10924                        &mut *scratch,
10925                        embd_dev,
10926                    )?;
10927                }
10928                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
10929                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
10930                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
10931                // col). Saves one MTP-block pass per round on top of the pairing fix.
10932                if !devacc_seeded {
10933                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
10934                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
10935                }
10936                pending = Some(bonus);
10937                if debug_spec {
10938                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
10939                }
10940            } else if !spec_replay && base + n_acc >= 1 {
10941                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
10942                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
10943                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
10944                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
10945                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
10946                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
10947                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
10948                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
10949                // accept (never compounds: the next verify recomputes true hiddens for all
10950                // committed columns).
10951                let j = base + n_acc;
10952                self.commit_verified_prefix(
10953                    e,
10954                    &mut *cache,
10955                    &snap,
10956                    ckpt.as_ref().unwrap(),
10957                    j,
10958                    devacc_seeded,
10959                    if devacc_seeded {
10960                        devacc_acc.as_ref().map(|a| (a, base, t_v))
10961                    } else {
10962                        None
10963                    },
10964                )?;
10965                let mut seed = e.zeros(n_embd)?;
10966                e.copy_view_into(
10967                    &mut seed,
10968                    0,
10969                    &vx.slice((j - 1) * n_embd..j * n_embd),
10970                    n_embd,
10971                )?;
10972                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
10973                // branch); without it the chain entries stand and only the tail truncates. Either
10974                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
10975                // (persistent mode), rope pos+j+1 (chain convention).
10976                if refresh {
10977                    scratch.set_len(e, pos)?;
10978                    let mut vxs = e.zeros(j * n_embd)?;
10979                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
10980                    if j > 1 {
10981                        e.copy_view_into(
10982                            &mut vxs,
10983                            n_embd,
10984                            &vx.slice(0..(j - 1) * n_embd),
10985                            (j - 1) * n_embd,
10986                        )?;
10987                    }
10988                    self.mtp_kv_fill(
10989                        e,
10990                        mtp,
10991                        &verify_tokens[0..j],
10992                        &vxs,
10993                        pos,
10994                        &mut *scratch,
10995                        embd_dev,
10996                    )?;
10997                } else {
10998                    scratch.set_len(e, pos + j)?;
10999                }
11000                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
11001                // bonus's predecessor (verify col j-1); no pseudo pass.
11002                if !devacc_seeded {
11003                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
11004                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
11005                }
11006                pending = Some(bonus);
11007                if debug_spec {
11008                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
11009                }
11010            } else if !spec_replay {
11011                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
11012                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
11013                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
11014                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
11015                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
11016                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
11017                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
11018                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
11019                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
11020                cache.rollback(e, &snap, 0)?;
11021                scratch.set_len(e, pos)?;
11022                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
11023                pending = Some(bonus);
11024                if debug_spec {
11025                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
11026                }
11027            } else {
11028                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
11029                // this round survives, only possible before the first pending exists, ~round 0):
11030                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
11031                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
11032                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
11033                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
11034                // trunk hidden.
11035                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
11036                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
11037                if let Some(b) = pending.take() {
11038                    replay.push(b);
11039                }
11040                replay.extend_from_slice(&draft[0..n_acc]);
11041                replay.push(bonus);
11042                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
11043                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
11044                // last col exactly as before (byte-identical to the old _h_emb_dev call).
11045                let (rl_d, rx) = if self.qwen35_serving_class() {
11046                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
11047                    let mut hidden = e.uninit(replay.len() * n_embd)?;
11048                    for (row, &token) in replay.iter().enumerate() {
11049                        let (row_logits, row_hidden) =
11050                            self.spec_target_step_h(e, token, &mut *cache)?;
11051                        logits.extend_from_slice(&row_logits);
11052                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
11053                    }
11054                    (e.htod(&logits)?, hidden)
11055                } else {
11056                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
11057                };
11058                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
11059                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
11060                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
11061                last_pred = e.dtoh_u32(&preds_d)?[0];
11062                if sampled {
11063                    let lr0 = replay.len();
11064                    let lc = last_col_logits
11065                        .as_mut()
11066                        .expect("sampled: last_col_logits unset");
11067                    e.copy_view_into(
11068                        lc,
11069                        0,
11070                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
11071                        n_vocab,
11072                    )?;
11073                }
11074                let lr = replay.len();
11075                if lr >= 2 {
11076                    e.copy_view_into(
11077                        &mut h_seed_buf,
11078                        0,
11079                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
11080                        n_embd,
11081                    )?;
11082                } else {
11083                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
11084                    // last_token, whose own-row hidden fill_prev still holds.
11085                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
11086                }
11087                // the bonus is COMMITTED here — it becomes the last committed row.
11088                let mut rh_last = e.zeros(n_embd)?;
11089                e.copy_view_into(
11090                    &mut rh_last,
11091                    0,
11092                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
11093                    n_embd,
11094                )?;
11095                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
11096                if debug_spec {
11097                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
11098                }
11099            }
11100            if devacc_seeded {
11101                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
11102                // consumed the old value (both slots carry the same value in every non-replay arm).
11103                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11104            }
11105            if successor_valid {
11106                let optimistic_scratch_len = successor_attempt
11107                    .as_ref()
11108                    .expect("valid controller successor disappeared")
11109                    .scratch_len;
11110                // The normal current-round commit refreshed/truncated the logical scratch tail.
11111                // Its optimistic successor row was already written physically, so restoring only
11112                // the retained logical length makes that row live for the carried round.
11113                scratch.set_len(e, optimistic_scratch_len)?;
11114            }
11115            if let Some(current) = current_opti.take() {
11116                opti_fork
11117                    .as_mut()
11118                    .ok_or("optipipe current retirement lost fork state")?
11119                    .retire(current.generation)?;
11120            }
11121            if successor_valid {
11122                let successor = successor_attempt
11123                    .take()
11124                    .expect("valid controller successor disappeared before promotion");
11125                let generation = successor.generation;
11126                opti_fork
11127                    .as_mut()
11128                    .ok_or("optipipe successor promotion lost fork state")?
11129                    .promote_successor_snapshot(&mut snap, generation);
11130                carried_opti = Some(successor);
11131            }
11132            if anatomy_on {
11133                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
11134                // only for this diagnostic so it does not disappear into the following draft's
11135                // first token readback.
11136                e.stream().synchronize()?;
11137                ph_commit += commit_started.elapsed().as_secs_f64();
11138            }
11139            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
11140            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
11141            // final position — the floor's position key reads the committed depth). Burst
11142            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
11143            // like gemma's burst arm.
11144            if adapt {
11145                let fl_now = floor_at(cache.pos);
11146                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
11147            }
11148            ph_mark(&mut ph_rest, phase_on);
11149            if let Some(p) = pipe {
11150                p.accept_end(round);
11151            }
11152            drop(pipe_accept);
11153            round += 1;
11154            // sse-cadence: this round's accepted drafts + bonus are committed (out is
11155            // append-only past step 4) — flush at round cadence.
11156            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11157        }
11158        if let Some(mut ticket) = carried_opti.take() {
11159            opti_fork
11160                .as_mut()
11161                .ok_or("optipipe tail drain lost fork state")?
11162                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
11163        }
11164        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
11165        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
11166        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
11167
11168        if spec_stats {
11169            let per_slot: Vec<String> = (0..k)
11170                .map(|j| {
11171                    if st_drafted[j] > 0 {
11172                        format!(
11173                            "{}/{}={:.3}",
11174                            st_accepted[j],
11175                            st_drafted[j],
11176                            st_accepted[j] as f64 / st_drafted[j] as f64
11177                        )
11178                    } else {
11179                        "0/0".into()
11180                    }
11181                })
11182                .collect();
11183            let acc = if total_drafted > 0 {
11184                total_accepted as f64 / total_drafted as f64
11185            } else {
11186                0.0
11187            };
11188            eprintln!(
11189                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
11190                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
11191                       tok_per_round={:.3}",
11192                per_slot.join(" "),
11193                (total_accepted + round) as f64 / round.max(1) as f64
11194            );
11195        }
11196        if constraint.is_some() {
11197            eprintln!(
11198                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
11199                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
11200                dm_clone_ns as f64 / 1e6,
11201                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
11202            );
11203        }
11204        if phase_on {
11205            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
11206            eprintln!(
11207                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
11208                ph_draft * 1e3,
11209                ph_draft / tot * 100.0,
11210                ph_verify * 1e3,
11211                ph_verify / tot * 100.0,
11212                ph_wait * 1e3,
11213                ph_wait / tot * 100.0,
11214                ph_rest * 1e3,
11215                ph_rest / tot * 100.0
11216            );
11217        }
11218        if anatomy_on {
11219            let rounds_f = round.max(1) as f64;
11220            let other = (ph_rest - ph_commit).max(0.0);
11221            eprintln!(
11222                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
11223                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
11224                ph_draft * 1e3 / rounds_f,
11225                ph_verify * 1e3 / rounds_f,
11226                ph_wait * 1e3 / rounds_f,
11227                ph_commit * 1e3 / rounds_f,
11228                other * 1e3 / rounds_f,
11229            );
11230        }
11231        let _pipe_tail = pipe.map(|p| p.primary());
11232        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
11233        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
11234        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
11235        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
11236        if let Some(slot) = sess_draft_slot.take() {
11237            *slot = Some(dctx);
11238        }
11239        let t_rounds = t_ent.elapsed();
11240        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
11241            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
11242            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
11243            // HERE, where the sampler, the session Philox counters and the penalty window are
11244            // all live and the boundary logits row still exists — that is the "make the state
11245            // available" half of the fix; the consuming burst then just emits it. `sctr` is
11246            // written to the session BELOW the draws so the advance is never lost.
11247            *next_pred_slot = Some(last_pred);
11248            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
11249            let mut stashed_pending = false;
11250            if let Some(b) = pending.take() {
11251                if !sampled {
11252                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
11253                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
11254                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
11255                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
11256                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
11257                    // OUT of `committed` (cache rows == committed); the consuming call
11258                    // prepends it once its verify commits the row. next_pred is unknowable
11259                    // without the commit pass — None; callers gate on pending_tok too.
11260                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
11261                    if let Some(slot) = sess_pending_slot.take() {
11262                        *slot = Some(b);
11263                    }
11264                    *next_pred_slot = None;
11265                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
11266                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
11267                    *last_h = Some(e.clone_dtod(&fill_prev)?);
11268                    stashed_pending = true;
11269                } else {
11270                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
11271                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
11272                    let pos_b = cache.pos;
11273                    scratch.set_len(e, pos_b)?;
11274                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
11275                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
11276                    // itself — the prediction AFTER the bonus never materialized; it would have
11277                    // been the next round's verify col 0). The commit's logits ARE that
11278                    // prediction — so they are also the row the next burst's boundary token
11279                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
11280                    *next_pred_slot = Some(if sample_boundary {
11281                        sample_boundary_token(
11282                            e,
11283                            &lg_b,
11284                            &sp,
11285                            &pen_hist,
11286                            &mut sctr,
11287                            "burst-tail-commit",
11288                        )?
11289                    } else {
11290                        argmax(&lg_b) as u32
11291                    });
11292                    self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
11293                    *last_h = Some(hb);
11294                }
11295            } else {
11296                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
11297                *last_h = Some(e.clone_dtod(&fill_prev)?);
11298                if sample_boundary {
11299                    // No pending to commit, so the boundary row is the one `last_pred` was
11300                    // argmaxed from and the sampled path keeps it on device: the init feed's
11301                    // logits when the burst ran zero rounds, else the legacy-replay path's
11302                    // last verify column (both predict the token AFTER the last committed
11303                    // row). It is retained precisely because round 0's accept test needs it,
11304                    // so the draw costs no extra D2H of the [n_vocab] row.
11305                    match last_col_logits.as_ref() {
11306                        Some(lc) => {
11307                            *next_pred_slot = Some(sample_boundary_token_dev(
11308                                e,
11309                                lc,
11310                                n_vocab,
11311                                &sp,
11312                                &pen_hist,
11313                                &mut sctr,
11314                                "burst-tail-nopending",
11315                            )?);
11316                        }
11317                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
11318                        // burst always feeds or replays, so the row exists — but if it ever
11319                        // is, the stream takes a greedy token and SAYS so rather than
11320                        // silently regressing to the pre-lane behaviour.
11321                        None => eprintln!(
11322                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
11323                             (reason: no retained boundary logits row)"
11324                        ),
11325                    }
11326                }
11327            }
11328            *sctr_slot = sctr;
11329            *uctr_slot = uctr;
11330            committed.extend_from_slice(prompt);
11331            if let Some(cb) = carried_pending {
11332                // the consumed carry's cache row landed in round 0's verify (every pending
11333                // round commits col 0) — it joins `committed` here, in sequence order.
11334                committed.push(cb);
11335            }
11336            if stashed_pending {
11337                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
11338                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
11339                // 18446744073709551615 out of range for slice of length 0", killing the
11340                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
11341                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
11342                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
11343                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
11344                // did). So a burst that stashes a pending without emitting anything of its own —
11345                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
11346                // guard skipping every token under a tight budget — arrives here with
11347                // out.len() == 0 and stashed_pending == true.
11348                //
11349                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
11350                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
11351                // just above is already accounted. Saturating, not a min/assert: an empty `out`
11352                // here is a legitimate burst shape, not a corrupt state.
11353                let emitted = out.len().saturating_sub(1);
11354                committed.extend_from_slice(&out[..emitted]);
11355            } else {
11356                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
11357            }
11358            debug_assert_eq!(
11359                cache.pos,
11360                committed.len(),
11361                "session invariant: cache rows == committed tokens"
11362            );
11363            if setup_trace {
11364                e.stream().synchronize()?; // bound the async tail fill in the trace
11365                let t_tail = t_ent.elapsed();
11366                eprintln!(
11367                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
11368                    t_init.as_secs_f64() * 1e3,
11369                    (t_cap - t_init).as_secs_f64() * 1e3,
11370                    (t_fill - t_cap).as_secs_f64() * 1e3,
11371                    (t_rounds - t_fill).as_secs_f64() * 1e3,
11372                    (t_tail - t_rounds).as_secs_f64() * 1e3,
11373                    t_tail.as_secs_f64() * 1e3,
11374                    out.len(),
11375                    continuation
11376                );
11377            }
11378            return Ok((out, total_drafted, total_accepted));
11379        }
11380        out.truncate(max_new);
11381        Ok((out, total_drafted, total_accepted))
11382    }
11383
11384    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
11385    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
11386    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
11387    pub fn extract_dspark_anchors(
11388        &self,
11389        e: &Engine,
11390        tokens: &[u32],
11391        anchor_positions: &[usize],
11392        gamma: usize,
11393        top_k: usize,
11394        chunk: usize,
11395        temperature: f32,
11396    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
11397        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
11398            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
11399        }
11400        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
11401            return Err("DSpark anchor positions must be sorted and unique".into());
11402        }
11403        for &position in anchor_positions {
11404            if position == 0 || position + gamma >= tokens.len() {
11405                return Err(format!(
11406                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
11407                    tokens.len()
11408                )
11409                .into());
11410            }
11411        }
11412
11413        let n_vocab = self.output.out_features();
11414        let n_embd = self.cfg.n_embd as usize;
11415        let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
11416        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11417        let embd_gpu = if spec_host_embd() {
11418            None
11419        } else {
11420            Some(
11421                self.embd_gpu
11422                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11423            )
11424        };
11425        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
11426
11427        struct PendingRecord {
11428            position: usize,
11429            hidden: Option<Vec<f32>>,
11430            tokens: Vec<u32>,
11431            target_top_ids: Vec<Option<Vec<u32>>>,
11432            target_top_logits: Vec<Option<Vec<f32>>>,
11433            target_top_probs: Vec<Option<Vec<f32>>>,
11434            target_tail_probs: Vec<Option<f32>>,
11435        }
11436
11437        let mut pending: Vec<PendingRecord> = anchor_positions
11438            .iter()
11439            .map(|&position| PendingRecord {
11440                position,
11441                hidden: None,
11442                tokens: tokens[position..=position + gamma].to_vec(),
11443                target_top_ids: vec![None; gamma],
11444                target_top_logits: vec![None; gamma],
11445                target_top_probs: vec![None; gamma],
11446                target_tail_probs: vec![None; gamma],
11447            })
11448            .collect();
11449
11450        let mut start = 0usize;
11451        while start < tokens.len() {
11452            let end = (start + chunk).min(tokens.len());
11453            let chunk_tokens = &tokens[start..end];
11454            let (target_logits, hidden_rows) =
11455                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
11456            for record in &mut pending {
11457                let hidden_position = record.position - 1;
11458                if hidden_position >= start && hidden_position < end {
11459                    let local = hidden_position - start;
11460                    record.hidden = Some(
11461                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
11462                    );
11463                }
11464                for slot in 0..gamma {
11465                    let target_row = record.position + slot;
11466                    if target_row < start || target_row >= end {
11467                        continue;
11468                    }
11469                    let local = target_row - start;
11470                    let logits =
11471                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
11472                    let (ids, top_logits, probs, tail) =
11473                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
11474                    record.target_top_ids[slot] = Some(ids);
11475                    record.target_top_logits[slot] = Some(top_logits);
11476                    record.target_top_probs[slot] = Some(probs);
11477                    record.target_tail_probs[slot] = Some(tail);
11478                }
11479            }
11480            start = end;
11481        }
11482
11483        pending
11484            .into_iter()
11485            .map(|record| {
11486                let hidden = record
11487                    .hidden
11488                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
11489                let target_top_ids =
11490                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
11491                let target_top_logits = flatten_dspark_rows(
11492                    record.target_top_logits,
11493                    record.position,
11494                    "target logits",
11495                )?;
11496                let target_top_probs =
11497                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
11498                let target_tail_probs = record
11499                    .target_tail_probs
11500                    .into_iter()
11501                    .enumerate()
11502                    .map(|(slot, value)| {
11503                        value.ok_or_else(|| {
11504                            format!("missing DSpark tail at {} slot {slot}", record.position)
11505                        })
11506                    })
11507                    .collect::<Result<Vec<_>, _>>()?;
11508                Ok(DsparkAnchorRecord {
11509                    position: record.position,
11510                    hidden,
11511                    tokens: record.tokens,
11512                    target_top_ids,
11513                    target_top_logits,
11514                    target_top_probs,
11515                    target_tail_probs,
11516                })
11517            })
11518            .collect()
11519    }
11520
11521    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
11522    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
11523    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
11524    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
11525    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
11526    /// quant-induced head/hidden-state mismatch from text drift.
11527    ///
11528    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
11529    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
11530    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
11531    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
11532    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
11533    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
11534    ///              conditions on the corpus — deterministic and arm-comparable by design.
11535    ///
11536    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
11537    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
11538    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
11539    ///
11540    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
11541    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
11542    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
11543    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
11544    /// agreement vs this path — not usable as a training-data source).
11545    pub fn replay_acceptance(
11546        &self,
11547        e: &Engine,
11548        tokens: &[u32],
11549        k: usize,
11550        stride: usize,
11551        chunk: usize,
11552        mut hdump: Option<&mut std::fs::File>,
11553    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
11554        assert!(k >= 1 && stride >= 1 && chunk >= 2);
11555        let mtp = self
11556            .mtp
11557            .as_ref()
11558            .expect("replay_acceptance requires an MTP head");
11559        let n_vocab = self.output.out_features();
11560        let d_vocab = mtp
11561            .shared_head_head
11562            .as_ref()
11563            .unwrap_or(&self.output)
11564            .out_features();
11565        let n_embd = self.cfg.n_embd as usize;
11566        let t_total = tokens.len();
11567        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
11568        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
11569        let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
11570        let mut scratch = MtpScratch::new(
11571            e,
11572            &self.cfg,
11573            t_total + k + 8,
11574            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
11575        )?;
11576        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11577        let embd_gpu = if spec_host_embd() {
11578            None
11579        } else {
11580            Some(
11581                self.embd_gpu
11582                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11583            )
11584        };
11585        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11586
11587        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
11588        let mut bg: Vec<u32> = vec![0; t_total + 1];
11589        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
11590        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
11591        let mut seed_buf = e.zeros(n_embd)?;
11592        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
11593        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
11594        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
11595        let mut s = 0usize;
11596        while s < t_total {
11597            let cend = (s + chunk).min(t_total);
11598            let tc = cend - s;
11599            let ch = &tokens[s..cend];
11600            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
11601            //    the chunk's true hiddens.
11602            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
11603            for j in 0..tc {
11604                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11605            }
11606            let preds = e.dtoh_u32(&preds_d)?;
11607            for j in 0..tc {
11608                bg[s + j + 1] = preds[j];
11609            }
11610            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
11611            // checkpoint-quality metric (position j's logits score the GOLD next token).
11612            if nll_on {
11613                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
11614                if jmax > 0 {
11615                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
11616                    let rows: Vec<i32> = (0..jmax as i32).collect();
11617                    let idsd = e.htod_u32_v(&ids)?;
11618                    let rowsd = e.htod_i32(&rows)?;
11619                    let mut outd = e.zeros(jmax)?;
11620                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
11621                    for pr in e.dtoh(&outd)? {
11622                        nll_sum += -((pr.max(1e-30)) as f64).ln();
11623                        nll_cnt += 1;
11624                    }
11625                }
11626            }
11627            if let Some(f) = hdump.as_deref_mut() {
11628                use std::io::Write;
11629                let host: Vec<f32> = e.dtoh(&vx)?;
11630                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
11631                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
11632                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
11633                for v in &host[..tc * n_embd] {
11634                    let b = v.to_bits();
11635                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
11636                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
11637                }
11638                f.write_all(&bytes)?;
11639            }
11640            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
11641            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
11642            // per token saved; the forced trunk pass + hdump is all the mode needs).
11643            let chainless = stride > t_total;
11644            if chainless {
11645                e.copy_view_into(
11646                    &mut prev_last_h,
11647                    0,
11648                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
11649                    n_embd,
11650                )?;
11651                s = cend;
11652                continue;
11653            }
11654            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
11655            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
11656            let mut vxs = e.zeros(tc * n_embd)?;
11657            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
11658            if tc > 1 {
11659                e.copy_view_into(
11660                    &mut vxs,
11661                    n_embd,
11662                    &vx.slice(0..(tc - 1) * n_embd),
11663                    (tc - 1) * n_embd,
11664                )?;
11665            }
11666            scratch.set_len(e, s)?;
11667            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
11668            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
11669            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
11670            //    truncates those approximate appends before they can ever be read.
11671            let ps: Vec<usize> = (s..cend)
11672                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
11673                .collect();
11674            for &p in ps.iter().rev() {
11675                scratch.set_len(e, p)?;
11676                if p == s {
11677                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
11678                } else {
11679                    e.copy_view_into(
11680                        &mut seed_buf,
11681                        0,
11682                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
11683                        n_embd,
11684                    )?;
11685                }
11686                let mut e_tok = tokens[p];
11687                let mut d_seed = e.clone_dtod(&seed_buf)?;
11688                let mut drafts: Vec<u32> = Vec::with_capacity(k);
11689                for j in 0..k {
11690                    let (dl_d, h_nextn) = self.mtp_head_forward_dev(
11691                        e,
11692                        mtp,
11693                        e_tok,
11694                        &d_seed,
11695                        &mut scratch,
11696                        p + 1 + j,
11697                        embd_dev,
11698                        None, // acceptance-oracle walk: no grammar
11699                    )?;
11700                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
11701                    let idx = e.dtoh_u32_one(&tok_d)?;
11702                    let d = match &mtp.d2t {
11703                        Some(map) => map[idx as usize],
11704                        None => idx,
11705                    };
11706                    drafts.push(d);
11707                    e_tok = d;
11708                    d_seed = h_nextn;
11709                }
11710                // targets may live in a LATER chunk's bg — resolved after the walk.
11711                rows.push((p, drafts, Vec::new()));
11712            }
11713            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
11714            //    expect scratch.len == cend with exact rows).
11715            scratch.set_len(e, s)?;
11716            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
11717            e.copy_view_into(
11718                &mut prev_last_h,
11719                0,
11720                &vx.slice((tc - 1) * n_embd..tc * n_embd),
11721                n_embd,
11722            )?;
11723            s = cend;
11724        }
11725        for (p, drafts, targets) in rows.iter_mut() {
11726            for j in 0..drafts.len() {
11727                targets.push(bg[*p + 1 + j]);
11728            }
11729        }
11730        rows.sort_by_key(|r| r.0);
11731        if nll_cnt > 0 {
11732            let mean = nll_sum / nll_cnt as f64;
11733            println!(
11734                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
11735                mean.exp()
11736            );
11737        }
11738        Ok((rows, bg))
11739    }
11740}
11741
11742#[cfg(test)]
11743mod dspark_sparse_tests {
11744    use super::dspark_sparse_softmax_topk;
11745
11746    #[test]
11747    fn topk_keeps_full_softmax_mass_and_stable_ties() {
11748        let logits = [1.0f32, 3.0, 3.0, -2.0];
11749        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
11750        assert_eq!(ids, vec![1, 2]);
11751        assert_eq!(top_logits, vec![3.0, 3.0]);
11752        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
11753        let expected = 1.0 / denominator;
11754        assert!((probs[0] - expected).abs() < 1.0e-6);
11755        assert!((probs[1] - expected).abs() < 1.0e-6);
11756        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
11757        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
11758    }
11759}
11760
11761#[cfg(test)]
11762mod spec_replay_env_tests {
11763    use super::spec_replay_env_on;
11764
11765    #[test]
11766    fn replay_requires_literal_one() {
11767        assert!(!spec_replay_env_on(None));
11768        assert!(!spec_replay_env_on(Some("")));
11769        assert!(!spec_replay_env_on(Some("0")));
11770        assert!(!spec_replay_env_on(Some("true")));
11771        assert!(!spec_replay_env_on(Some("2")));
11772        assert!(spec_replay_env_on(Some("1")));
11773    }
11774}
11775
11776#[cfg(test)]
11777mod telem_tests {
11778    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
11779
11780    #[test]
11781    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
11782        let counters = SpecTelemetryCounters::default();
11783        for mask in [
11784            [true, true, true],
11785            [true, true, false],
11786            [true, false, false],
11787            [false, false, false],
11788        ] {
11789            let accepted = mask.iter().take_while(|&&value| value).count();
11790            counters.record_round(mask.len(), accepted);
11791        }
11792
11793        let snapshot = counters.snapshot();
11794        assert_eq!(
11795            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
11796            (4, 12, 6)
11797        );
11798        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
11799        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
11800        assert_eq!(snapshot.tau(), 1.5);
11801        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11802        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
11803    }
11804
11805    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
11806    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
11807    #[test]
11808    fn delta_isolates_burst_contribution() {
11809        let mut t = SpecTelemetry::default();
11810        // "previous request": 2 rounds of k=3, accepts 3 then 1.
11811        for (kr, na) in [(3usize, 3usize), (3, 1)] {
11812            t.rounds += 1;
11813            t.drafted += kr as u64;
11814            t.accepted += na as u64;
11815            for j in 0..kr {
11816                t.pos_drafted[j] += 1;
11817            }
11818            for j in 0..na {
11819                t.pos_accepted[j] += 1;
11820            }
11821        }
11822        let before = t;
11823        // "this burst": 1 round k=3, accepts 2.
11824        t.rounds += 1;
11825        t.drafted += 3;
11826        t.accepted += 2;
11827        for j in 0..3 {
11828            t.pos_drafted[j] += 1;
11829        }
11830        for j in 0..2 {
11831            t.pos_accepted[j] += 1;
11832        }
11833        let d = t.delta_since(&before);
11834        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
11835        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
11836        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
11837        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
11838    }
11839
11840    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
11841    /// aggregation invariant.
11842    #[test]
11843    fn merge_accumulates_fieldwise() {
11844        let mut agg = SpecTelemetry::default();
11845        let mut d1 = SpecTelemetry {
11846            rounds: 2,
11847            drafted: 6,
11848            accepted: 4,
11849            ..Default::default()
11850        };
11851        d1.pos_drafted[0] = 2;
11852        d1.pos_accepted[0] = 2;
11853        let mut d2 = SpecTelemetry {
11854            rounds: 1,
11855            drafted: 3,
11856            accepted: 1,
11857            ..Default::default()
11858        };
11859        d2.pos_drafted[0] = 1;
11860        d2.pos_accepted[0] = 1;
11861        d2.pos_drafted[1] = 1;
11862        agg.merge(&d1);
11863        agg.merge(&d2);
11864        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
11865        assert_eq!(agg.pos_drafted[0], 3);
11866        assert_eq!(agg.pos_accepted[0], 3);
11867        assert_eq!(agg.pos_drafted[1], 1);
11868        assert_eq!(agg.pos_accepted[1], 0);
11869    }
11870
11871    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
11872    /// public metrics surface and must never publish a u64-wrapped garbage value.
11873    #[test]
11874    fn delta_saturates_never_wraps() {
11875        let small = SpecTelemetry {
11876            rounds: 1,
11877            drafted: 2,
11878            accepted: 1,
11879            ..Default::default()
11880        };
11881        let big = SpecTelemetry {
11882            rounds: 5,
11883            drafted: 15,
11884            accepted: 9,
11885            ..Default::default()
11886        };
11887        let d = small.delta_since(&big);
11888        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
11889    }
11890}
11891
11892#[cfg(test)]
11893mod opti_fork_tests {
11894    use super::{
11895        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
11896    };
11897
11898    #[test]
11899    fn controller_threshold_and_three_miss_breaker_are_exact() {
11900        let mut policy = OptiControllerPolicy {
11901            threshold: 0.7,
11902            consecutive_misses: 0,
11903            breaker_tripped: false,
11904        };
11905        assert!(!policy.admit(0.699_999));
11906        assert!(policy.admit(0.7));
11907        assert!(!policy.resolve(false));
11908        assert!(!policy.resolve(false));
11909        assert!(policy.resolve(false));
11910        assert!(policy.breaker_tripped);
11911        assert!(!policy.admit(1.0));
11912        assert!(
11913            !policy.resolve(true),
11914            "a resolved hit cannot re-arm a tripped request"
11915        );
11916        assert!(policy.breaker_tripped);
11917    }
11918
11919    #[test]
11920    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
11921        let mut policy = OptiControllerPolicy {
11922            threshold: 0.0,
11923            consecutive_misses: 0,
11924            breaker_tripped: false,
11925        };
11926        for _ in 0..16 {
11927            assert!(policy.admit(0.0));
11928            assert!(!policy.resolve(false));
11929        }
11930        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
11931            assert!(
11932                !policy.admit(invalid),
11933                "invalid q proxy must fail closed: {invalid}"
11934            );
11935        }
11936        assert!(!policy.breaker_tripped);
11937        assert_eq!(policy.consecutive_misses, 0);
11938    }
11939
11940    #[test]
11941    fn alternating_mode_flips_by_generation_not_round_parity() {
11942        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
11943        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
11944        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
11945        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
11946    }
11947
11948    #[test]
11949    fn live_generation_cannot_be_overwritten() {
11950        let mut tracker = OptiForkGenerationTracker::default();
11951        let g0 = tracker.reserve().unwrap();
11952        let g1 = tracker.reserve().unwrap();
11953        let err = tracker.reserve().unwrap_err().to_string();
11954        assert!(
11955            err.contains("still owns generation 0"),
11956            "unexpected error: {err}"
11957        );
11958        tracker.retire(g0).unwrap();
11959        let g2 = tracker.reserve().unwrap();
11960        assert_eq!((g2.id, g2.slot), (2, 0));
11961        tracker.retire(g1).unwrap();
11962        tracker.retire(g2).unwrap();
11963    }
11964
11965    #[test]
11966    fn teardown_rejects_a_stale_generation_tag() {
11967        let mut tracker = OptiForkGenerationTracker::default();
11968        let g0 = tracker.reserve().unwrap();
11969        tracker.retire(g0).unwrap();
11970        let err = tracker.retire(g0).unwrap_err().to_string();
11971        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
11972    }
11973}
11974
11975#[cfg(test)]
11976mod draft_graph_fallback_tests {
11977    use super::DraftGraphFallback;
11978
11979    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
11980    #[test]
11981    fn flip_is_loud_once_and_memoized_after() {
11982        let mut f = DraftGraphFallback::default();
11983        let line = f
11984            .mark_greedy("out of memory")
11985            .expect("first flip must return the warn line");
11986        assert!(
11987            line.contains("WARN"),
11988            "flip line must be warn-level: {line}"
11989        );
11990        assert!(
11991            line.contains("out of memory"),
11992            "flip line must carry the reason: {line}"
11993        );
11994        assert!(f.greedy_failed());
11995        // re-marking an already-failed graph is the memoization: quiet, still failed.
11996        assert!(f.mark_greedy("out of memory").is_none());
11997        assert!(f.greedy_failed());
11998        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
11999        assert!(!f.sampled_failed());
12000        let line_s = f
12001            .mark_sampled("capture unsupported")
12002            .expect("sampled flip is its own flip");
12003        assert!(
12004            line_s.contains("sampled"),
12005            "sampled flip names itself: {line_s}"
12006        );
12007        assert!(f.mark_sampled("capture unsupported").is_none());
12008    }
12009
12010    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
12011    /// and says so exactly when there was something to reset.
12012    #[test]
12013    fn reset_on_resume_clears_flags_and_logs_once() {
12014        let mut f = DraftGraphFallback::default();
12015        // clean session: resume is silent, nothing to reset.
12016        assert!(f.reset_on_resume().is_none());
12017        f.mark_greedy("oom").unwrap();
12018        f.mark_sampled("oom").unwrap();
12019        let note = f
12020            .reset_on_resume()
12021            .expect("a set flag must produce the reset note");
12022        assert!(
12023            note.contains("greedy+sampled"),
12024            "note names what was reset: {note}"
12025        );
12026        assert!(
12027            !f.greedy_failed() && !f.sampled_failed(),
12028            "both flags cleared"
12029        );
12030        // and the NEXT failure after a reset is a fresh flip — loud again.
12031        assert!(f.mark_greedy("oom again").is_some());
12032        let note2 = f.reset_on_resume().expect("greedy-only reset");
12033        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
12034    }
12035
12036    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
12037    /// they precede a fresh capture attempt whose own failure re-flips loudly.
12038    #[test]
12039    fn shape_change_clears_are_silent() {
12040        let mut f = DraftGraphFallback::default();
12041        f.mark_greedy("oom").unwrap();
12042        f.clear_greedy();
12043        assert!(!f.greedy_failed());
12044        f.mark_sampled("oom").unwrap();
12045        f.clear_sampled();
12046        assert!(!f.sampled_failed());
12047        // after a silent clear there is nothing left for resume to report.
12048        assert!(f.reset_on_resume().is_none());
12049    }
12050}
12051
12052/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
12053///
12054/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
12055/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
12056/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
12057/// than remembered.
12058#[cfg(test)]
12059mod sampled_graph_key_tests {
12060    use super::{SampledGraphKey, debug_t_pred0};
12061
12062    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
12063    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
12064        (k.seed, k.temp_bits, k.k)
12065    }
12066
12067    fn pure_temp_key() -> SampledGraphKey {
12068        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
12069        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
12070    }
12071
12072    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
12073    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
12074    #[test]
12075    fn vendor_filters_change_the_key() {
12076        let parked = pure_temp_key();
12077        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
12078        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
12079        assert_eq!(
12080            legacy_key(&parked),
12081            legacy_key(&vendor),
12082            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
12083        );
12084        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
12085        assert!(parked.pure_temp());
12086        assert!(!vendor.pure_temp());
12087    }
12088
12089    /// Each distribution-shaping field alone is enough to drop the parked graph.
12090    #[test]
12091    fn every_filter_field_is_keyed() {
12092        let base = pure_temp_key();
12093        for (what, other) in [
12094            (
12095                "top_k",
12096                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
12097            ),
12098            (
12099                "top_p",
12100                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
12101            ),
12102            (
12103                "min_p",
12104                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
12105            ),
12106            (
12107                "penalties",
12108                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
12109            ),
12110        ] {
12111            assert_ne!(base, other, "{what} must be part of the key");
12112            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
12113            assert_eq!(
12114                legacy_key(&base),
12115                legacy_key(&other),
12116                "{what} was invisible to the pre-fix key",
12117            );
12118        }
12119    }
12120
12121    /// The baked constants stay keyed (this half was always right — regression cover for it).
12122    #[test]
12123    fn baked_constants_stay_keyed() {
12124        let base = pure_temp_key();
12125        assert_ne!(
12126            base,
12127            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
12128            "seed"
12129        );
12130        assert_ne!(
12131            base,
12132            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
12133            "temp"
12134        );
12135        assert_ne!(
12136            base,
12137            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
12138            "k"
12139        );
12140        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
12141        assert_eq!(
12142            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
12143            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
12144        );
12145    }
12146
12147    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
12148    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
12149    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
12150    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
12151    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
12152    ///
12153    /// This test is the other end of that argument, asserted here rather than remembered in a
12154    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
12155    /// would silently become the unsound thing it is documented not to be.
12156    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
12157    #[test]
12158    fn seed_alone_still_rekeys_the_draft_graph() {
12159        let parked = pure_temp_key();
12160        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
12161        assert_ne!(
12162            parked, reseeded,
12163            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
12164             decision not to compare seed rests on exactly this",
12165        );
12166        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
12167        // because of a filter difference.
12168        assert!(parked.pure_temp() && reseeded.pure_temp());
12169    }
12170
12171    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
12172    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
12173    /// agree on the regime, so a graph that survives the drop is legal to launch.
12174    #[test]
12175    fn equal_keys_agree_on_the_regime() {
12176        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
12177        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
12178        assert_eq!(a, b);
12179        assert_eq!(a.pure_temp(), b.pure_temp());
12180        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
12181        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
12182        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
12183        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
12184    }
12185
12186    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
12187    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
12188    #[test]
12189    fn debug_print_survives_the_sampled_arm() {
12190        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
12191        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
12192        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
12193        // round 0 without a pending bonus still reports last_pred, in both arms.
12194        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
12195        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
12196        // greedy keeps the real prediction it always printed.
12197        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
12198        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
12199    }
12200}