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 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
212/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
213/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
214/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
215/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
216///
217/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
218/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
219/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
220/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
221/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
222/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
223/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
224/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
225/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
226/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
227/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
228/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
229/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
230/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
231/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
232/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
233/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
234/// ratification on the serve-surface battery.
235pub(crate) fn dspark_verify_graph_on() -> bool {
236    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
237    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
238}
239/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
240/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
241/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
242/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
243/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
244/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
245/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
246/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
247/// 256-token run vs the serve session's thousands of rounds), and the two
248/// instruments must keep their own measured dispositions rather than share one flag.
249pub(crate) fn dspark_verify_graph_serve_on() -> bool {
250    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
251    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
252}
253/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
254/// pool's memory policy STATED instead of silently unbounded. The keyspace is
255/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
256/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
257/// on the q38 export — so the default (256) never engages there; the knob is the
258/// safety valve for a future export with a wider ladder. At the ceiling the pool
259/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
260/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
261/// cols-stashed layers inside one commit). No eviction by design: destroying a live
262/// exec graph re-opens the stale-address class the indirect tables exist to close,
263/// and the bounded keyspace makes reclaim worthless.
264pub(crate) fn dspark_vg_cap() -> usize {
265    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
266    *CAP.get_or_init(|| {
267        std::env::var("MEMRA_DSPARK_VG_MAX")
268            .ok()
269            .and_then(|v| v.parse().ok())
270            .unwrap_or(256)
271    })
272}
273/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
274/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
275/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
276/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
277/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
278/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
279/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
280/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
281/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
282/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
283/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
284/// empty partial the combine never reads, so the shared n_splits_max stride changes no
285/// bytes) and re-gated e2e by this lane's battery.
286pub(crate) fn dspark_fa_rows_on() -> bool {
287    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
288    *ON.get_or_init(|| {
289        std::env::var("MEMRA_DSPARK_FA_ROWS")
290            .map(|v| v != "0")
291            .unwrap_or(true)
292    })
293}
294
295/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
296///
297/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
298/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
299/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
300/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
301/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
302/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
303/// the flag crashed precisely the regime it exists to investigate.
304///
305/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
306/// indexing (an out-of-range pred there is a real bug and must still be loud).
307fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
308    if base == 0 {
309        return last_pred.to_string();
310    }
311    match preds.get(base - 1) {
312        Some(p) => p.to_string(),
313        // sampled: the greedy per-column argmax was never run for this round.
314        None => {
315            debug_assert!(
316                sampled,
317                "greedy spec: preds[{}] missing at base {base}",
318                base - 1
319            );
320            "n/a".to_string()
321        }
322    }
323}
324
325/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
326///
327/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
328/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
329/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
330/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
331/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
332/// not believe in — and `u * 0 < p` then accepts it unconditionally.
333///
334/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
335/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
336pub(crate) fn skey_probe() -> bool {
337    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
338    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
339}
340
341/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
342/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
343/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
344/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
345/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
346/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
347/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
348/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
349/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
350pub trait SpecConstraint {
351    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
352    /// masked argmax).
353    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
354    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
355    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
356    /// Is `tok` consumable in the CURRENT state?
357    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
358    /// Advance the state with an emitted token.
359    fn consume(&mut self, tok: u32) -> Result<(), String>;
360
361    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
362    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
363    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
364    // loose, research/constrained-full-20260803). These three methods let the engine mask the
365    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
366    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
367    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
368    // stays the correctness backstop and the emitted stream is unchanged by construction
369    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
370    // argmax; a cut slot is recomputed as the masked argmax either way).
371    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
372
373    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
374    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
375    fn draft_mask_enabled(&self) -> bool {
376        false
377    }
378    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
379    /// slot. Called once per spec round, before the first draft position.
380    fn draft_begin(&mut self) -> Result<(), String> {
381        Ok(())
382    }
383    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
384    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
385    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
386        Ok(None)
387    }
388    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
389    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
390    /// engine stops drafting; the token already pushed still goes through verify.
391    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
392        Ok(false)
393    }
394}
395
396/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
397/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
398/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
399/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
400/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
401/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
402/// verify emits the masked argmax as usual).
403fn upload_draft_mask(
404    e: &Engine,
405    c: &mut dyn SpecConstraint,
406    dst: &mut CudaSlice<u32>,
407    d2t: Option<&Vec<u32>>,
408    d_vocab: usize,
409    words: usize,
410) -> Result<bool, Box<dyn std::error::Error>> {
411    let Some(tw) = c
412        .draft_mask_words()
413        .map_err(|e2| format!("constraint: {e2}"))?
414    else {
415        return Ok(false);
416    };
417    let bit = |t: usize| -> bool {
418        let w = t >> 5;
419        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
420    };
421    let mut buf = vec![0u32; words];
422    match d2t {
423        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
424        Some(map) => {
425            for (i, &t) in map.iter().enumerate().take(d_vocab) {
426                if bit(t as usize) {
427                    buf[i >> 5] |= 1u32 << (i & 31);
428                }
429            }
430        }
431        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
432        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
433        None => {
434            let n = tw.len().min(words);
435            buf[..n].copy_from_slice(&tw[..n]);
436        }
437    }
438    if buf.iter().all(|w| *w == 0) {
439        return Ok(false);
440    }
441    e.htod_u32_into(dst, &buf)?;
442    Ok(true)
443}
444
445/// Keep the full token-embedding table in host memory and upload only the rows needed by each
446/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
447/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
448/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
449pub(crate) fn spec_host_embd() -> bool {
450    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
451    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
452}
453
454/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
455/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
456/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
457/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
458/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
459/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
460/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
461/// run-spec K=1..8 + acceptance identity arbitrate e2e).
462pub(crate) fn spec_fused_t() -> bool {
463    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
464    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
465    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
466    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
467    *F.get_or_init(|| {
468        std::env::var("MEMRA_SPEC_FUSED_T")
469            .map(|v| v != "0")
470            .unwrap_or(true)
471    })
472}
473
474/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
475/// Only call this on such buffers — the lean contract is "identical bytes by construction".
476fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
477    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
478}
479
480/// Scratch KV for the MTP block (one full-attn layer).
481///
482/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
483/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
484/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
485/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
486/// engine's "mtp_update" design). Entries come from two sources:
487///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
488///     hidden chain-approximate — the reference engine accepts the same);
489///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
490///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
491/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
492/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
493/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
494/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
495/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
496/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
497/// committed row across turns (the predecessor-pairing seed + fill anchor).
498/// Per-request sampling config for the sampled-spec serve path.
499#[derive(Clone, Copy, Debug)]
500pub struct SpecSampling {
501    pub temp: f32,
502    pub seed: u64,
503    pub top_k: i32,            // 0 = off
504    pub top_p: f32,            // 1.0 = off
505    pub min_p: f32,            // 0.0 = off
506    pub penalty_last_n: usize, // 0 = penalties off
507    pub penalty_repeat: f32,
508    pub penalty_freq: f32,
509    pub penalty_present: f32,
510}
511
512impl SpecSampling {
513    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
514    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
515    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
516    /// key their penalty arms off this.
517    pub fn pen_on(&self) -> bool {
518        self.penalty_last_n > 0
519            && (self.penalty_repeat != 1.0
520                || self.penalty_freq != 0.0
521                || self.penalty_present != 0.0)
522    }
523}
524
525/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
526/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
527/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
528/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
529/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
530/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
531/// is a distributional bug, not a style problem).
532pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
533    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
534    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
535    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
536    for _ in 0..10 {
537        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
538        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
539        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
540        c0 = n0;
541        c1 = n1;
542        c2 = n2;
543        c3 = n3;
544        k0 = k0.wrapping_add(0x9E3779B9);
545        k1 = k1.wrapping_add(0xBB67AE85);
546    }
547    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
548}
549
550/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
551/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
552pub const SPEC_TELEM_POS: usize = 8;
553
554/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
555/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
556/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
557/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
558/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
559/// in NEITHER drafted nor accepted.
560#[derive(Clone, Copy, Default, Debug)]
561pub struct SpecTelemetry {
562    /// verify rounds completed (a round-stream burst counts each of its M rounds).
563    pub rounds: u64,
564    /// tokens drafted / accepted across all rounds.
565    pub drafted: u64,
566    pub accepted: u64,
567    /// how often draft position j (0-based within a round's chain) was offered / accepted.
568    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
569    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
570    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
571    pub pos_drafted: [u64; SPEC_TELEM_POS],
572    pub pos_accepted: [u64; SPEC_TELEM_POS],
573}
574
575impl SpecTelemetry {
576    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
577    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
578    /// a wrapped counter.
579    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
580        let mut d = SpecTelemetry {
581            rounds: self.rounds.saturating_sub(prev.rounds),
582            drafted: self.drafted.saturating_sub(prev.drafted),
583            accepted: self.accepted.saturating_sub(prev.accepted),
584            ..Default::default()
585        };
586        for j in 0..SPEC_TELEM_POS {
587            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
588            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
589        }
590        d
591    }
592    /// Fieldwise `self += d` — the worker's per-model aggregation.
593    pub fn merge(&mut self, d: &SpecTelemetry) {
594        self.rounds += d.rounds;
595        self.drafted += d.drafted;
596        self.accepted += d.accepted;
597        for j in 0..SPEC_TELEM_POS {
598            self.pos_drafted[j] += d.pos_drafted[j];
599            self.pos_accepted[j] += d.pos_accepted[j];
600        }
601    }
602
603    /// Mean accepted draft-prefix length per verify round (tau).
604    pub fn tau(&self) -> f64 {
605        if self.rounds > 0 {
606            self.accepted as f64 / self.rounds as f64
607        } else {
608            0.0
609        }
610    }
611}
612
613/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
614/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
615/// launch, synchronization, allocation, or ordering dependency to the numeric path.
616struct SpecTelemetryCounters {
617    rounds: AtomicU64,
618    drafted: AtomicU64,
619    accepted: AtomicU64,
620    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
621    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
622}
623
624impl Default for SpecTelemetryCounters {
625    fn default() -> Self {
626        Self {
627            rounds: AtomicU64::new(0),
628            drafted: AtomicU64::new(0),
629            accepted: AtomicU64::new(0),
630            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
631            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
632        }
633    }
634}
635
636impl SpecTelemetryCounters {
637    fn record_round(&self, drafted: usize, accepted: usize) {
638        debug_assert!(accepted <= drafted);
639        self.rounds.fetch_add(1, Ordering::Relaxed);
640        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
641        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
642        for counter in self.pos_drafted.iter().take(drafted) {
643            counter.fetch_add(1, Ordering::Relaxed);
644        }
645        for counter in self.pos_accepted.iter().take(accepted) {
646            counter.fetch_add(1, Ordering::Relaxed);
647        }
648    }
649
650    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
651    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
652    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
653        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
654        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
655        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
656    }
657
658    fn snapshot(&self) -> SpecTelemetry {
659        SpecTelemetry {
660            rounds: self.rounds.load(Ordering::Relaxed),
661            drafted: self.drafted.load(Ordering::Relaxed),
662            accepted: self.accepted.load(Ordering::Relaxed),
663            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
664            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
665        }
666    }
667}
668
669pub struct SpecSession {
670    pub(crate) cache: Cache,
671    pub(crate) scratch: MtpScratch,
672    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
673    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
674    /// session must count them. Callers render output from this, not from their own echo.
675    pub committed: Vec<u32>,
676    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
677    pub(crate) last_h: Option<CudaSlice<f32>>,
678    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
679    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
680    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
681    pub next_pred: Option<u32>,
682    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
683    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
684    pub sctr: u32,
685    pub uctr: u32,
686    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
687    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
688    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
689    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
690    /// research/spec-serving-20260801). None before the first turn; error paths drop it
691    /// (next burst recaptures — serve retires errored sessions anyway).
692    pub(crate) draft_ctx: Option<DraftGraphCtx>,
693    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
694    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
695    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
696    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
697    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
698    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
699    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
700    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
701    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
702    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
703    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
704    pub pending_tok: Option<u32>,
705    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
706    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
707    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
708    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
709    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
710    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
711    /// accounting the loop already does — no syncs, no allocation. NOTE a
712    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
713    /// diff with [`SpecTelemetry::delta_since`] around each burst.
714    telem: SpecTelemetryCounters,
715    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
716    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
717    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
718    /// prime, result lands in `boundary_captures`.
719    pub capture_at: Option<usize>,
720    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
721    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
722    /// publication just isn't available for that request. Plural since
723    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
724    /// split (the shared-prefix class) and the stable pre-generation boundary (the
725    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
726    /// prefill tick publishes/checkpoints.
727    pub boundary_captures: Vec<SpecBoundaryCapture>,
728    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
729    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
730    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
731    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
732    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
733    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
734    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
735    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
736    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
737    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
738    /// prompt-end capture.
739    pub ckpt_at: Option<usize>,
740}
741impl SpecSession {
742    /// Context capacity of the session's caches (the server's ContextFull guard).
743    pub fn cache_max_ctx(&self) -> usize {
744        self.cache.max_ctx
745    }
746    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
747    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
748    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
749    /// the prime boundary), so no copy was taken at prime time.
750    pub fn cache_ref(&self) -> &Cache {
751        &self.cache
752    }
753    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
754    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
755    /// like the trunk KV — draft rows below the prompt end are append-only for the
756    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
757    /// committed length, never below the prime boundary, and the true-hidden refresh
758    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
759    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
760    /// prefix-addressable; the prefix cache already refuses that class end to end).
761    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
762        if self.scratch.kv.ring.is_some() {
763            return None;
764        }
765        Some((
766            &self.scratch.kv.k,
767            &self.scratch.kv.v,
768            self.scratch.kv.k_tok_bytes,
769            self.scratch.kv.v_tok_bytes,
770        ))
771    }
772    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
773    pub fn telemetry(&self) -> SpecTelemetry {
774        self.telem.snapshot()
775    }
776    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
777    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
778    /// `spec_rewind_to_checkpoint`.
779    pub fn rewind_pos(&self) -> Option<usize> {
780        self.turn_ckpt.as_ref().map(|c| c.pos)
781    }
782    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
783    pub fn rewind_is_resident(&self) -> bool {
784        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
785            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
786        })
787    }
788    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
789    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
790    /// session has never run a turn and has no prediction to hand over.
791    pub fn demote_ready(&self) -> bool {
792        self.pending_tok.is_none() && self.next_pred.is_some()
793    }
794    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
795    pub fn has_pending(&self) -> bool {
796        self.pending_tok.is_some()
797    }
798    /// Committed row count == cache rows (the session invariant), for the caller's own
799    /// `fed`-length cross-check at a handoff boundary.
800    pub fn committed_len(&self) -> usize {
801        self.committed.len()
802    }
803    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
804    /// cache + next-token prediction to the plain batched-decode path.
805    ///
806    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
807    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
808    /// tokenwise prime of the same `committed` sequence would have left it (that is the
809    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
810    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
811    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
812    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
813    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
814    /// a state indistinguishable from one the batched path produced itself: the batched tick
815    /// emits `next_pred`, feeds it into this same cache, and decodes on.
816    ///
817    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
818    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
819    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
820    /// path would silently skip a token.
821    ///
822    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
823    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
824    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
825    /// would mean an `mtp_kv_fill` over the whole committed history).
826    pub fn into_demoted(self) -> Option<(Cache, u32)> {
827        if self.pending_tok.is_some() {
828            return None;
829        }
830        let np = self.next_pred?;
831        debug_assert_eq!(
832            self.cache.pos,
833            self.committed.len(),
834            "demotion handoff: cache rows != committed tokens"
835        );
836        Some((self.cache, np))
837    }
838    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
839    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
840    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
841    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
842    pub fn reset_graph_fallback_on_resume(&mut self) {
843        if let Some(line) = self
844            .draft_ctx
845            .as_mut()
846            .and_then(|c| c.failed.reset_on_resume())
847        {
848            eprintln!("{line}");
849        }
850    }
851}
852
853/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
854///
855/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
856/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
857/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
858/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
859/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
860/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
861///
862/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
863/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
864/// position index, so it must be a real device COPY — that copy is the entire reason a spec
865/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
866/// below the boundary were written by this turn's fill and are never revisited (the per-round
867/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
868/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
869/// predecessor-pairing anchor the next prime's fill reads for its first row.
870///
871/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
872pub(crate) struct SpecCheckpoint {
873    snap: crate::cache::CacheSnapshot,
874    /// Committed length at the boundary (== cache.pos there, the session invariant).
875    pos: usize,
876    /// Pre-output_norm hidden of row `pos - 1`.
877    last_h: CudaSlice<f32>,
878}
879
880/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
881/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
882/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
883/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
884/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
885/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
886/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
887/// so the worker slices those from the live caches post-burst instead of copying at prime time.
888pub struct SpecBoundaryCapture {
889    pub snap: crate::cache::CacheSnapshot,
890    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
891    pub pos: usize,
892    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
893    pub logits: Vec<f32>,
894    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
895    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
896    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
897    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
898    pub last_h: Vec<f32>,
899}
900
901/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
902/// spec boundary capture carries for later restored-session fills. Failure is silent
903/// (`turn_ckpt` convention): the capture publishes without an anchor.
904fn capture_boundary_hidden(
905    e: &Engine,
906    h_rows: &CudaSlice<f32>,
907    pos: usize,
908    n_embd: usize,
909) -> Vec<f32> {
910    if pos == 0 || h_rows.len() < pos * n_embd {
911        return Vec::new();
912    }
913    let Ok(mut row) = e.uninit(n_embd) else {
914        return Vec::new();
915    };
916    if e.copy_view_into(
917        &mut row,
918        0,
919        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
920        n_embd,
921    )
922    .is_err()
923    {
924        return Vec::new();
925    }
926    e.dtoh(&row).unwrap_or_default()
927}
928
929/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
930/// Default ON: the token a burst emits at its own boundary is drawn from the request's
931/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
932/// every boundary) without touching greedy, which is byte-unaffected either way.
933pub fn spec_sampled_boundary_on() -> bool {
934    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
935    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
936}
937
938/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
939/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
940/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
941/// restores the pre-lane posture (each burst restarts the window from its own prompt
942/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
943/// must keep refusing penalized sampled prefix-cache restores, because the restored
944/// session's continuation burst is handed no prompt slice at all.
945pub fn spec_pen_session_on() -> bool {
946    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
947    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
948}
949
950/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
951/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
952/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
953/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
954/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
955/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
956pub fn spec_restore_republish_on() -> bool {
957    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
958    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
959}
960
961/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
962/// the argmax the pre-lane code would have emitted from the same row. This is how the
963/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
964fn spec_boundary_trace() -> bool {
965    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
966    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
967}
968
969/// llama-parity floor for the penalty window when the request does not ask for a bigger
970/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
971/// non-identity penalty, so this floor only matters to explicit small windows and to the
972/// CLI env path.
973const PEN_WINDOW_FLOOR: usize = 64;
974
975/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
976/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
977/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
978/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
979/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
980/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
981/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
982/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
983/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
984/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
985/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
986/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
987/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
988/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
989/// is a second thing to drift.
990pub const PEN_WINDOW_MAX: usize = 8192;
991
992/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
993/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
994/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
995/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
996/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
997/// client actually asked us to penalize, where the pre-lane code had NOTHING.
998/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
999/// window through the SAME function (one definition of "the window" across both spec
1000/// routes and the gate binary's trunk-only reference arm).
1001pub fn pen_window_seed(
1002    session_committed: &[u32],
1003    burst_prompt: &[u32],
1004    penalty_last_n: usize,
1005) -> Vec<u32> {
1006    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1007    let take_prompt = burst_prompt.len().min(win);
1008    let take_sess = (win - take_prompt).min(session_committed.len());
1009    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1010    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1011    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1012    hist
1013}
1014
1015/// Draw a BOUNDARY token from the target distribution the request asked for
1016/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1017/// every burst boundary".
1018///
1019/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1020/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1021/// row after the last committed token on a continuation burst; the prefix-cache entry's
1022/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1023/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1024/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1025/// customer asked for a sampled token, so this draws one.
1026///
1027/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1028/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1029/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1030/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1031/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1032/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1033///
1034/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1035/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1036/// stream the accept walk uses — never a second, independently seeded stream (which would be
1037/// a new distributional bug: two streams from one seed correlate wherever their counters
1038/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1039/// to the cold session's own first draw from the same logits row, which is what preserves the
1040/// sampled-hit lane's per-seed hit==cold byte identity.
1041#[allow(clippy::too_many_arguments)]
1042pub fn sample_boundary_token_dev(
1043    e: &Engine,
1044    logits: &CudaSlice<f32>,
1045    n_vocab: usize,
1046    sp: &SpecSampling,
1047    pen_hist: &[u32],
1048    sctr: &mut u32,
1049    site: &str,
1050) -> Result<u32, Box<dyn std::error::Error>> {
1051    debug_assert!(
1052        sp.temp > 0.0,
1053        "boundary sampling is the sampled regime only"
1054    );
1055    // Own copy: penalize_logits mutates in place and the caller's row is live state
1056    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1057    let mut col = e.zeros(n_vocab)?;
1058    e.copy_into(&mut col, 0, logits, n_vocab)?;
1059    let pen_on = sp.penalty_last_n > 0
1060        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1061    if pen_on && !pen_hist.is_empty() {
1062        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1063        let w0 = pen_hist
1064            .len()
1065            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1066        let hist = &pen_hist[w0..];
1067        let hd = e.htod_u32_v(hist)?;
1068        e.penalize_logits(
1069            &mut col,
1070            &hd,
1071            hist.len(),
1072            sp.penalty_repeat,
1073            sp.penalty_freq,
1074            sp.penalty_present,
1075            n_vocab,
1076        )?;
1077    }
1078    let rows0 = e.htod_i32(&[0])?;
1079    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1080    e.filter_stats(
1081        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1082        sp.top_p, sp.min_p,
1083    )?;
1084    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1085    let mut perturb = e.zeros(n_vocab)?;
1086    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1087    *sctr = sctr.wrapping_add(1);
1088    let td = e.argmax_token_device(&perturb, n_vocab)?;
1089    let tok = e.dtoh_u32_one(&td)?;
1090    if spec_boundary_trace() {
1091        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1092        let raw = e.argmax_token_device(logits, n_vocab)?;
1093        let greedy = e.dtoh_u32_one(&raw)?;
1094        eprintln!(
1095            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1096             deviates={} temp={} sctr={}",
1097            (tok != greedy) as u8,
1098            sp.temp,
1099            sctr.wrapping_sub(1),
1100        );
1101    }
1102    Ok(tok)
1103}
1104
1105/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1106/// host `Vec<f32>`).
1107#[allow(clippy::too_many_arguments)]
1108pub fn sample_boundary_token(
1109    e: &Engine,
1110    logits: &[f32],
1111    sp: &SpecSampling,
1112    pen_hist: &[u32],
1113    sctr: &mut u32,
1114    site: &str,
1115) -> Result<u32, Box<dyn std::error::Error>> {
1116    let n_vocab = logits.len();
1117    let d = e.htod(logits)?;
1118    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1119}
1120
1121struct SpecPipeTraceClock {
1122    pair: usize,
1123    started: std::time::Instant,
1124}
1125
1126#[derive(Clone)]
1127struct SpecPipeTraceCtx {
1128    clock: std::sync::Arc<SpecPipeTraceClock>,
1129    round: usize,
1130    lane: usize,
1131}
1132
1133struct SpecPipeTraceMarker {
1134    trace: SpecPipeTraceCtx,
1135    phase: &'static str,
1136    edge: &'static str,
1137    slot: Option<usize>,
1138}
1139
1140unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1141    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1142    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1143    let slot = marker
1144        .slot
1145        .map(|v| v.to_string())
1146        .unwrap_or_else(|| "-".into());
1147    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1148    use std::io::Write as _;
1149    let stderr = std::io::stderr();
1150    let mut stderr = stderr.lock();
1151    let _ = writeln!(
1152        stderr,
1153        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1154         slot={slot} t_ms={t_ms:.3}",
1155        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1156    );
1157}
1158
1159fn enqueue_spec_pipe_trace_marker(
1160    stream: &cudarc::driver::CudaStream,
1161    trace: Option<&SpecPipeTraceCtx>,
1162    phase: &'static str,
1163    edge: &'static str,
1164    slot: Option<usize>,
1165) -> Result<(), Box<dyn std::error::Error>> {
1166    let Some(trace) = trace else {
1167        return Ok(());
1168    };
1169    let marker = Box::new(SpecPipeTraceMarker {
1170        trace: trace.clone(),
1171        phase,
1172        edge,
1173        slot,
1174    });
1175    let raw = Box::into_raw(marker);
1176    let result = unsafe {
1177        cudarc::driver::result::stream::launch_host_function(
1178            stream.cu_stream(),
1179            spec_pipe_trace_marker,
1180            raw.cast(),
1181        )
1182    };
1183    if let Err(err) = result {
1184        unsafe {
1185            drop(Box::from_raw(raw));
1186        }
1187        return Err(err.into());
1188    }
1189    Ok(())
1190}
1191
1192#[derive(Default)]
1193struct SpecPipeProgress {
1194    setup_done: [bool; 2],
1195    draft_done: [usize; 2],
1196    stage0_done: [usize; 2],
1197    verify_done: [usize; 2],
1198    accept_done: [usize; 2],
1199    finished: [bool; 2],
1200    aborted: bool,
1201}
1202
1203/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1204/// keeps its existing call stack and round locals; this object only orders phase entry. The
1205/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1206/// cannot be interleaved by the two host threads.
1207struct SpecPipeSync {
1208    progress: std::sync::Mutex<SpecPipeProgress>,
1209    changed: std::sync::Condvar,
1210    primary: std::sync::Mutex<()>,
1211    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1212}
1213
1214impl SpecPipeSync {
1215    fn new() -> Self {
1216        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1217        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1218            std::sync::Arc::new(SpecPipeTraceClock {
1219                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1220                started: std::time::Instant::now(),
1221            })
1222        });
1223        Self {
1224            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1225            changed: std::sync::Condvar::new(),
1226            primary: std::sync::Mutex::new(()),
1227            trace,
1228        }
1229    }
1230}
1231
1232#[derive(Clone)]
1233struct SpecPipeLane {
1234    sync: std::sync::Arc<SpecPipeSync>,
1235    lane: usize,
1236}
1237
1238impl SpecPipeLane {
1239    fn peer(&self) -> usize {
1240        1 - self.lane
1241    }
1242
1243    fn aborted() -> Box<dyn std::error::Error> {
1244        "paired speculative peer aborted".into()
1245    }
1246
1247    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1248        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1249            clock: clock.clone(),
1250            round,
1251            lane: self.lane,
1252        })
1253    }
1254
1255    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1256        let mut p = self.sync.progress.lock().unwrap();
1257        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1258            p = self.sync.changed.wait(p).unwrap();
1259        }
1260        if p.aborted {
1261            Err(Self::aborted())
1262        } else {
1263            Ok(())
1264        }
1265    }
1266
1267    fn setup_end(&self) {
1268        let mut p = self.sync.progress.lock().unwrap();
1269        p.setup_done[self.lane] = true;
1270        self.sync.changed.notify_all();
1271    }
1272
1273    fn draft_begin(
1274        &self,
1275        round: usize,
1276    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1277        let peer = self.peer();
1278        let mut p = self.sync.progress.lock().unwrap();
1279        loop {
1280            if p.aborted {
1281                return Err(Self::aborted());
1282            }
1283            let setup_ready =
1284                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1285            let prior_ready = p.accept_done[self.lane] >= round
1286                && (p.accept_done[peer] >= round || p.finished[peer]);
1287            let turn_ready = if self.lane == 0 {
1288                true
1289            } else {
1290                p.draft_done[0] > round || p.finished[0]
1291            };
1292            if setup_ready && prior_ready && turn_ready {
1293                break;
1294            }
1295            p = self.sync.changed.wait(p).unwrap();
1296        }
1297        drop(p);
1298        Ok(self.sync.primary.lock().unwrap())
1299    }
1300
1301    fn draft_end(&self, round: usize) {
1302        let mut p = self.sync.progress.lock().unwrap();
1303        p.draft_done[self.lane] = round + 1;
1304        self.sync.changed.notify_all();
1305    }
1306
1307    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1308    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1309    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1310        let peer = self.peer();
1311        let mut p = self.sync.progress.lock().unwrap();
1312        loop {
1313            if p.aborted {
1314                return Err(Self::aborted());
1315            }
1316            let ready = if self.lane == 0 {
1317                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1318            } else {
1319                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1320            };
1321            if ready {
1322                return Ok(self.lane == 0 || p.finished[peer]);
1323            }
1324            p = self.sync.changed.wait(p).unwrap();
1325        }
1326    }
1327
1328    fn stage0_end(&self, round: usize) {
1329        let mut p = self.sync.progress.lock().unwrap();
1330        p.stage0_done[self.lane] = round + 1;
1331        self.sync.changed.notify_all();
1332    }
1333
1334    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1335    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1336    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1337        let mut p = self.sync.progress.lock().unwrap();
1338        while !p.aborted
1339            && !(p.stage0_done[self.lane] > round
1340                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1341        {
1342            p = self.sync.changed.wait(p).unwrap();
1343        }
1344        if p.aborted {
1345            Err(Self::aborted())
1346        } else {
1347            Ok(())
1348        }
1349    }
1350
1351    fn verify_end(&self, round: usize) {
1352        let mut p = self.sync.progress.lock().unwrap();
1353        p.verify_done[self.lane] = round + 1;
1354        self.sync.changed.notify_all();
1355    }
1356
1357    fn accept_begin(
1358        &self,
1359        round: usize,
1360    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1361        let mut p = self.sync.progress.lock().unwrap();
1362        loop {
1363            if p.aborted {
1364                return Err(Self::aborted());
1365            }
1366            let ready = if self.lane == 0 {
1367                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1368            } else {
1369                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1370            };
1371            if ready {
1372                break;
1373            }
1374            p = self.sync.changed.wait(p).unwrap();
1375        }
1376        drop(p);
1377        Ok(self.sync.primary.lock().unwrap())
1378    }
1379
1380    fn accept_end(&self, round: usize) {
1381        let mut p = self.sync.progress.lock().unwrap();
1382        p.accept_done[self.lane] = round + 1;
1383        self.sync.changed.notify_all();
1384    }
1385
1386    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1387        self.sync.primary.lock().unwrap()
1388    }
1389
1390    fn finish(&self, failed: bool) {
1391        let mut p = self.sync.progress.lock().unwrap();
1392        p.finished[self.lane] = true;
1393        p.aborted |= failed;
1394        self.sync.changed.notify_all();
1395    }
1396}
1397
1398struct SpecPipeFinish<'a> {
1399    lane: &'a SpecPipeLane,
1400    closed: bool,
1401}
1402
1403impl<'a> SpecPipeFinish<'a> {
1404    fn new(lane: &'a SpecPipeLane) -> Self {
1405        Self {
1406            lane,
1407            closed: false,
1408        }
1409    }
1410
1411    fn close(&mut self, failed: bool) {
1412        self.lane.finish(failed);
1413        self.closed = true;
1414    }
1415}
1416
1417impl Drop for SpecPipeFinish<'_> {
1418    fn drop(&mut self) {
1419        if !self.closed {
1420            self.lane.finish(true);
1421        }
1422    }
1423}
1424
1425/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1426/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1427/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1428/// binds that context before touching the session, joins before returning, and never aliases the
1429/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1430/// session type Send.
1431struct SpecPipeSessionPtr(*mut SpecSession);
1432
1433unsafe impl Send for SpecPipeSessionPtr {}
1434
1435impl SpecPipeSessionPtr {
1436    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1437        unsafe { &mut *self.0 }
1438    }
1439}
1440
1441/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1442/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1443/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1444/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1445/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1446/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1447/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1448/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1449/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1450///
1451/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1452/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1453/// load-bearing:
1454///
1455/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1456///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1457///   This is all the key used to carry.
1458/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1459///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1460///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1461///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1462///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1463///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1464///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1465///
1466/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1467/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1468/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1469/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1470/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1471#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1472pub(crate) struct SampledGraphKey {
1473    seed: u64,
1474    temp_bits: u32,
1475    k: usize,
1476    top_k: i32,
1477    top_p_bits: u32,
1478    min_p_bits: u32,
1479    pen_on: bool,
1480}
1481
1482impl SampledGraphKey {
1483    pub(crate) fn new(
1484        seed: u64,
1485        temp: f32,
1486        k: usize,
1487        top_k: i32,
1488        top_p: f32,
1489        min_p: f32,
1490        pen_on: bool,
1491    ) -> Self {
1492        SampledGraphKey {
1493            seed,
1494            temp_bits: temp.to_bits(),
1495            k,
1496            top_k,
1497            top_p_bits: top_p.to_bits(),
1498            min_p_bits: min_p.to_bits(),
1499            pen_on,
1500        }
1501    }
1502
1503    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1504    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1505    /// the key can never drift apart (they were three separate expressions before this lane, and
1506    /// the launch site simply forgot to ask).
1507    pub(crate) fn pure_temp(&self) -> bool {
1508        self.top_k == 0
1509            && f32::from_bits(self.top_p_bits) >= 1.0
1510            && f32::from_bits(self.min_p_bits) <= 0.0
1511            && !self.pen_on
1512    }
1513}
1514
1515pub(crate) struct DraftGraphCtx {
1516    g_tok: CudaSlice<u32>,
1517    g_pos: CudaSlice<i32>,
1518    g_seed: CudaSlice<f32>,
1519    g_p: CudaSlice<f32>,
1520    g_ctr: CudaSlice<u32>,
1521    g_q: CudaSlice<f32>,
1522    g_perturb: CudaSlice<f32>,
1523    q_slots: Vec<CudaSlice<f32>>,
1524    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1525    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1526    /// per-position contents the host re-uploads before each replay (the graph-promote
1527    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1528    g_dmask: CudaSlice<u32>,
1529    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1530    graph_masked: bool,
1531    graph: Option<cudarc::driver::CudaGraph>,
1532    graph_s: Option<cudarc::driver::CudaGraph>,
1533    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1534    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1535    failed: DraftGraphFallback,
1536    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1537    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1538    s_key: Option<SampledGraphKey>,
1539    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1540    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1541    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1542    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1543    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1544    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1545    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1546    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1547    keeper: Vec<Box<dyn std::any::Any + Send>>,
1548    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1549}
1550
1551/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1552/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1553///
1554/// Three contracts:
1555/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1556///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1557///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1558///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1559///   fallback from paying a doomed capture attempt every burst).
1560/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1561///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1562///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1563///   actually set (quiet on the common clean-resume path).
1564/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1565///   capture attempt whose own failure would re-flip loudly.
1566#[derive(Default)]
1567pub(crate) struct DraftGraphFallback {
1568    greedy: bool,
1569    sampled: bool,
1570}
1571impl DraftGraphFallback {
1572    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1573        if self.greedy {
1574            return None;
1575        }
1576        self.greedy = true;
1577        Some(format!(
1578            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1579        ))
1580    }
1581    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1582        if self.sampled {
1583            return None;
1584        }
1585        self.sampled = true;
1586        Some(format!(
1587            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1588        ))
1589    }
1590    fn greedy_failed(&self) -> bool {
1591        self.greedy
1592    }
1593    fn sampled_failed(&self) -> bool {
1594        self.sampled
1595    }
1596    fn clear_greedy(&mut self) {
1597        self.greedy = false;
1598    }
1599    fn clear_sampled(&mut self) {
1600        self.sampled = false;
1601    }
1602    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1603    /// was set (so clean resumes stay quiet).
1604    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1605        if !self.greedy && !self.sampled {
1606            return None;
1607        }
1608        let which = match (self.greedy, self.sampled) {
1609            (true, true) => "greedy+sampled",
1610            (true, false) => "greedy",
1611            _ => "sampled",
1612        };
1613        self.greedy = false;
1614        self.sampled = false;
1615        Some(format!(
1616            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1617        ))
1618    }
1619}
1620
1621impl DraftGraphCtx {
1622    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1623        Ok(DraftGraphCtx {
1624            g_tok: e.alloc_u32_zeroed(1)?,
1625            g_pos: e.htod_i32(&[0])?,
1626            g_seed: e.zeros(n_embd)?,
1627            g_p: e.zeros(1)?,
1628            g_ctr: e.alloc_u32_zeroed(1)?,
1629            g_q: e.zeros(qlen)?,
1630            g_perturb: e.zeros(qlen)?,
1631            q_slots: Vec::new(),
1632            g_dmask: e.alloc_u32_zeroed(1)?,
1633            graph_masked: false,
1634            graph: None,
1635            graph_s: None,
1636            failed: DraftGraphFallback::default(),
1637            s_key: None,
1638            keeper: Vec::new(),
1639            keeper_s: Vec::new(),
1640        })
1641    }
1642}
1643
1644pub(crate) struct MtpScratch {
1645    kv: KvLayer,
1646    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1647    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1648    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1649    /// smaller host-indexed SWA ring instead.
1650    cap: usize,
1651}
1652
1653fn mtp_scratch_layout(
1654    cfg: &memra_gguf::config::ModelConfig,
1655    geom: Option<&crate::hybrid::DraftGeom>,
1656) -> (usize, usize, usize, usize) {
1657    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1658    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1659    let head_dim_k = cfg.head_dim_k as usize;
1660    let head_dim_v = cfg.head_dim_v as usize;
1661    assert!(
1662        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1663        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1664    );
1665    let kv_dim_k = head_dim_k * n_head_kv;
1666    let kv_dim_v = head_dim_v * n_head_kv;
1667    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1668    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1669    let (kbb, vbb) = crate::kv_blk_bytes();
1670    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1671    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1672    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1673}
1674
1675impl MtpScratch {
1676    fn new(
1677        e: &Engine,
1678        cfg: &memra_gguf::config::ModelConfig,
1679        cap: usize,
1680        geom: Option<&crate::hybrid::DraftGeom>,
1681    ) -> Result<Self, Box<dyn std::error::Error>> {
1682        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1683        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1684        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1685        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1686        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1687        let ring = if crate::cache::swa_ring_on() && cfg.arch.is_step35() {
1688            let window = cfg.step35.as_ref().unwrap().sliding_window as usize;
1689            Some(crate::cache::KvRing::new(
1690                crate::cache::swa_ring_rows(window, cap),
1691                window,
1692            ))
1693        } else {
1694            None
1695        };
1696        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1697        Ok(MtpScratch {
1698            kv: KvLayer {
1699                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1700                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1701                kv_dim_k,
1702                kv_dim_v,
1703                k_tok_bytes,
1704                v_tok_bytes,
1705                len: 0,
1706                ring,
1707                len_d: e.htod_i32(&[0])?,
1708            },
1709            cap,
1710        })
1711    }
1712    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1713    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1714    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1715    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1716        if self
1717            .kv
1718            .ring
1719            .as_ref()
1720            .is_some_and(|ring| !ring.can_rewind_to(n))
1721        {
1722            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1723        }
1724        self.kv.len = n;
1725        e.set_i32_one(&mut self.kv.len_d, n as i32)
1726    }
1727
1728    fn can_rewind_to(&self, n: usize) -> bool {
1729        self.kv
1730            .ring
1731            .as_ref()
1732            .is_none_or(|ring| ring.can_rewind_to(n))
1733    }
1734}
1735
1736/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1737/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1738/// full weight reads per round — recomputing columns the verify had already produced
1739/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1740/// to "after the first j verify columns" WITHOUT re-running the trunk:
1741/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1742///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1743///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1744///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1745///   pure-copy ring rebuild.
1746/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1747///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1748///   target: j <= t-1).
1749/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1750/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1751struct GdnStash {
1752    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1753    q_l2: CudaSlice<f32>,
1754    k_l2: CudaSlice<f32>,
1755    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1756    g_log: CudaSlice<f32>,
1757    beta: CudaSlice<f32>, // [t, num_v]
1758}
1759pub(crate) struct VerifyCkpt {
1760    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1761    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1762}
1763/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1764pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1765
1766/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1767/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1768/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1769/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1770/// layers between full-attention layers are shape-static given vt — no positions, no
1771/// t_kv, state addressed through pointer tables — so runs of them capture per
1772/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1773/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1774///
1775/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1776/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1777/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1778/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1779/// before and restored after — the graph's first real launch starts from the exact
1780/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1781/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1782/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1783pub(crate) struct DsparkVerifyGraphs {
1784    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1785    lin: Vec<usize>,
1786    lin_pos: std::collections::HashMap<usize, usize>,
1787    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1788    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1789    table_all: CudaSlice<u64>,
1790    host_table: Vec<u64>,
1791    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1792    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1793    stash_conv: Vec<CudaSlice<f32>>,
1794    stash_ssm: Vec<CudaSlice<f32>>,
1795    conv_words: usize,
1796    ssm_words: usize,
1797    /// Per-vt input/output staging (stable addresses the graphs bake).
1798    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1799    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1800    /// so the sink buffer must live (and persist) with the graphs, not with the round.
1801    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1802    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1803    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
1804    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
1805    save_conv: CudaSlice<f32>,
1806    save_ssm: CudaSlice<f32>,
1807    max_run: usize,
1808    n_embd: usize,
1809    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1810    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1811    pub(crate) round_slab: bool,
1812    // ---- slice 4c: full-verify single graph per (vt, rung) ----
1813    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
1814    fa: Vec<usize>,
1815    fa_pos: std::collections::HashMap<usize, usize>,
1816    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
1817    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
1818    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
1819    fa_table: CudaSlice<u64>,
1820    fa_host_table: Vec<u64>,
1821    t_cap: usize,
1822    /// Per-vt position staging for the captured bodies — contents refreshed per round
1823    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
1824    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
1825    /// Full-verify graphs keyed (vt, rung_end, hi).
1826    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
1827    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
1828    covered: usize,
1829    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
1830    /// full-verify capture walks all of them.
1831    walk_uniform: bool,
1832}
1833
1834struct DsparkSegGraph {
1835    graph: cudarc::driver::CudaGraph,
1836    _keeper: Vec<Box<dyn std::any::Any + Send>>,
1837}
1838
1839/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
1840/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
1841/// modes without a second copy of the math.
1842pub(crate) struct FaLayerArgs<'a> {
1843    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
1844    /// them per-z (append slot = pos, T_kv = pos + 1).
1845    pub pos_d: &'a CudaSlice<i32>,
1846    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
1847    /// arm builds/uses them (graph mode refuses that arm).
1848    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
1849    pub pos0: usize,
1850    pub seqs_append: bool,
1851    pub batch_fa_on: bool,
1852    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
1853    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
1854    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
1855    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
1856    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
1857    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
1858    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
1859    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
1860    /// for FA layers that never touch it.
1861    pub ckpt: Option<&'a mut VerifyCkpt>,
1862}
1863
1864// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
1865// no automatic trait; CUDA driver graph handles are context-scoped rather than
1866// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
1867// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
1868// single decode-stream thread.
1869unsafe impl Send for DsparkVerifyGraphs {}
1870
1871impl DsparkVerifyGraphs {
1872    /// Build for this cache's shape. None when there are no linear layers, sizes are
1873    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
1874    pub(crate) fn new(
1875        e: &Engine,
1876        cache: &Cache,
1877        t_max: usize,
1878        n_embd: usize,
1879    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
1880        let lin: Vec<usize> = (0..cache.recur.len())
1881            .filter(|&il| cache.recur[il].is_some())
1882            .collect();
1883        if lin.is_empty() || t_max < 2 {
1884            return Ok(None);
1885        }
1886        let first = cache.recur[lin[0]].as_ref().unwrap();
1887        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
1888        for &il in &lin {
1889            let rl = cache.recur[il].as_ref().unwrap();
1890            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
1891                return Ok(None);
1892            }
1893        }
1894        let n = lin.len();
1895        let mut lin_pos = std::collections::HashMap::with_capacity(n);
1896        for (k, &il) in lin.iter().enumerate() {
1897            lin_pos.insert(il, k);
1898        }
1899        // longest run of consecutive linear layers (save-scratch sizing)
1900        let mut max_run = 1usize;
1901        let mut run = 1usize;
1902        for w in lin.windows(2) {
1903            if w[1] == w[0] + 1 {
1904                run += 1;
1905                max_run = max_run.max(run);
1906            } else {
1907                run = 1;
1908            }
1909        }
1910        let rows = t_max - 1;
1911        let mut stash_conv = Vec::with_capacity(n);
1912        let mut stash_ssm = Vec::with_capacity(n);
1913        for _ in 0..n {
1914            stash_conv.push(e.uninit(rows * conv_words)?);
1915            stash_ssm.push(e.uninit(rows * ssm_words)?);
1916        }
1917        let host_table = vec![0u64; n * 6];
1918        let table_all = e.htod_u64(&host_table)?;
1919        // slice 4c: full-attention census for the full-verify graphs.
1920        let fa: Vec<usize> = (0..cache.kv.len())
1921            .filter(|&il| cache.kv[il].is_some())
1922            .collect();
1923        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
1924        for (k, &il) in fa.iter().enumerate() {
1925            fa_pos.insert(il, k);
1926        }
1927        let n_layers = cache.kv.len().max(cache.recur.len());
1928        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
1929        let walk_uniform = (0..n_layers).all(|il| {
1930            cache.recur.get(il).is_some_and(|r| r.is_some())
1931                != cache.kv.get(il).is_some_and(|k| k.is_some())
1932        });
1933        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
1934        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
1935        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
1936        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
1937        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
1938        let covered = (0..n_layers)
1939            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
1940            .count();
1941        let t_cap = t_max;
1942        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
1943        let fa_table = e.htod_u64(&fa_host_table)?;
1944        Ok(Some(Self {
1945            lin,
1946            lin_pos,
1947            table_all,
1948            host_table,
1949            stash_conv,
1950            stash_ssm,
1951            conv_words,
1952            ssm_words,
1953            stage: std::collections::HashMap::new(),
1954            tap_bufs: std::collections::HashMap::new(),
1955            graphs: std::collections::HashMap::new(),
1956            save_conv: e.uninit(n * conv_words)?,
1957            save_ssm: e.uninit(n * ssm_words)?,
1958            max_run,
1959            n_embd,
1960            round_slab: false,
1961            fa,
1962            fa_pos,
1963            fa_table,
1964            fa_host_table,
1965            t_cap,
1966            pos_stage: std::collections::HashMap::new(),
1967            full: std::collections::HashMap::new(),
1968            covered,
1969            walk_uniform,
1970        }))
1971    }
1972
1973    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
1974    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
1975    /// cache buffers land at new addresses; a stale table would read the wrong state).
1976    pub(crate) fn refresh_tables(
1977        &mut self,
1978        e: &Engine,
1979        cache: &Cache,
1980    ) -> Result<(), Box<dyn std::error::Error>> {
1981        use cudarc::driver::DevicePtr;
1982        {
1983            let s = &e.gpu.stream();
1984            for (k, &il) in self.lin.iter().enumerate() {
1985                let rl = cache.recur[il].as_ref().unwrap();
1986                let (pc, _g0) = rl.conv_state.device_ptr(s);
1987                let (p0, _g1) = rl.ssm_state.device_ptr(s);
1988                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
1989                let o = k * 6;
1990                self.host_table[o] = pc as u64;
1991                self.host_table[o + 1] = p0 as u64;
1992                self.host_table[o + 2] = p1 as u64;
1993                self.host_table[o + 3] = pc as u64;
1994                self.host_table[o + 4] = p1 as u64;
1995                self.host_table[o + 5] = p0 as u64;
1996            }
1997            for (k, &il) in self.fa.iter().enumerate() {
1998                let kvl = cache.kv[il].as_ref().unwrap();
1999                let (pk, _g0) = kvl.k.device_ptr(s);
2000                let (pv, _g1) = kvl.v.device_ptr(s);
2001                let o = k * 2 * self.t_cap;
2002                for z in 0..self.t_cap {
2003                    self.fa_host_table[o + 2 * z] = pk as u64;
2004                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2005                }
2006            }
2007        }
2008        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2009        if !self.fa_host_table.is_empty() {
2010            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2011        }
2012        Ok(())
2013    }
2014
2015    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2016    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2017    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2018    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2019    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2020    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2021    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2022    /// captured graph is bit-identical for every round the rung covers.
2023    #[allow(clippy::too_many_arguments)]
2024    pub(crate) fn full_rung(
2025        &self,
2026        model: &crate::hybrid::HybridModel,
2027        cache: &Cache,
2028        lo: usize,
2029        hi: usize,
2030        t: usize,
2031        seqs_arms_on: bool,
2032    ) -> Option<usize> {
2033        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2034            static ONCE: std::sync::Once = std::sync::Once::new();
2035            let len0 = self
2036                .fa
2037                .first()
2038                .and_then(|&il| cache.kv[il].as_ref())
2039                .map(|k| k.len);
2040            ONCE.call_once(|| {
2041                eprintln!(
2042                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2043                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2044                    self.lin.len(), self.fa.len(), self.t_cap, len0
2045                );
2046            });
2047        }
2048        if !self.walk_uniform
2049            || !seqs_arms_on
2050            || !dspark_fa_rows_on()
2051            || t < 2
2052            || lo != 0
2053            || hi > self.covered
2054            || t > self.t_cap
2055            || self.fa.is_empty()
2056        {
2057            return None;
2058        }
2059        let cfg = &model.cfg;
2060        let head_dim_global = cfg.head_dim_k as usize;
2061        let nkv = cfg.n_head_kv as usize;
2062        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2063        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2064        // projection stride (the body's guard, hoisted so ineligible models fall back
2065        // instead of refusing mid-capture).
2066        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2067        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2068        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2069            return None;
2070        }
2071        let len0 = kvl0.len;
2072        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2073        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2074            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2075            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2076        {
2077            return None;
2078        }
2079        let rung = t_kv_last.next_power_of_two().max(256);
2080        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2081            return None;
2082        }
2083        Some(rung)
2084    }
2085
2086    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2087    /// the residual + refresh the per-vt position staging, capture on first encounter
2088    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2089    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2090    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2091    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2092    #[allow(clippy::too_many_arguments)]
2093    pub(crate) fn run_full(
2094        &mut self,
2095        model: &crate::hybrid::HybridModel,
2096        e: &Engine,
2097        lo: usize,
2098        hi: usize,
2099        x: &CudaSlice<f32>,
2100        t: usize,
2101        pos0: usize,
2102        rung: usize,
2103        cache: &mut Cache,
2104    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2105        let n_embd = self.n_embd;
2106        if !self.stage.contains_key(&t) {
2107            let xin = e.uninit(t * n_embd)?;
2108            let xout = e.uninit(t * n_embd)?;
2109            self.stage.insert(t, (xin, xout));
2110        }
2111        if !self.pos_stage.contains_key(&t) {
2112            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2113        }
2114        // Per-round refresh: position contents + input staging (both addresses are baked
2115        // by the captured bodies; only their CONTENTS change round to round).
2116        {
2117            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2118            let pb = self.pos_stage.get_mut(&t).unwrap();
2119            e.htod_i32_into(pb, &pos_host)?;
2120            let (xin, _) = self.stage.get_mut(&t).unwrap();
2121            e.copy_into(xin, 0, x, t * n_embd)?;
2122        }
2123        let key = (t, rung, hi);
2124        if !self.full.contains_key(&key) {
2125            // The warmups EXECUTE the whole walk on live state — save every linear
2126            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2127            // graph mode never bumps host lens and the appends write this round's own
2128            // slots).
2129            for (k, &il) in self.lin.iter().enumerate() {
2130                let rl = cache.recur[il].as_ref().unwrap();
2131                e.copy_into(
2132                    &mut self.save_conv,
2133                    k * self.conv_words,
2134                    &rl.conv_state,
2135                    self.conv_words,
2136                )?;
2137                e.copy_into(
2138                    &mut self.save_ssm,
2139                    k * self.ssm_words,
2140                    &rl.ssm_state,
2141                    self.ssm_words,
2142                )?;
2143            }
2144            let (graph, keeper) = {
2145                let table_all = &self.table_all;
2146                let lin_pos = &self.lin_pos;
2147                let fa_pos = &self.fa_pos;
2148                let fa_table = &self.fa_table;
2149                let t_cap = self.t_cap;
2150                let stash_conv = &mut self.stash_conv;
2151                let stash_ssm = &mut self.stash_ssm;
2152                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2153                let (xin, xout) = self
2154                    .stage
2155                    .get_mut(&t)
2156                    .map(|(a, b)| (&*a, b))
2157                    .expect("stage bucket created above");
2158                let cache_ref: &mut Cache = cache;
2159                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2160                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2161                } else {
2162                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2163                };
2164                e.capture_graph_retained_flags(iflag, move |e| {
2165                    let mut xc: Option<CudaSlice<f32>> = None;
2166                    for il in lo..hi {
2167                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2168                        let nx = if let Some(&k) = lin_pos.get(&il) {
2169                            model.qwen35_tparallel_linear_layer(
2170                                e,
2171                                il,
2172                                xr,
2173                                t,
2174                                cache_ref,
2175                                None,
2176                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2177                                Some((table_all, k * 6)),
2178                            )?
2179                        } else if let Some(&kf) = fa_pos.get(&il) {
2180                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2181                            model.qwen35_tparallel_fa_layer(
2182                                e,
2183                                il,
2184                                xr,
2185                                t,
2186                                cache_ref,
2187                                FaLayerArgs {
2188                                    pos_d,
2189                                    pos_rows: &mut no_rows,
2190                                    pos0,
2191                                    seqs_append: true,
2192                                    batch_fa_on: true,
2193                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2194                                    stream: None,
2195                                    ckpt: None,
2196                                },
2197                            )?
2198                        } else {
2199                            return Err(format!(
2200                                "run_full: layer {il} is neither linear nor full-attention"
2201                            )
2202                            .into());
2203                        };
2204                        xc = Some(nx);
2205                    }
2206                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2207                    Ok(())
2208                })?
2209            };
2210            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2211            // is odd -> 3 runs = net one swap), then restore the device state the
2212            // warmups consumed (walk scope only — layers past hi never executed). The
2213            // launch below then behaves exactly like one run.
2214            if t % 2 == 1 {
2215                for &il in &self.lin {
2216                    if il < lo || il >= hi {
2217                        continue;
2218                    }
2219                    let rl = cache.recur[il].as_mut().unwrap();
2220                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2221                }
2222            }
2223            for (k, &il) in self.lin.iter().enumerate() {
2224                if il < lo || il >= hi {
2225                    continue;
2226                }
2227                let rl = cache.recur[il].as_mut().unwrap();
2228                let (cw, sw) = (self.conv_words, self.ssm_words);
2229                {
2230                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2231                    let win = sv.slice(k * cw..(k + 1) * cw);
2232                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2233                }
2234                {
2235                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2236                    let win = sv.slice(k * sw..(k + 1) * sw);
2237                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2238                }
2239            }
2240            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2241                if let Ok(c) = crate::graph_update::node_census(&graph) {
2242                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2243                }
2244            }
2245            self.full.insert(
2246                key,
2247                DsparkSegGraph {
2248                    graph,
2249                    _keeper: keeper,
2250                },
2251            );
2252        }
2253        self.full[&key].graph.launch()?;
2254        // Host bookkeeping for the replayed body (captured host code does not re-run):
2255        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2256        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2257        // head layer's kv) that the walk never touches.
2258        if t % 2 == 1 {
2259            for &il in &self.lin {
2260                if il < lo || il >= hi {
2261                    continue;
2262                }
2263                let rl = cache.recur[il].as_mut().unwrap();
2264                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2265            }
2266        }
2267        for &il in &self.fa {
2268            if il < lo || il >= hi {
2269                continue;
2270            }
2271            cache.kv[il].as_mut().unwrap().len += t;
2272        }
2273        let (_, xout) = self.stage.get(&t).unwrap();
2274        let mut out = e.uninit(t * n_embd)?;
2275        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2276        Ok(out)
2277    }
2278
2279    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2280    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2281    /// bracketed by a segment state save/restore), launch, then apply the host parity
2282    /// bookkeeping the captured body would have done. Returns the fresh residual.
2283    #[allow(clippy::too_many_arguments)]
2284    fn run_segment(
2285        &mut self,
2286        model: &crate::hybrid::HybridModel,
2287        e: &Engine,
2288        start: usize,
2289        end: usize,
2290        x: &CudaSlice<f32>,
2291        t: usize,
2292        cache: &mut Cache,
2293    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2294        let n_embd = self.n_embd;
2295        debug_assert!(end - start <= self.max_run);
2296        if !self.stage.contains_key(&t) {
2297            let xin = e.uninit(t * n_embd)?;
2298            let xout = e.uninit(t * n_embd)?;
2299            self.stage.insert(t, (xin, xout));
2300        }
2301        // Stage the residual at the bucket's baked input address.
2302        {
2303            let (xin, _) = self.stage.get_mut(&t).unwrap();
2304            e.copy_into(xin, 0, x, t * n_embd)?;
2305        }
2306        let key = (start, t);
2307        if !self.graphs.contains_key(&key) {
2308            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2309            // ssm of every segment layer first, restore after, so the graph's first real
2310            // launch starts from the exact pre-round state (bytes gated e2e).
2311            for (k, il) in (start..end).enumerate() {
2312                let rl = cache.recur[il].as_ref().unwrap();
2313                e.copy_into(
2314                    &mut self.save_conv,
2315                    k * self.conv_words,
2316                    &rl.conv_state,
2317                    self.conv_words,
2318                )?;
2319                e.copy_into(
2320                    &mut self.save_ssm,
2321                    k * self.ssm_words,
2322                    &rl.ssm_state,
2323                    self.ssm_words,
2324                )?;
2325            }
2326            let (graph, keeper) = {
2327                let table_all = &self.table_all;
2328                let lin_pos = &self.lin_pos;
2329                let stash_conv = &mut self.stash_conv;
2330                let stash_ssm = &mut self.stash_ssm;
2331                let (xin, xout) = self
2332                    .stage
2333                    .get_mut(&t)
2334                    .map(|(a, b)| (&*a, b))
2335                    .expect("stage bucket created above");
2336                let cache_ref: &mut Cache = cache;
2337                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2338                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2339                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2340                // = ~0.41 ms/round, most of the eager-launch savings. The captured
2341                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2342                // (every transient drops inside the capture region — the generic
2343                // capture path's census precedent, 1589/1589), so AUTO_FREE has
2344                // nothing to reclaim and the graph is legal to instantiate without
2345                // it; PRIORITY is the flag the gemma slotted door ships for exactly
2346                // this reason (both alternatives drop the scan; UPLOAD via
2347                // cuGraphInstantiateWithFlags is WithParams-only and refused).
2348                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2349                // the node census at capture (the ALLOC==FREE receipt).
2350                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2351                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2352                } else {
2353                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2354                };
2355                e.capture_graph_retained_flags(iflag, move |e| {
2356                    let mut xc: Option<CudaSlice<f32>> = None;
2357                    for il in start..end {
2358                        let k = lin_pos[&il];
2359                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2360                        let nx = model.qwen35_tparallel_linear_layer(
2361                            e,
2362                            il,
2363                            xr,
2364                            t,
2365                            cache_ref,
2366                            None,
2367                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
2368                            Some((table_all, k * 6)),
2369                        )?;
2370                        xc = Some(nx);
2371                    }
2372                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2373                    Ok(())
2374                })?
2375            };
2376            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2377            // is odd -> 3 runs = net one swap), then restore the device state the
2378            // warmups consumed. The launch below then behaves exactly like one run.
2379            if t % 2 == 1 {
2380                for il in start..end {
2381                    let rl = cache.recur[il].as_mut().unwrap();
2382                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2383                }
2384            }
2385            for (k, il) in (start..end).enumerate() {
2386                let rl = cache.recur[il].as_mut().unwrap();
2387                let (cw, sw) = (self.conv_words, self.ssm_words);
2388                {
2389                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2390                    let win = sv.slice(k * cw..(k + 1) * cw);
2391                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2392                }
2393                {
2394                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2395                    let win = sv.slice(k * sw..(k + 1) * sw);
2396                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2397                }
2398            }
2399            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2400                if let Ok(c) = crate::graph_update::node_census(&graph) {
2401                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
2402                }
2403            }
2404            self.graphs.insert(
2405                key,
2406                DsparkSegGraph {
2407                    graph,
2408                    _keeper: keeper,
2409                },
2410            );
2411        }
2412        self.graphs[&key].graph.launch()?;
2413        // Host parity bookkeeping for the replayed body (the captured host swaps do not
2414        // re-run at replay).
2415        if t % 2 == 1 {
2416            for il in start..end {
2417                let rl = cache.recur[il].as_mut().unwrap();
2418                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2419            }
2420        }
2421        let (_, xout) = self.stage.get(&t).unwrap();
2422        let mut out = e.uninit(t * n_embd)?;
2423        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2424        Ok(out)
2425    }
2426
2427    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
2428    fn can_capture(&self) -> bool {
2429        self.graphs.len() + self.full.len() < dspark_vg_cap()
2430    }
2431
2432    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
2433    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
2434    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
2435    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
2436    /// refusal would stash some layers in the ctx slabs and others in the round's cols
2437    /// while one commit reads only one of them.
2438    pub(crate) fn segments_ready(
2439        &self,
2440        model: &crate::hybrid::HybridModel,
2441        lo: usize,
2442        hi: usize,
2443        t: usize,
2444    ) -> bool {
2445        if self.can_capture() {
2446            return true;
2447        }
2448        let mut il = lo;
2449        while il < hi {
2450            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2451                let start = il;
2452                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2453                    il += 1;
2454                }
2455                if !self.graphs.contains_key(&(start, t)) {
2456                    return false;
2457                }
2458            } else {
2459                il += 1;
2460            }
2461        }
2462        true
2463    }
2464
2465    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
2466    /// `row` (0-based) of layer `il`. None for non-linear layers.
2467    pub(crate) fn slab_row(
2468        &self,
2469        e: &Engine,
2470        il: usize,
2471        row: usize,
2472    ) -> Option<(u64, u64, usize, usize)> {
2473        use cudarc::driver::DevicePtr;
2474        let k = *self.lin_pos.get(&il)?;
2475        let s = &e.gpu.stream();
2476        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
2477        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
2478        Some((
2479            pc as u64 + (row * self.conv_words * 4) as u64,
2480            ps as u64 + (row * self.ssm_words * 4) as u64,
2481            self.conv_words,
2482            self.ssm_words,
2483        ))
2484    }
2485}
2486
2487impl VerifyCkpt {
2488    fn new(n_layer: usize) -> Self {
2489        VerifyCkpt {
2490            gdn: (0..n_layer).map(|_| None).collect(),
2491            cols: (0..n_layer).map(|_| None).collect(),
2492        }
2493    }
2494}
2495
2496/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
2497/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
2498/// a logical round number.
2499struct VerifyBoundaryTicket {
2500    rt: &'static crate::pp::PpNRt,
2501    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2502    slot: usize,
2503    pos0: usize,
2504    t: usize,
2505    payload: usize,
2506    n_st: usize,
2507    pipelined: bool,
2508    pp_anatomy: bool,
2509    pp_started: std::time::Instant,
2510    reverse_ms: f64,
2511    stage0_ms: f64,
2512    tx_ms: f64,
2513    trace: Option<SpecPipeTraceCtx>,
2514}
2515
2516/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
2517/// increment-2 controller can also be armed by the server's fresh-process research door.
2518#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2519pub enum OptiForkGateMode {
2520    Disabled,
2521    Hit,
2522    Miss,
2523    Alternate,
2524    Abort,
2525    Controller,
2526}
2527
2528static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
2529static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
2530    std::sync::atomic::AtomicU32::new(0);
2531static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2532static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2533static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2534static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2535static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2536static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2537static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2538static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2539static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2540static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2541    std::sync::atomic::AtomicU64::new(0);
2542static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2543    std::sync::atomic::AtomicU64::new(0);
2544static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2545
2546impl OptiForkGateMode {
2547    fn code(self) -> u8 {
2548        match self {
2549            Self::Disabled => 0,
2550            Self::Hit => 1,
2551            Self::Miss => 2,
2552            Self::Alternate => 3,
2553            Self::Abort => 4,
2554            Self::Controller => 5,
2555        }
2556    }
2557
2558    fn configured() -> Self {
2559        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2560            1 => Self::Hit,
2561            2 => Self::Miss,
2562            3 => Self::Alternate,
2563            4 => Self::Abort,
2564            5 => Self::Controller,
2565            _ => Self::Disabled,
2566        }
2567    }
2568
2569    fn action(self, generation: u64) -> OptiForkAction {
2570        match self {
2571            Self::Hit => OptiForkAction::Hit,
2572            Self::Miss => OptiForkAction::Miss,
2573            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2574            Self::Alternate => OptiForkAction::Miss,
2575            Self::Abort => OptiForkAction::Abort,
2576            Self::Disabled | Self::Controller => {
2577                unreachable!("non-forced mode cannot choose a forced fork action")
2578            }
2579        }
2580    }
2581
2582    fn is_forced(self) -> bool {
2583        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2584    }
2585}
2586
2587/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2588pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2589    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2590}
2591
2592/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2593/// two-token draft-probability product. Serving can call this only through its explicit
2594/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2595pub fn set_optipipe_controller_threshold(threshold: f32) {
2596    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2597    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2598    set_optipipe_gate_mode(OptiForkGateMode::Controller);
2599}
2600
2601#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2602pub struct OptiForkGateStats {
2603    pub attempts: u64,
2604    pub hits: u64,
2605    pub misses: u64,
2606    pub abort_drains: u64,
2607    pub refusals: u64,
2608    pub gate_checks: u64,
2609    pub gate_admits: u64,
2610    pub gate_rejects: u64,
2611    pub reconciles: u64,
2612    pub wasted_draft_tokens: u64,
2613    pub shadow_draft_tokens: u64,
2614    pub breaker_trips: u64,
2615}
2616
2617#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2618pub struct OptiForkStateIdentity {
2619    pub trunk_kv_bytes: usize,
2620    pub recurrent_bytes: usize,
2621    pub scratch_kv_bytes: usize,
2622    pub hidden_bytes: usize,
2623}
2624
2625pub fn reset_optipipe_gate_stats() {
2626    for counter in [
2627        &OPTI_FORK_ATTEMPTS,
2628        &OPTI_FORK_HITS,
2629        &OPTI_FORK_MISSES,
2630        &OPTI_FORK_ABORT_DRAINS,
2631        &OPTI_FORK_REFUSALS,
2632        &OPTI_GATE_CHECKS,
2633        &OPTI_GATE_ADMITS,
2634        &OPTI_GATE_REJECTS,
2635        &OPTI_RECONCILES,
2636        &OPTI_WASTED_DRAFT_TOKENS,
2637        &OPTI_SHADOW_DRAFT_TOKENS,
2638        &OPTI_BREAKER_TRIPS,
2639    ] {
2640        counter.store(0, std::sync::atomic::Ordering::Relaxed);
2641    }
2642}
2643
2644pub fn optipipe_gate_stats() -> OptiForkGateStats {
2645    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2646    OptiForkGateStats {
2647        attempts: load(&OPTI_FORK_ATTEMPTS),
2648        hits: load(&OPTI_FORK_HITS),
2649        misses: load(&OPTI_FORK_MISSES),
2650        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2651        refusals: load(&OPTI_FORK_REFUSALS),
2652        gate_checks: load(&OPTI_GATE_CHECKS),
2653        gate_admits: load(&OPTI_GATE_ADMITS),
2654        gate_rejects: load(&OPTI_GATE_REJECTS),
2655        reconciles: load(&OPTI_RECONCILES),
2656        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2657        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2658        breaker_trips: load(&OPTI_BREAKER_TRIPS),
2659    }
2660}
2661
2662#[derive(Clone, Copy, Debug)]
2663struct OptiControllerPolicy {
2664    threshold: f32,
2665    consecutive_misses: u8,
2666    breaker_tripped: bool,
2667}
2668
2669impl OptiControllerPolicy {
2670    fn configured() -> Self {
2671        Self {
2672            threshold: f32::from_bits(
2673                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2674            ),
2675            consecutive_misses: 0,
2676            breaker_tripped: false,
2677        }
2678    }
2679
2680    fn admit(&self, q_proxy: f32) -> bool {
2681        q_proxy.is_finite()
2682            && (0.0..=1.0).contains(&q_proxy)
2683            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2684    }
2685
2686    /// Returns true exactly when this resolution newly trips the three-miss breaker.
2687    fn resolve(&mut self, hit: bool) -> bool {
2688        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2689        // every optimistic opportunity, so the safety breaker is measured separately and must
2690        // not silently turn this arm into "three attempts then serial".
2691        if self.threshold == 0.0 {
2692            self.consecutive_misses = 0;
2693            return false;
2694        }
2695        if hit {
2696            self.consecutive_misses = 0;
2697            return false;
2698        }
2699        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2700        if !self.breaker_tripped && self.consecutive_misses >= 3 {
2701            self.breaker_tripped = true;
2702            return true;
2703        }
2704        false
2705    }
2706}
2707
2708#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2709enum OptiForkAction {
2710    Hit,
2711    Miss,
2712    Abort,
2713}
2714
2715#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2716struct OptiForkGeneration {
2717    id: u64,
2718    slot: usize,
2719}
2720
2721#[derive(Default)]
2722struct OptiForkGenerationTracker {
2723    next: u64,
2724    live: [Option<u64>; 2],
2725}
2726
2727impl OptiForkGenerationTracker {
2728    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2729        let generation = OptiForkGeneration {
2730            id: self.next,
2731            slot: (self.next & 1) as usize,
2732        };
2733        if let Some(live) = self.live[generation.slot] {
2734            return Err(format!(
2735                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2736                generation.slot,
2737            )
2738            .into());
2739        }
2740        self.next += 1;
2741        self.live[generation.slot] = Some(generation.id);
2742        Ok(generation)
2743    }
2744
2745    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2746        match self.live[generation.slot] {
2747            Some(id) if id == generation.id => {
2748                self.live[generation.slot] = None;
2749                Ok(())
2750            }
2751            other => Err(format!(
2752                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2753                generation.id, generation.slot,
2754            )
2755            .into()),
2756        }
2757    }
2758}
2759
2760struct OptiForkSeedGeneration {
2761    h_seed: CudaSlice<f32>,
2762    fill_prev: CudaSlice<f32>,
2763    scratch_len: usize,
2764}
2765
2766/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2767/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2768/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2769/// device ownership.
2770fn opti_snapshot_stage_owned(
2771    e: &Engine,
2772    cache: &Cache,
2773    rt: &'static crate::pp::PpNRt,
2774    fence: &[usize],
2775) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2776    let n = cache.kv.len();
2777    let mut snapshot = crate::cache::CacheSnapshot {
2778        kv_len: vec![None; n],
2779        conv: (0..n).map(|_| None).collect(),
2780        ssm: (0..n).map(|_| None).collect(),
2781        pos: cache.pos,
2782    };
2783    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
2784    Ok(snapshot)
2785}
2786
2787fn opti_snapshot_stage_owned_into(
2788    e: &Engine,
2789    cache: &Cache,
2790    rt: &'static crate::pp::PpNRt,
2791    fence: &[usize],
2792    snapshot: &mut crate::cache::CacheSnapshot,
2793) -> Result<(), Box<dyn std::error::Error>> {
2794    if fence.len() != rt.n_stages() + 1 || snapshot.kv_len.len() != cache.kv.len() {
2795        return Err("optipipe stage-owned snapshot shape mismatch".into());
2796    }
2797    for stage in 0..rt.n_stages() {
2798        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
2799    }
2800    snapshot.pos = cache.pos;
2801    Ok(())
2802}
2803
2804/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
2805/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
2806/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
2807/// either point would capture one side of the fork at the wrong generation.
2808fn opti_snapshot_one_stage_owned_into(
2809    e: &Engine,
2810    cache: &Cache,
2811    rt: &'static crate::pp::PpNRt,
2812    fence: &[usize],
2813    stage: usize,
2814    snapshot: &mut crate::cache::CacheSnapshot,
2815) -> Result<(), Box<dyn std::error::Error>> {
2816    if fence.len() != rt.n_stages() + 1
2817        || snapshot.kv_len.len() != cache.kv.len()
2818        || stage >= rt.n_stages()
2819    {
2820        return Err("optipipe single-stage snapshot shape mismatch".into());
2821    }
2822    let _scope = rt.enter(stage);
2823    let owner = rt.engine(stage, e);
2824    for il in fence[stage]..fence[stage + 1] {
2825        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
2826        match &cache.recur[il] {
2827            Some(recur) => {
2828                match snapshot.conv[il].as_mut() {
2829                    Some(dst) => {
2830                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
2831                    }
2832                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
2833                }
2834                match snapshot.ssm[il].as_mut() {
2835                    Some(dst) => {
2836                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
2837                    }
2838                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
2839                }
2840            }
2841            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
2842                return Err(
2843                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
2844                );
2845            }
2846            None => {}
2847        }
2848    }
2849    snapshot.pos = cache.pos;
2850    Ok(())
2851}
2852
2853/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
2854/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
2855/// resolve, so the reconcile tables and conditional restores are stage-local.
2856struct OptiForkState {
2857    mode: OptiForkGateMode,
2858    controller: Option<OptiControllerPolicy>,
2859    generations: OptiForkGenerationTracker,
2860    active_snapshot_slot: usize,
2861    alternate_snapshot: crate::cache::CacheSnapshot,
2862    seeds: [OptiForkSeedGeneration; 2],
2863    rt: &'static crate::pp::PpNRt,
2864    fence: [usize; 3],
2865    split: usize,
2866    len_ptrs: CudaSlice<u64>,
2867    saved_lens: CudaSlice<i32>,
2868    forced_acc: CudaSlice<u32>,
2869    valid: CudaSlice<u32>,
2870    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2871    logical_payload_bytes: [usize; 2],
2872}
2873
2874struct OptiForkTicket {
2875    generation: OptiForkGeneration,
2876    boundary: Option<VerifyBoundaryTicket>,
2877    drain: std::sync::Arc<cudarc::driver::CudaStream>,
2878    settled: bool,
2879}
2880
2881struct OptiControllerTicket {
2882    generation: OptiForkGeneration,
2883    boundary: Option<VerifyBoundaryTicket>,
2884    ckpt: Option<VerifyCkpt>,
2885    verify_tokens: [u32; 2],
2886    draft_prob: f32,
2887    eager_seed: Option<CudaSlice<f32>>,
2888    q_proxy: f32,
2889    scratch_len: usize,
2890    issued_at: std::time::Instant,
2891    drain: std::sync::Arc<cudarc::driver::CudaStream>,
2892    settled: bool,
2893}
2894
2895struct OptiControllerPrepared {
2896    verify_tokens: [u32; 2],
2897    draft_prob: f32,
2898    eager_seed: Option<CudaSlice<f32>>,
2899    q_proxy: f32,
2900    scratch_len: usize,
2901}
2902
2903impl OptiControllerTicket {
2904    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2905        self.boundary
2906            .take()
2907            .expect("controller boundary ticket already consumed")
2908    }
2909
2910    fn take_ckpt(&mut self) -> VerifyCkpt {
2911        self.ckpt
2912            .take()
2913            .expect("controller verify checkpoint already consumed")
2914    }
2915
2916    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
2917        self.eager_seed.take()
2918    }
2919
2920    fn settle(&mut self) {
2921        self.settled = true;
2922    }
2923}
2924
2925impl Drop for OptiControllerTicket {
2926    fn drop(&mut self) {
2927        if !self.settled {
2928            let _ = self.drain.synchronize();
2929            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2930        }
2931    }
2932}
2933
2934impl OptiForkTicket {
2935    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
2936        self.boundary
2937            .take()
2938            .expect("fork ticket boundary already consumed")
2939    }
2940
2941    fn settle(&mut self) {
2942        self.settled = true;
2943    }
2944}
2945
2946impl Drop for OptiForkTicket {
2947    fn drop(&mut self) {
2948        if !self.settled {
2949            let _ = self.drain.synchronize();
2950            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2951        }
2952    }
2953}
2954
2955impl OptiForkState {
2956    #[allow(clippy::too_many_arguments)]
2957    fn new(
2958        e: &Engine,
2959        cache: &Cache,
2960        mode: OptiForkGateMode,
2961        alternate_snapshot: crate::cache::CacheSnapshot,
2962        h_seed: &CudaSlice<f32>,
2963        fill_prev: &CudaSlice<f32>,
2964        rt: &'static crate::pp::PpNRt,
2965        split: usize,
2966        n_layer: usize,
2967    ) -> Result<Self, Box<dyn std::error::Error>> {
2968        let fence = [0, split, n_layer];
2969        let mut logical_payload_bytes = [0usize; 2];
2970        for stage in 0..2 {
2971            for il in fence[stage]..fence[stage + 1] {
2972                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
2973                    .as_ref()
2974                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2975                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
2976                    .as_ref()
2977                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
2978            }
2979        }
2980        let seeds = [
2981            OptiForkSeedGeneration {
2982                h_seed: e.clone_dtod(h_seed)?,
2983                fill_prev: e.clone_dtod(fill_prev)?,
2984                scratch_len: 0,
2985            },
2986            OptiForkSeedGeneration {
2987                h_seed: e.clone_dtod(h_seed)?,
2988                fill_prev: e.clone_dtod(fill_prev)?,
2989                scratch_len: 0,
2990            },
2991        ];
2992        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
2993            let _stage = rt.enter(0);
2994            let e0 = rt.engine(0, e);
2995            (
2996                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
2997                e0.htod_i32(&vec![0; split])?,
2998                e0.alloc_u32_zeroed(2)?,
2999                e0.alloc_u32_zeroed(1)?,
3000                e0.stream(),
3001            )
3002        };
3003        logical_payload_bytes[0] += seeds
3004            .iter()
3005            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3006            .sum::<usize>();
3007        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3008            + saved_lens.len() * std::mem::size_of::<i32>()
3009            + forced_acc.len() * std::mem::size_of::<u32>()
3010            + valid.len() * std::mem::size_of::<u32>();
3011        Ok(Self {
3012            mode,
3013            controller: (mode == OptiForkGateMode::Controller)
3014                .then(OptiControllerPolicy::configured),
3015            generations: OptiForkGenerationTracker::default(),
3016            active_snapshot_slot: 0,
3017            alternate_snapshot,
3018            seeds,
3019            rt,
3020            fence,
3021            split,
3022            len_ptrs,
3023            saved_lens,
3024            forced_acc,
3025            valid,
3026            stage0_stream,
3027            logical_payload_bytes,
3028        })
3029    }
3030
3031    fn reserve(
3032        &mut self,
3033        current_snapshot: &mut crate::cache::CacheSnapshot,
3034    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3035        let generation = self.generations.reserve()?;
3036        if generation.slot != self.active_snapshot_slot {
3037            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3038            self.active_snapshot_slot = generation.slot;
3039        }
3040        Ok(generation)
3041    }
3042
3043    fn capture_seed(
3044        &mut self,
3045        e: &Engine,
3046        generation: OptiForkGeneration,
3047        h_seed: &CudaSlice<f32>,
3048        fill_prev: &CudaSlice<f32>,
3049        scratch_len: usize,
3050    ) -> Result<(), Box<dyn std::error::Error>> {
3051        let seed = &mut self.seeds[generation.slot];
3052        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3053        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3054        seed.scratch_len = scratch_len;
3055        Ok(())
3056    }
3057
3058    fn ticket(
3059        &self,
3060        generation: OptiForkGeneration,
3061        boundary: VerifyBoundaryTicket,
3062    ) -> OptiForkTicket {
3063        OptiForkTicket {
3064            generation,
3065            boundary: Some(boundary),
3066            drain: self.stage0_stream.clone(),
3067            settled: false,
3068        }
3069    }
3070
3071    #[allow(clippy::too_many_arguments)]
3072    fn controller_ticket(
3073        &self,
3074        generation: OptiForkGeneration,
3075        boundary: VerifyBoundaryTicket,
3076        ckpt: VerifyCkpt,
3077        verify_tokens: [u32; 2],
3078        draft_prob: f32,
3079        eager_seed: Option<CudaSlice<f32>>,
3080        q_proxy: f32,
3081        scratch_len: usize,
3082    ) -> OptiControllerTicket {
3083        OptiControllerTicket {
3084            generation,
3085            boundary: Some(boundary),
3086            ckpt: Some(ckpt),
3087            verify_tokens,
3088            draft_prob,
3089            eager_seed,
3090            q_proxy,
3091            scratch_len,
3092            issued_at: std::time::Instant::now(),
3093            drain: self.stage0_stream.clone(),
3094            settled: false,
3095        }
3096    }
3097
3098    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3099        self.generations.reserve()
3100    }
3101
3102    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3103        &mut self.alternate_snapshot
3104    }
3105
3106    fn promote_successor_snapshot(
3107        &mut self,
3108        current_snapshot: &mut crate::cache::CacheSnapshot,
3109        generation: OptiForkGeneration,
3110    ) {
3111        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3112        self.active_snapshot_slot = generation.slot;
3113    }
3114
3115    fn queue_actual_reconcile(
3116        &mut self,
3117        e: &Engine,
3118        snapshot: &crate::cache::CacheSnapshot,
3119        acc: &CudaSlice<u32>,
3120        optimistic_pending: u32,
3121        base: usize,
3122    ) -> Result<(), Box<dyn std::error::Error>> {
3123        let saved: Vec<i32> = (0..self.split)
3124            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3125            .collect();
3126        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3127        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3128        // the validity/reconcile kernels must never peer-read acc before it is written. The
3129        // increment-1 harness uses primary stage 0, where stream order already provides this.
3130        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3131            self.rt.fence_stages_behind(&e.stream())?;
3132        }
3133        let _stage = self.rt.enter(0);
3134        let e0 = self.rt.engine(0, e);
3135        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3136        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3137        e0.spec_fork_reconcile_kv(
3138            &self.len_ptrs,
3139            &self.saved_lens,
3140            acc,
3141            &self.valid,
3142            base,
3143            self.split,
3144        )
3145    }
3146
3147    fn finish_actual_reconcile(
3148        &mut self,
3149        e: &Engine,
3150        cache: &mut Cache,
3151        snapshot: &crate::cache::CacheSnapshot,
3152        n_acc: usize,
3153        base: usize,
3154        hit: bool,
3155    ) -> Result<(), Box<dyn std::error::Error>> {
3156        if hit {
3157            return Ok(());
3158        }
3159        let len_delta = base + n_acc;
3160        for il in 0..self.split {
3161            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3162                kv.len = saved + len_delta;
3163            }
3164        }
3165        {
3166            let _stage = self.rt.enter(1);
3167            let e1 = self.rt.engine(1, e);
3168            for il in self.split..self.fence[2] {
3169                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3170                    kv.len = saved + len_delta;
3171                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3172                }
3173            }
3174        }
3175        self.rt.publish_to(0, &e.stream())?;
3176        Ok(())
3177    }
3178
3179    fn cancel_controller_ticket(
3180        &mut self,
3181        e: &Engine,
3182        cache: &mut Cache,
3183        scratch: &mut MtpScratch,
3184        snapshot: &crate::cache::CacheSnapshot,
3185        ticket: &mut OptiControllerTicket,
3186    ) -> Result<(), Box<dyn std::error::Error>> {
3187        {
3188            let _stage = self.rt.enter(0);
3189            let e0 = self.rt.engine(0, e);
3190            for il in 0..self.split {
3191                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3192                    kv.len = saved;
3193                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3194                }
3195            }
3196        }
3197        scratch.set_len(e, snapshot.pos)?;
3198        ticket.settle();
3199        self.generations.retire(ticket.generation)?;
3200        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3201        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3202        eprintln!(
3203            "[opti-controller] tail-drain generation={} slot={}",
3204            ticket.generation.id, ticket.generation.slot,
3205        );
3206        Ok(())
3207    }
3208
3209    #[allow(clippy::too_many_arguments)]
3210    fn reconcile(
3211        &mut self,
3212        e: &Engine,
3213        cache: &mut Cache,
3214        scratch: &mut MtpScratch,
3215        snapshot: &crate::cache::CacheSnapshot,
3216        h_seed: &mut CudaSlice<f32>,
3217        fill_prev: &mut CudaSlice<f32>,
3218        generation: OptiForkGeneration,
3219        action: OptiForkAction,
3220        optimistic_pending: u32,
3221    ) -> Result<(), Box<dyn std::error::Error>> {
3222        debug_assert!(action != OptiForkAction::Abort);
3223        let miss_started = std::time::Instant::now();
3224        let keep = action == OptiForkAction::Hit;
3225        let saved: Vec<i32> = (0..self.split)
3226            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3227            .collect();
3228        let seed = &self.seeds[generation.slot];
3229        {
3230            let _stage = self.rt.enter(0);
3231            let e0 = self.rt.engine(0, e);
3232            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3233            let forced = if keep {
3234                [1u32, optimistic_pending]
3235            } else {
3236                [0u32, optimistic_pending]
3237            };
3238            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3239            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3240            e0.spec_fork_reconcile_kv(
3241                &self.len_ptrs,
3242                &self.saved_lens,
3243                &self.forced_acc,
3244                &self.valid,
3245                0,
3246                self.split,
3247            )?;
3248            for il in 0..self.split {
3249                if let Some(recur) = cache.recur[il].as_mut() {
3250                    let conv = snapshot.conv[il]
3251                        .as_ref()
3252                        .ok_or("optipipe stage0 snapshot missing conv state")?;
3253                    let ssm = snapshot.ssm[il]
3254                        .as_ref()
3255                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
3256                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3257                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3258                }
3259            }
3260            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3261            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3262        }
3263
3264        if keep {
3265            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3266            return Ok(());
3267        }
3268
3269        for il in 0..self.split {
3270            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3271                kv.len = saved;
3272            }
3273        }
3274        scratch.set_len(e, seed.scratch_len)?;
3275        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3276        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3277        let caller = e.stream();
3278        self.rt.publish_to(0, &caller)?;
3279        caller.synchronize()?;
3280        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3281        eprintln!(
3282            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3283            generation.id, generation.slot,
3284        );
3285        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3286        Ok(())
3287    }
3288
3289    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3290        self.generations.retire(generation)
3291    }
3292}
3293
3294impl HybridModel {
3295    fn opti_graph_draft_step(
3296        &self,
3297        e: &Engine,
3298        mtp: &MtpHead,
3299        dctx: &mut DraftGraphCtx,
3300        scratch: &mut MtpScratch,
3301        d_vocab: usize,
3302    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3303        dctx.graph
3304            .as_ref()
3305            .ok_or("optipipe controller requires the greedy draft graph")?
3306            .launch()?;
3307        scratch.kv.len += 1;
3308        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
3309        if (idx as usize) >= d_vocab {
3310            return Err(
3311                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
3312            );
3313        }
3314        let probability = e.dtoh(&dctx.g_p)?[0];
3315        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3316            return Err(format!("optipipe draft probability is invalid: {probability}").into());
3317        }
3318        let token = match &mtp.d2t {
3319            Some(map) => map[idx as usize],
3320            None => idx,
3321        };
3322        if token != idx {
3323            e.set_u32_one(&mut dctx.g_tok, token)?;
3324        }
3325        Ok((token, probability))
3326    }
3327
3328    #[allow(clippy::too_many_arguments)]
3329    fn opti_controller_draft_step(
3330        &self,
3331        e: &Engine,
3332        mtp: &MtpHead,
3333        dctx: &mut DraftGraphCtx,
3334        scratch: &mut MtpScratch,
3335        d_vocab: usize,
3336        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
3337        eager_pos: usize,
3338        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3339    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3340        if dctx.graph.is_some() {
3341            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
3342        }
3343        let (input_token, input_seed) = eager_state
3344            .take()
3345            .ok_or("optipipe eager continuation seed is unavailable")?;
3346        let (logits, next_seed) = self.mtp_head_forward_dev(
3347            e,
3348            mtp,
3349            input_token,
3350            &input_seed,
3351            scratch,
3352            eager_pos,
3353            embd_dev,
3354            None,
3355        )?;
3356        let token_d = e.argmax_token_device(&logits, d_vocab)?;
3357        let idx = e.dtoh_u32_one(&token_d)?;
3358        if (idx as usize) >= d_vocab {
3359            return Err(format!(
3360                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
3361            )
3362            .into());
3363        }
3364        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
3365        let probability = e.dtoh(&probability_d)?[0];
3366        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3367            return Err(
3368                format!("optipipe eager draft probability is invalid: {probability}").into(),
3369            );
3370        }
3371        let token = match &mtp.d2t {
3372            Some(map) => map[idx as usize],
3373            None => idx,
3374        };
3375        *eager_state = Some((token, next_seed));
3376        Ok((token, probability))
3377    }
3378
3379    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3380    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3381    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3382    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3383    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3384    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3385    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3386    /// transfer + host argmax per draft token from the K-token draft chain.
3387    #[allow(clippy::too_many_arguments)]
3388    fn mtp_head_forward_dev(
3389        &self,
3390        e: &Engine,
3391        mtp: &MtpHead,
3392        e_tok: u32,
3393        h_seed: &CudaSlice<f32>,
3394        scratch: &mut MtpScratch,
3395        mtp_pos: usize,
3396        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3397        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3398        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3399        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3400        mask: Option<(&CudaSlice<u32>, usize)>,
3401    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3402        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3403        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3404        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3405        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3406        static ANAT_NS: [AtomicU64; 5] = [
3407            AtomicU64::new(0),
3408            AtomicU64::new(0),
3409            AtomicU64::new(0),
3410            AtomicU64::new(0),
3411            AtomicU64::new(0),
3412        ];
3413        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3414        let anat = {
3415            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3416            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3417        };
3418        if anat {
3419            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3420        }
3421        let t_all = std::time::Instant::now();
3422        let mut t_ph = std::time::Instant::now();
3423        let mut anat_mark = |i: usize,
3424                             e: &Engine,
3425                             t: &mut std::time::Instant|
3426         -> Result<(), Box<dyn std::error::Error>> {
3427            if anat {
3428                e.stream().synchronize()?;
3429                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3430                *t = std::time::Instant::now();
3431            }
3432            Ok(())
3433        };
3434        let cfg = &self.cfg;
3435        let n_embd = cfg.n_embd as usize;
3436        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3437        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3438        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3439        let eps = cfg.rms_eps;
3440        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3441
3442        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3443        // expands this one row on CPU and transfers n_embd f32 values instead.
3444        let e_emb = match embd_dev {
3445            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3446            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
3447        };
3448
3449        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3450        let mut e_norm = e.zeros(n_embd)?;
3451        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3452        let mut h_norm = e.zeros(n_embd)?;
3453        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3454
3455        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3456        let mut concat = e.zeros(2 * n_embd)?;
3457        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3458        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3459
3460        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3461        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3462
3463        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3464        let mut a_norm = e.zeros(di)?;
3465        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3466        anat_mark(0, e, &mut t_ph)?;
3467
3468        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3469        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3470        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3471        // advances only the device counter).
3472        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3473            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3474            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
3475            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
3476            // whose host-side mirror the caller does).
3477            (Mixer::Full(fa), Some(g)) => {
3478                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch)?
3479            }
3480            (Mixer::Full(fa), None) => {
3481                let out =
3482                    self.mtp_full_attn_dc(e, fa, &a_norm, &pos_d, scratch, mtp.geom.as_ref())?;
3483                scratch.kv.len += 1;
3484                out
3485            }
3486            (Mixer::Linear(_), _) => {
3487                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3488            }
3489            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
3490        };
3491        anat_mark(1, e, &mut t_ph)?;
3492
3493        // op 7: x1 = inpSA + attn_out
3494        let mut x1 = e.zeros(di)?;
3495        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3496
3497        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
3498        let mut z = e.zeros(di)?;
3499        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3500
3501        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3502        let ffn_out = match &mtp.ffn {
3503            crate::hybrid::Ffn::Dense {
3504                ffn_gate,
3505                ffn_up,
3506                ffn_down,
3507            } => {
3508                let n_ff = ffn_gate.out_features();
3509                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3510                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3511                    (
3512                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3513                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3514                    )
3515                } else {
3516                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3517                };
3518                let mut act = e.zeros(n_ff)?;
3519                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3520                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3521                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3522                // passes None, which is `ffn_act`'s dispatch verbatim.
3523                Self::ffn_act_lim(
3524                    e,
3525                    &self.cfg,
3526                    &gate,
3527                    &up,
3528                    1.0,
3529                    1.0,
3530                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
3531                    &mut act,
3532                    n_ff,
3533                )?;
3534                e.matmul(ffn_down, &act, 1)?
3535            }
3536            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3537            // so they never alias trunk layer 0's cache keys.
3538            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3539        };
3540        anat_mark(2, e, &mut t_ph)?;
3541
3542        // op 10: h_nextn = x1 + ffn_out (at di)
3543        let mut h_inner = e.zeros(di)?;
3544        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3545
3546        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3547        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3548        let h_nextn = match mtp.geom.as_ref() {
3549            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3550            None => h_inner,
3551        };
3552
3553        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3554        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3555        let mut final_h = e.zeros(n_embd)?;
3556        e.rms_norm(
3557            &h_nextn,
3558            final_norm.float_data(),
3559            &mut final_h,
3560            n_embd,
3561            1,
3562            eps,
3563        )?;
3564
3565        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3566        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3567        let mut logits = e.matmul(head, &final_h, 1)?;
3568        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3569        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3570        if let Some((mask_d, mw)) = mask {
3571            let d_vocab = head.out_features();
3572            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3573        }
3574        anat_mark(3, e, &mut t_ph)?;
3575        if anat {
3576            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3577            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3578            if n % 128 == 0 {
3579                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3580                eprintln!(
3581                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3582                    us(0),
3583                    us(1),
3584                    us(2),
3585                    us(3),
3586                    us(4)
3587                );
3588            }
3589        }
3590        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3591        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3592        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3593    }
3594
3595    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3596    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3597    /// the dc path, and all three are properties of this arch's MTP block:
3598    ///
3599    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3600    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3601    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3602    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3603    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3604    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3605    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
3606    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3607    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3608    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3609    ///    resolved `Step35MtpGeom`, never from `cfg`.
3610    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3611    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3612    ///    fused-into-wq `q_gate_split` form the dc arm handles.
3613    ///
3614    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3615    /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3616    /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3617    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3618    ///
3619    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3620    /// caller must not mirror.
3621    fn mtp_step35_attn(
3622        &self,
3623        e: &Engine,
3624        fa: &FullAttnLayer,
3625        g: &crate::hybrid::Step35MtpGeom,
3626        h: &CudaSlice<f32>,
3627        pos_d: &CudaSlice<i32>,
3628        scratch: &mut MtpScratch,
3629    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3630        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3631        let eps = self.cfg.rms_eps;
3632        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3633        let n_embd = self.cfg.n_embd as usize;
3634        let gw = fa
3635            .attn_gate
3636            .as_ref()
3637            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3638
3639        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3640            && e.uses_q8_1_fast(&fa.wk)
3641            && e.uses_q8_1_fast(&fa.wv)
3642            && e.uses_q8_1_fast(gw)
3643        {
3644            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3645            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3646                Some(t3) => t3,
3647                None => (
3648                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3649                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3650                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3651                ),
3652            };
3653            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3654        } else {
3655            (
3656                e.matmul(&fa.wq, h, 1)?,
3657                e.matmul(&fa.wk, h, 1)?,
3658                e.matmul(&fa.wv, h, 1)?,
3659                e.matmul(gw, h, 1)?,
3660            )
3661        };
3662
3663        let mut q = e.uninit(nh * hd)?;
3664        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3665        let mut k = e.uninit(nkv * hd)?;
3666        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3667        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3668        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3669        // the resolved flag, not the constant, so an all-full sibling stays correct.
3670        let ff = if g.swa {
3671            None
3672        } else {
3673            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3674        };
3675        #[cfg(debug_assertions)]
3676        if let Some(ff) = ff {
3677            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3678        }
3679        e.rope_neox2(
3680            &mut q,
3681            &mut k,
3682            pos_d,
3683            hd,
3684            g.n_rot,
3685            nh,
3686            nkv,
3687            1,
3688            g.rope_base,
3689            1.0,
3690            ff,
3691        )?;
3692
3693        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3694        // length on the host anyway, and the windowed view below needs it there to compute the
3695        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3696        // dc-family consumer of this scratch still agree.
3697        let kv = &mut scratch.kv;
3698        assert!(
3699            kv.len < scratch.cap,
3700            "step35 MTP scratch overflow ({} >= {})",
3701            kv.len,
3702            scratch.cap
3703        );
3704        let next_len = kv.len + 1;
3705        let (off, t_kv) = if g.swa && next_len > g.window {
3706            (next_len - g.window, g.window)
3707        } else {
3708            (0, next_len)
3709        };
3710        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
3711        e.append_kv_quantized(
3712            &k,
3713            &v0,
3714            &mut kv.k,
3715            &mut kv.v,
3716            write_row,
3717            kv.kv_dim_k,
3718            kv.kv_dim_v,
3719            kv.k_tok_bytes,
3720            kv.v_tok_bytes,
3721            false,
3722        )?;
3723        kv.len = next_len;
3724        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3725        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
3726        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
3727        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
3728        // therefore live, not theoretical.
3729        let physical = kv.physical_rows(off, off + t_kv)?;
3730        let k_view = e.view_u8_range(
3731            &kv.k,
3732            physical.start * kv.k_tok_bytes,
3733            physical.end * kv.k_tok_bytes,
3734        );
3735        let v_view = e.view_u8_range(
3736            &kv.v,
3737            physical.start * kv.v_tok_bytes,
3738            physical.end * kv.v_tok_bytes,
3739        );
3740        let mut attn = e.uninit(nh * hd)?;
3741        e.fa_decode_kvmod(
3742            &q,
3743            &k_view,
3744            &v_view,
3745            &mut attn,
3746            hd,
3747            nh,
3748            nkv,
3749            t_kv,
3750            scale,
3751            kv.k_tok_bytes,
3752            kv.v_tok_bytes,
3753            false,
3754        )?;
3755
3756        let mut ag = e.uninit(nh * hd)?;
3757        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
3758        Ok(e.matmul(&fa.wo, &ag, 1)?)
3759    }
3760
3761    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
3762    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
3763    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
3764    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
3765    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
3766    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
3767    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
3768    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
3769    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
3770    fn mtp_full_attn_dc(
3771        &self,
3772        e: &Engine,
3773        fa: &FullAttnLayer,
3774        h: &CudaSlice<f32>,
3775        pos_d: &CudaSlice<i32>,
3776        scratch: &mut MtpScratch,
3777        geom: Option<&crate::hybrid::DraftGeom>,
3778    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3779        let cfg = &self.cfg;
3780        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3781        let geometry = cfg.full_attention_geometry_at(mtp_il);
3782        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
3783        let n_head_kv = geom
3784            .map(|g| g.n_head_kv)
3785            .unwrap_or(geometry.n_head_kv as usize);
3786        let head_dim = geometry.head_dim_k as usize;
3787        let eps = cfg.rms_eps;
3788        let scale = geometry.attention_scale();
3789        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
3790        let bucket_max = scratch.cap; // < 96 guaranteed by the graph_draft eligibility gate
3791
3792        let (qf, mut k, v) =
3793            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
3794                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
3795                (
3796                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
3797                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
3798                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
3799                )
3800            } else {
3801                (
3802                    e.matmul(&fa.wq, h, 1)?,
3803                    e.matmul(&fa.wk, h, 1)?,
3804                    e.matmul(&fa.wv, h, 1)?,
3805                )
3806            };
3807        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
3808        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3809        let (mut q, gate) = if gated {
3810            let mut q = e.zeros(n_head * head_dim)?;
3811            let mut gate = e.zeros(n_head * head_dim)?;
3812            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
3813            (q, Some(gate))
3814        } else {
3815            (qf, None)
3816        };
3817
3818        let mut qn = e.zeros(n_head * head_dim)?;
3819        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
3820        q = qn;
3821        let mut kn = e.zeros(n_head_kv * head_dim)?;
3822        e.rms_norm(
3823            &k,
3824            fa.k_norm.float_data(),
3825            &mut kn,
3826            head_dim,
3827            n_head_kv,
3828            eps,
3829        )?;
3830        k = kn;
3831        let rope_dims = geometry.n_rot as usize;
3832        e.rope_neox(
3833            &mut q,
3834            pos_d,
3835            head_dim,
3836            rope_dims,
3837            n_head,
3838            1,
3839            geometry.rope_base,
3840            1.0,
3841        )?;
3842        e.rope_neox(
3843            &mut k,
3844            pos_d,
3845            head_dim,
3846            rope_dims,
3847            n_head_kv,
3848            1,
3849            geometry.rope_base,
3850            1.0,
3851        )?;
3852
3853        let kv = &mut scratch.kv;
3854        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
3855        e.append_kv_quantized_dc(
3856            &k,
3857            &v,
3858            &mut kv.k,
3859            &mut kv.v,
3860            &kv.len_d,
3861            kv.kv_dim_k,
3862            kv.kv_dim_v,
3863            kv.k_tok_bytes,
3864            kv.v_tok_bytes,
3865            false,
3866        )?;
3867        e.inc_seqlen(&mut kv.len_d)?;
3868        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
3869        // key range from the device counter.
3870        let k_view = e.view_u8(&kv.k, kv.k.len());
3871        let v_view = e.view_u8(&kv.v, kv.v.len());
3872        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
3873        let mut attn = e.zeros(n_head * head_dim)?;
3874        e.fa_decode_dc(
3875            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
3876            scale, ktb, vtb, false,
3877        )?;
3878
3879        let attn_g = match &gate {
3880            Some(gate) => {
3881                let mut gsig = e.zeros(n_head * head_dim)?;
3882                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
3883                let mut ag = e.zeros(n_head * head_dim)?;
3884                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
3885                ag
3886            }
3887            None => attn,
3888        };
3889        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
3890    }
3891
3892    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
3893    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
3894    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
3895    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
3896    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
3897    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
3898    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
3899    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
3900    #[allow(clippy::too_many_arguments)]
3901    fn mtp_kv_fill(
3902        &self,
3903        e: &Engine,
3904        mtp: &MtpHead,
3905        tokens: &[u32],
3906        h: &CudaSlice<f32>,
3907        pos0: usize,
3908        scratch: &mut MtpScratch,
3909        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3910    ) -> Result<(), Box<dyn std::error::Error>> {
3911        let cfg = &self.cfg;
3912        let n_embd = cfg.n_embd as usize;
3913        let eps = cfg.rms_eps;
3914        let t = tokens.len();
3915        assert_eq!(scratch.kv.len, pos0, "mtp_kv_fill: append slot mismatch");
3916        assert!(pos0 + t <= scratch.cap, "mtp_kv_fill: scratch overflow");
3917        let Mixer::Full(fa) = &mtp.mixer else {
3918            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3919        };
3920        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
3921        let pos_d = e.htod_i32(&pos_vec)?;
3922
3923        // ops A/1/2: embed + the two input norms, T-wide.
3924        let e_emb = match embd_dev {
3925            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
3926            None => e.htod(&self.embd.gather(n_embd, tokens))?,
3927        };
3928        let mut e_norm = e.zeros(t * n_embd)?;
3929        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
3930        let mut h_norm = e.zeros(t * n_embd)?;
3931        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
3932
3933        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
3934        let mut concat = e.zeros(t * 2 * n_embd)?;
3935        for i in 0..t {
3936            e.copy_view_into(
3937                &mut concat,
3938                i * 2 * n_embd,
3939                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
3940                n_embd,
3941            )?;
3942            e.copy_view_into(
3943                &mut concat,
3944                i * 2 * n_embd + n_embd,
3945                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
3946                n_embd,
3947            )?;
3948        }
3949
3950        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
3951        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3952        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
3953        let mut a_norm = e.zeros(t * di)?;
3954        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
3955
3956        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
3957        // the fill only has to leave correct K/V rows behind for later chains to attend over.
3958        let n_head_kv = mtp
3959            .geom
3960            .as_ref()
3961            .map(|g| g.n_head_kv)
3962            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
3963            .unwrap_or_else(|| {
3964                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3965                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
3966            });
3967        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
3968        let geometry = cfg.full_attention_geometry_at(mtp_il);
3969        let head_dim = geometry.head_dim_k as usize;
3970        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
3971        let v = e.matmul(&fa.wv, &a_norm, t)?;
3972        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
3973        e.rms_norm(
3974            &k,
3975            fa.k_norm.float_data(),
3976            &mut kn,
3977            head_dim,
3978            n_head_kv * t,
3979            eps,
3980        )?;
3981        k = kn;
3982        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
3983        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
3984        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
3985        // writes K rows the attention arm then re-derives at a different theta: correct-looking
3986        // output with dead acceptance, invisible to the exactness gates.
3987        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
3988            Some(s) => (
3989                s.n_rot,
3990                s.rope_base,
3991                if s.swa {
3992                    None
3993                } else {
3994                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3995                },
3996            ),
3997            None => (geometry.n_rot as usize, geometry.rope_base, None),
3998        };
3999        #[cfg(debug_assertions)]
4000        if let Some(ff) = ff {
4001            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4002        }
4003        match ff {
4004            Some(f) => e.rope_neox_ff(
4005                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4006            )?,
4007            None => e.rope_neox(
4008                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4009            )?,
4010        }
4011
4012        let kv = &mut scratch.kv;
4013        // Match the trunk prime contract: a chunk may need the aligned window immediately before
4014        // its first row, so preserve that prefix when the physical tail rebases at wrap.
4015        let retain_from = kv
4016            .ring
4017            .as_ref()
4018            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
4019            .unwrap_or(0);
4020        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4021        for i in 0..t {
4022            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4023            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4024            e.append_kv_quantized_view(
4025                &k_row,
4026                &v_row,
4027                &mut kv.k,
4028                &mut kv.v,
4029                write_row + i,
4030                kv.kv_dim_k,
4031                kv.kv_dim_v,
4032                kv.k_tok_bytes,
4033                kv.v_tok_bytes,
4034                false,
4035            )?;
4036        }
4037        kv.len = pos0 + t;
4038        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4039        Ok(())
4040    }
4041
4042    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4043    /// every varying input device-resident —
4044    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4045    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4046    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4047    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4048    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4049    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4050    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4051    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4052    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4053    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4054    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4055    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4056    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4057    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4058    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4059    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4060    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4061    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4062    #[allow(clippy::too_many_arguments)]
4063    fn mtp_head_forward_cap(
4064        &self,
4065        e: &Engine,
4066        mtp: &MtpHead,
4067        tok_d: &mut CudaSlice<u32>,
4068        pos_d: &mut CudaSlice<i32>,
4069        h_seed_d: &mut CudaSlice<f32>,
4070        p_d: &mut CudaSlice<f32>,
4071        scratch: &mut MtpScratch,
4072        with_prob: bool,
4073        with_head: bool,
4074        embd_gpu: &CudaSlice<u8>,
4075        embd_qt: i32,
4076        embd_rb: usize,
4077        d_vocab: usize,
4078        sampled_cap: Option<(
4079            &mut CudaSlice<u32>,
4080            &mut CudaSlice<f32>,
4081            &mut CudaSlice<f32>,
4082            u64,
4083            f32,
4084        )>,
4085        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4086        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4087        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4088        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4089        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4090        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4091        mask_cap: Option<(&CudaSlice<u32>, usize)>,
4092    ) -> Result<(), Box<dyn std::error::Error>> {
4093        let cfg = &self.cfg;
4094        let n_embd = cfg.n_embd as usize;
4095        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
4096        // whose device-counter key bound always starts at row 0 — it cannot express this block's
4097        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
4098        // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
4099        // refuses step35 heads explicitly (SWA refusal), so the eager chain
4100        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
4101        // panic) is what the two capture sites and the round-stream capture already handle by
4102        // degrading to eager / stream-off.
4103        if mtp.step35.is_some() {
4104            return Err(
4105                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4106                        block's SWA view offset; same root cause as the dc decode refusal) — the \
4107                        eager draft chain serves this arch"
4108                    .into(),
4109            );
4110        }
4111        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4112        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4113        let eps = cfg.rms_eps;
4114        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4115        let mut e_norm = e.zeros(n_embd)?;
4116        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4117        let mut h_norm = e.zeros(n_embd)?;
4118        e.rms_norm(
4119            &*h_seed_d,
4120            mtp.hnorm.float_data(),
4121            &mut h_norm,
4122            n_embd,
4123            1,
4124            eps,
4125        )?;
4126        let mut concat = e.zeros(2 * n_embd)?;
4127        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4128        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4129        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4130        let mut a_norm = e.zeros(di)?;
4131        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4132        let attn_out = match &mtp.mixer {
4133            Mixer::Full(fa) => {
4134                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, mtp.geom.as_ref())?
4135            }
4136            Mixer::Linear(_) => {
4137                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4138            }
4139            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4140        };
4141        let mut x1 = e.zeros(di)?;
4142        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4143        let mut z = e.zeros(di)?;
4144        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4145        let ffn_out = match &mtp.ffn {
4146            crate::hybrid::Ffn::Dense {
4147                ffn_gate,
4148                ffn_up,
4149                ffn_down,
4150            } => {
4151                let n_ff = ffn_gate.out_features();
4152                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4153                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4154                    (
4155                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4156                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4157                    )
4158                } else {
4159                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4160                };
4161                let mut act = e.zeros(n_ff)?;
4162                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
4163                e.matmul(ffn_down, &act, 1)?
4164            }
4165            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
4166            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
4167            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
4168            // error arm degrades the caller to eager/stream-off.
4169            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
4170                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4171            }
4172            crate::hybrid::Ffn::Moe(_) => {
4173                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
4174            }
4175        };
4176        let mut h_inner = e.zeros(di)?;
4177        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4178        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4179        let h_nextn = match mtp.geom.as_ref() {
4180            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4181            None => h_inner,
4182        };
4183        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4184        let final_h = if with_head || spec_hpost() {
4185            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4186            let mut fh = e.zeros(n_embd)?;
4187            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4188            Some(fh)
4189        } else {
4190            None
4191        };
4192        if with_head {
4193            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4194            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4195            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4196            // before the argmax — proposals become legal by construction. Contents-only
4197            // per-replay upload keeps the capture valid.
4198            if let Some((mask_d, mw)) = mask_cap {
4199                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4200            }
4201            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
4202                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4203                // own buffer is pool-recycled after the capture body returns, so it can't be the
4204                // retention target), bump the device event counter, gumbel-perturb reading it,
4205                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4206                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4207                e.sctr_inc(ctr_d)?;
4208                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
4209                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4210                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4211                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4212                if with_prob {
4213                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4214                }
4215            } else {
4216                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4217                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4218                // p-min under a draft mask reads the MASKED row: confidence relative to the
4219                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4220                // is the right semantics for "does the drafter know what comes next here" and
4221                // the same row the pick came from. Draft-quality only — verify arbitrates.
4222                if with_prob {
4223                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4224                }
4225            }
4226        }
4227        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4228        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4229        if let Some((out, slot, d2t)) = stream_pack {
4230            e.pack_tok_p(tok_d, p_d, out, slot)?;
4231            if let Some(map) = d2t {
4232                e.tok_map_u32(tok_d, map)?;
4233            }
4234        }
4235        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4236        if spec_hpost() {
4237            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4238        } else {
4239            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4240        }
4241        // advance the draft rope position in-graph.
4242        e.inc_seqlen(pos_d)?;
4243        Ok(())
4244    }
4245
4246    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4247    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4248    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4249    /// Advances `cache.pos` by T.
4250    pub fn decode_step_t(
4251        &self,
4252        e: &Engine,
4253        tokens: &[u32],
4254        pos0: usize,
4255        cache: &mut Cache,
4256    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4257        if self.is_gemma4_e4b() {
4258            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4259        }
4260        if self.cfg.gemma4.is_some() {
4261            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4262        }
4263        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4264    }
4265
4266    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4267    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4268    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4269    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4270    pub fn decode_step_t_h(
4271        &self,
4272        e: &Engine,
4273        tokens: &[u32],
4274        pos0: usize,
4275        cache: &mut Cache,
4276    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4277        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4278    }
4279
4280    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4281    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4282    pub fn decode_step_t_h_emb(
4283        &self,
4284        e: &Engine,
4285        tokens: &[u32],
4286        pos0: usize,
4287        cache: &mut Cache,
4288        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4289    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4290        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4291        Ok((e.dtoh(&logits_d)?, h_seed))
4292    }
4293
4294    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4295    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4296    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4297    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4298    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4299    pub fn decode_step_t_h_emb_dev(
4300        &self,
4301        e: &Engine,
4302        tokens: &[u32],
4303        pos0: usize,
4304        cache: &mut Cache,
4305        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4306    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4307        let n_embd = self.cfg.n_embd as usize;
4308        let t = tokens.len();
4309        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4310        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4311        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4312        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4313        Ok((logits, hs))
4314    }
4315
4316    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4317    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4318    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4319    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4320    /// retains/copies — they never change what any kernel computes).
4321    fn decode_step_t_core(
4322        &self,
4323        e: &Engine,
4324        tokens: &[u32],
4325        pos0: usize,
4326        cache: &mut Cache,
4327        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4328        mut ckpt: Option<&mut VerifyCkpt>,
4329    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4330        self.decode_step_t_core_stream(
4331            e,
4332            tokens,
4333            pos0,
4334            cache,
4335            embd_dev,
4336            ckpt.take(),
4337            None,
4338            None,
4339            None,
4340            None,
4341        )
4342    }
4343
4344    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
4345    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
4346    fn decode_step_t_core_pipelined(
4347        &self,
4348        e: &Engine,
4349        tokens: &[u32],
4350        pos0: usize,
4351        cache: &mut Cache,
4352        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4353        mut ckpt: Option<&mut VerifyCkpt>,
4354        pipe: &SpecPipeLane,
4355        round: usize,
4356    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4357        let fence = crate::pp::pp_cuts(self.layers.len())
4358            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
4359        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
4360            return Err("two-session speculative pipeline requires the PP verify split".into());
4361        }
4362        let interval_fence = pipe.stage0_begin(round)?;
4363        let ticket = self.verify_stage0_issue(
4364            e,
4365            tokens,
4366            pos0,
4367            cache,
4368            embd_dev,
4369            ckpt.as_deref_mut(),
4370            None,
4371            &fence,
4372            Some(interval_fence),
4373            pipe.trace(round),
4374        )?;
4375        pipe.stage0_end(round);
4376        pipe.stage1_begin(round)?;
4377        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
4378        pipe.verify_end(round);
4379        Ok(result)
4380    }
4381
4382    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4383    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4384    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4385    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4386    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4387    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4388    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4389    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4390    #[allow(clippy::too_many_arguments)]
4391    fn decode_step_t_core_stream(
4392        &self,
4393        e: &Engine,
4394        tokens: &[u32],
4395        pos0: usize,
4396        cache: &mut Cache,
4397        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4398        mut ckpt: Option<&mut VerifyCkpt>,
4399        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4400        pp_pipe: Option<bool>,
4401        vtok_dev: Option<&CudaSlice<u32>>,
4402        graphs: Option<&mut DsparkVerifyGraphs>,
4403    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4404        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4405        // exactly as the eager and batched steps do. This is the single funnel every verify
4406        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4407        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4408        // is untouched.
4409        //
4410        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4411        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4412        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4413        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4414        // or a placement whose PpNRt fails to build — so a config that would still walk the
4415        // whole trunk on one stream refuses instead of regressing 28x.
4416        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4417            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
4418                if vtok_dev.is_some() {
4419                    return Err(
4420                        "device-token dspark verify (slice-2 deferred readback) has no PP \
4421                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
4422                         route on one device"
4423                            .into(),
4424                    );
4425                }
4426                return self.decode_step_t_core_ppn(
4427                    e,
4428                    tokens,
4429                    pos0,
4430                    cache,
4431                    embd_dev,
4432                    ckpt.take(),
4433                    stream,
4434                    &fence,
4435                    pp_pipe,
4436                );
4437            }
4438        }
4439        crate::pp::refuse_unsplit_if_remote(
4440            "decode_step_t (spec verify)",
4441            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4442             split (decode_step_t_core_ppn); or run spec on one device",
4443        )?;
4444        let cfg = &self.cfg;
4445        let n_embd = cfg.n_embd as usize;
4446        let eps = cfg.rms_eps;
4447        let t = tokens.len();
4448        let pos_d = match stream {
4449            Some((_, ctr)) => {
4450                let mut p = e.alloc_uninit::<i32>(t)?;
4451                e.pos_iota(ctr, &mut p, t)?;
4452                p
4453            }
4454            None => {
4455                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4456                e.htod_i32(&pos_vec)?
4457            }
4458        };
4459
4460        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
4461        let x = match (stream, embd_dev) {
4462            (Some((vtok, _)), Some((g, qt, rb))) => {
4463                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4464            }
4465            (None, Some((g, qt, rb))) => match vtok_dev {
4466                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
4467                // bit-identical rows to the host-token arm (same per-dtype deq).
4468                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
4469                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4470            },
4471            _ => {
4472                assert!(
4473                    vtok_dev.is_none(),
4474                    "device-token verify requires the resident embed table (embd_dev)"
4475                );
4476                e.htod(&self.embd.gather(n_embd, tokens))?
4477            }
4478        };
4479
4480        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
4481        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
4482        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
4483        let x = self.verify_layers(
4484            e,
4485            x,
4486            0,
4487            self.layers.len(),
4488            &pos_d,
4489            pos0,
4490            t,
4491            cache,
4492            ckpt.take(),
4493            stream,
4494            graphs,
4495        )?;
4496
4497        let mut hn = vbuf(e, t * n_embd)?;
4498        let serving_head = self.cfg.step35.is_some() || self.qwen35_serving_class();
4499        let logits = if serving_head {
4500            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
4501            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
4502            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
4503            // serve one batched numeric class at every live width, including B=1. Keep the
4504            // verify head in that same class; other generic families retain the decode-exact
4505            // head that their run-spec contract pins.
4506            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4507            e.matmul(&self.output, &hn, t)?
4508        } else {
4509            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4510            e.matmul_decode_exact(&self.output, &hn, t)?
4511        };
4512        // stream: the device pos counter owns position; host mirror reconciles at drain.
4513        if stream.is_none() {
4514            cache.pos += t;
4515        }
4516        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
4517        Ok((logits, if spec_hpost() { hn } else { x }))
4518    }
4519
4520    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
4521    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
4522    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
4523    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
4524    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
4525    /// the payload).
4526    ///
4527    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
4528    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
4529    /// receipts):
4530    ///
4531    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4532    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4533    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4534    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4535    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
4536    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4537    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4538    ///
4539    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4540    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4541    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4542    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4543    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
4544    ///
4545    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4546    ///    sharded loader leaves the table with stage 0 by construction).
4547    ///
4548    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4549    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4550    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4551    ///    model, every round.
4552    ///
4553    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4554    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4555    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4556    /// through the primary context by UVA — the same read the batched serving epilogue's
4557    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4558    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4559    ///
4560    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4561    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4562    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4563    ///
4564    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4565    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4566    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4567    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4568    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4569    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4570    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4571    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4572    #[allow(clippy::too_many_arguments)]
4573    fn decode_step_t_core_ppn(
4574        &self,
4575        e: &Engine,
4576        tokens: &[u32],
4577        pos0: usize,
4578        cache: &mut Cache,
4579        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4580        mut ckpt: Option<&mut VerifyCkpt>,
4581        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4582        fence: &[usize],
4583        pp_pipe: Option<bool>,
4584    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4585        let ticket = self.verify_stage0_issue(
4586            e,
4587            tokens,
4588            pos0,
4589            cache,
4590            embd_dev,
4591            ckpt.as_deref_mut(),
4592            stream,
4593            fence,
4594            pp_pipe,
4595            None,
4596        )?;
4597        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
4598    }
4599
4600    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
4601    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
4602    #[allow(clippy::too_many_arguments)]
4603    fn verify_stage0_issue(
4604        &self,
4605        e: &Engine,
4606        tokens: &[u32],
4607        pos0: usize,
4608        cache: &mut Cache,
4609        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4610        mut ckpt: Option<&mut VerifyCkpt>,
4611        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4612        fence: &[usize],
4613        pp_pipe: Option<bool>,
4614        trace: Option<SpecPipeTraceCtx>,
4615    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
4616        assert!(
4617            !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
4618            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
4619             (the gemma4 arms have their own decode_step_t twins)"
4620        );
4621        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
4622            return Err(
4623                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
4624                 boundary itself is host-staged, but device-resident verify still peer-reads \
4625                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
4626                 serving on this host class; spec requires local per-stage inputs first."
4627                    .into(),
4628            );
4629        }
4630        let rt = crate::pp::PpNRt::get(e)?;
4631        let n_st = fence.len() - 1;
4632        assert_eq!(
4633            rt.n_stages(),
4634            n_st,
4635            "PpNRt stage count {} != fence stages {n_st}",
4636            rt.n_stages()
4637        );
4638        let n_embd = self.cfg.n_embd as usize;
4639        let t = tokens.len();
4640        let payload = t * n_embd;
4641        if pp_pipe.is_some() {
4642            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
4643        }
4644        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
4645        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
4646        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
4647        // the report below names exactly two stages and must never imply it measured middle ones.
4648        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
4649        let pp_started = std::time::Instant::now();
4650        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
4651        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
4652        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
4653        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
4654        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
4655        // stage stream and the wait would self-order into a no-op.
4656        let caller_stream = e.stream();
4657        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
4658        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
4659        // the primary stream still holds queued reads of them — with event tracking elided,
4660        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
4661        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
4662        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
4663        // stage stream behind the caller before enqueueing new stage work.
4664        let reverse_started = std::time::Instant::now();
4665        if pp_pipe != Some(false) {
4666            rt.fence_stages_behind(&caller_stream)?;
4667        }
4668        if pp_pipe == Some(true) {
4669            // Both session verifies must alternate boundary slots even when the ordinary
4670            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
4671            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
4672            rt.prepare_overlap_slots(0, payload)?;
4673        }
4674        if pp_anatomy {
4675            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
4676            // prices any primary-stream rollback/refresh tail inherited from the prior round.
4677            for s in 0..n_st {
4678                let _st = rt.enter(s);
4679                rt.engine(s, e).stream().synchronize()?;
4680            }
4681            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
4682        }
4683
4684        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
4685        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
4686        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4687            match stream {
4688                Some((_, ctr)) => {
4689                    let mut p = es.alloc_uninit::<i32>(t)?;
4690                    es.pos_iota(ctr, &mut p, t)?;
4691                    Ok(p)
4692                }
4693                None => {
4694                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4695                    es.htod_i32(&pos_vec)
4696                }
4697            }
4698        };
4699
4700        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
4701        let slot = {
4702            let _st0 = rt.enter(0);
4703            let e0 = rt.engine(0, e);
4704            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
4705            let stage0_started = std::time::Instant::now();
4706            let pos_d = stage_pos(e0)?;
4707            let x = match (stream, embd_dev) {
4708                (Some((vtok, _)), Some((g, qt, rb))) => {
4709                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4710                }
4711                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4712                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
4713            };
4714            let x = self.verify_layers(
4715                e0,
4716                x,
4717                fence[0],
4718                fence[1],
4719                &pos_d,
4720                pos0,
4721                t,
4722                cache,
4723                ckpt.as_deref_mut(),
4724                stream,
4725                None,
4726            )?;
4727            if pp_anatomy {
4728                e0.stream().synchronize()?;
4729                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
4730            }
4731            let tx_started = std::time::Instant::now();
4732            let slot = if pp_pipe.is_some() {
4733                rt.tx_pipelined(0, &x, payload)?
4734            } else {
4735                rt.tx(0, &x, payload)?
4736            };
4737            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
4738            if pp_anatomy {
4739                e0.stream().synchronize()?;
4740                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
4741            }
4742            slot
4743            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
4744        };
4745
4746        Ok(VerifyBoundaryTicket {
4747            rt,
4748            caller_stream,
4749            slot,
4750            pos0,
4751            t,
4752            payload,
4753            n_st,
4754            pipelined: pp_pipe.is_some(),
4755            pp_anatomy,
4756            pp_started,
4757            reverse_ms,
4758            stage0_ms,
4759            tx_ms,
4760            trace,
4761        })
4762    }
4763
4764    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
4765    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
4766    #[allow(clippy::too_many_arguments)]
4767    fn verify_stage1_finish(
4768        &self,
4769        e: &Engine,
4770        ticket: VerifyBoundaryTicket,
4771        cache: &mut Cache,
4772        mut ckpt: Option<&mut VerifyCkpt>,
4773        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4774        fence: &[usize],
4775        publish_to_caller: bool,
4776    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4777        let VerifyBoundaryTicket {
4778            rt,
4779            caller_stream,
4780            slot,
4781            pos0,
4782            t,
4783            payload,
4784            n_st,
4785            pipelined,
4786            pp_anatomy,
4787            pp_started,
4788            reverse_ms,
4789            stage0_ms,
4790            tx_ms,
4791            trace,
4792        } = ticket;
4793        let n_embd = self.cfg.n_embd as usize;
4794        let eps = self.cfg.rms_eps;
4795        let mut slot = slot;
4796        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
4797        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
4798            match stream {
4799                Some((_, ctr)) => {
4800                    let mut p = es.alloc_uninit::<i32>(t)?;
4801                    es.pos_iota(ctr, &mut p, t)?;
4802                    Ok(p)
4803                }
4804                None => {
4805                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4806                    es.htod_i32(&pos_vec)
4807                }
4808            }
4809        };
4810
4811        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
4812        for s in 1..n_st - 1 {
4813            let _st = rt.enter(s);
4814            let es = rt.engine(s, e);
4815            let pos_d = stage_pos(es)?;
4816            let x = rt.rx(s - 1, slot, payload)?;
4817            let x = self.verify_layers(
4818                es,
4819                x,
4820                fence[s],
4821                fence[s + 1],
4822                &pos_d,
4823                pos0,
4824                t,
4825                cache,
4826                ckpt.as_deref_mut(),
4827                stream,
4828                None,
4829            )?;
4830            slot = if pipelined {
4831                rt.tx_pipelined(s, &x, payload)?
4832            } else {
4833                rt.tx(s, &x, payload)?
4834            };
4835        }
4836
4837        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
4838        let _stl = rt.enter(n_st - 1);
4839        let el = rt.engine(n_st - 1, e);
4840        let pos_d = stage_pos(el)?;
4841        let rx_started = std::time::Instant::now();
4842        let x = rt.rx(n_st - 2, slot, payload)?;
4843        if pp_anatomy {
4844            el.stream().synchronize()?;
4845            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
4846        }
4847        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
4848        let stage1_started = std::time::Instant::now();
4849        let x = self.verify_layers(
4850            el,
4851            x,
4852            fence[n_st - 1],
4853            fence[n_st],
4854            &pos_d,
4855            pos0,
4856            t,
4857            cache,
4858            ckpt.as_deref_mut(),
4859            stream,
4860            None,
4861        )?;
4862
4863        let mut hn = vbuf(el, payload)?;
4864        let logits = if self.cfg.step35.is_some() {
4865            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
4866            // Verify must not switch numeric class merely because the same session speculates.
4867            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4868            el.matmul(&self.output, &hn, t)?
4869        } else {
4870            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4871            el.matmul_decode_exact(&self.output, &hn, t)?
4872        };
4873        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
4874        if pp_anatomy {
4875            el.stream().synchronize()?;
4876            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
4877        }
4878        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
4879        // stream. Order the caller's stream behind that work before the buffers escape this
4880        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
4881        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
4882        // the following arm's KV in the same process).
4883        if publish_to_caller {
4884            rt.publish_to(n_st - 1, &caller_stream)?;
4885        }
4886        if pp_anatomy {
4887            if publish_to_caller {
4888                caller_stream.synchronize()?;
4889            }
4890            eprintln!(
4891                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
4892                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
4893                pp_started.elapsed().as_secs_f64() * 1e3,
4894            );
4895        }
4896        // stream: the device pos counter owns position; host mirror reconciles at drain.
4897        if stream.is_none() {
4898            cache.pos += t;
4899        }
4900        Ok((logits, if spec_hpost() { hn } else { x }))
4901    }
4902
4903    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
4904    ///
4905    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
4906    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
4907    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
4908    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
4909    /// bytes when a request moves from batched plain serving into speculative verify. Run the
4910    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
4911    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
4912    /// every norm/projection/FFN uses exactly the live serving dispatch.
4913    #[allow(clippy::too_many_arguments)]
4914    fn step35_verify_batch_layers(
4915        &self,
4916        e: &Engine,
4917        mut x: CudaSlice<f32>,
4918        lo: usize,
4919        hi: usize,
4920        pos0: usize,
4921        t: usize,
4922        cache: &mut Cache,
4923    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4924        let n_embd = self.cfg.n_embd as usize;
4925        self.cfg
4926            .step35
4927            .as_ref()
4928            .ok_or("step35 verify batch requires step35 cfg")?;
4929        let mut ph_last = std::time::Instant::now();
4930        for il in lo..hi {
4931            let mut next = e.uninit(t * n_embd)?;
4932            for r in 0..t {
4933                let mut row = e.uninit(n_embd)?;
4934                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4935                // The caller owns this verify's position. During controller overlap, cache.pos
4936                // still describes generation N while this stage-0 walk belongs to N+1.
4937                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
4938                let mut one = [&mut *cache];
4939                let out = self.step35_decode_batch_layers(
4940                    e,
4941                    row,
4942                    &mut one,
4943                    &row_pos,
4944                    il,
4945                    il + 1,
4946                    &mut ph_last,
4947                )?;
4948                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
4949            }
4950            self.dflash_tap(e, cache, il, &next, t)?;
4951            x = next;
4952        }
4953        Ok(x)
4954    }
4955
4956    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
4957    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
4958    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
4959    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
4960    /// prefix-keep, not all-or-nothing).
4961    pub(crate) fn dspark_verify_t_am(
4962        &self,
4963        e: &Engine,
4964        tokens: &[u32],
4965        pos0: usize,
4966        cache: &mut Cache,
4967    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
4968        let (logits, _hn) = self.decode_step_t_core_stream(
4969            e, tokens, pos0, cache, None, None, None, None, None, None,
4970        )?;
4971        let t = tokens.len();
4972        let v = self.output.out_features();
4973        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
4974        for r in 0..t {
4975            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
4976        }
4977        Ok(e.dtoh_u32(&am_d)?)
4978    }
4979
4980    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
4981    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
4982    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
4983    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
4984    pub(crate) fn dspark_verify_t_logits(
4985        &self,
4986        e: &Engine,
4987        tokens: &[u32],
4988        pos0: usize,
4989        cache: &mut Cache,
4990    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4991        let (logits, _hn) = self.decode_step_t_core_stream(
4992            e, tokens, pos0, cache, None, None, None, None, None, None,
4993        )?;
4994        Ok(logits)
4995    }
4996
4997    /// DSpark verify with the MTP column-stash armed: identical forward to
4998    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
4999    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
5000    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
5001    pub(crate) fn dspark_verify_t_am_ckpt(
5002        &self,
5003        e: &Engine,
5004        tokens: &[u32],
5005        pos0: usize,
5006        cache: &mut Cache,
5007    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5008        let mut ck = VerifyCkpt::new(self.layers.len());
5009        let (logits, _hn) = self.decode_step_t_core_stream(
5010            e,
5011            tokens,
5012            pos0,
5013            cache,
5014            None,
5015            Some(&mut ck),
5016            None,
5017            None,
5018            None,
5019            None,
5020        )?;
5021        let t = tokens.len();
5022        let v = self.output.out_features();
5023        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5024        for r in 0..t {
5025            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5026        }
5027        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
5028    }
5029
5030    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
5031    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
5032    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
5033    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
5034    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
5035    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
5036    pub(crate) fn dspark_verify_t_am_ckpt_dev(
5037        &self,
5038        e: &Engine,
5039        vtok: &CudaSlice<u32>,
5040        t: usize,
5041        pos0: usize,
5042        cache: &mut Cache,
5043        embd_dev: (&CudaSlice<u8>, i32, usize),
5044        graphs: Option<&mut DsparkVerifyGraphs>,
5045    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5046        debug_assert!(
5047            vtok.len() >= t,
5048            "verify window exceeds the device token buffer"
5049        );
5050        // The slab flag is a per-round statement: clear it here so a verify that never
5051        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
5052        // stale `true` steering the commit at slabs the round never wrote.
5053        let mut graphs = graphs;
5054        if let Some(g) = graphs.as_deref_mut() {
5055            g.round_slab = false;
5056        }
5057        let mut ck = VerifyCkpt::new(self.layers.len());
5058        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
5059        // arm's established pattern — spec.rs stream-mode verify does the same).
5060        let dummy = vec![0u32; t];
5061        let (logits, _hn) = self.decode_step_t_core_stream(
5062            e,
5063            &dummy,
5064            pos0,
5065            cache,
5066            Some(embd_dev),
5067            Some(&mut ck),
5068            None,
5069            None,
5070            Some(vtok),
5071            graphs,
5072        )?;
5073        let v = self.output.out_features();
5074        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5075        for r in 0..t {
5076            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5077        }
5078        Ok((am_d, DsparkVerifyCkpt(ck)))
5079    }
5080
5081    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
5082    pub(crate) fn dspark_verify_t_logits_ckpt(
5083        &self,
5084        e: &Engine,
5085        tokens: &[u32],
5086        pos0: usize,
5087        cache: &mut Cache,
5088    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5089        let mut ck = VerifyCkpt::new(self.layers.len());
5090        let (logits, _hn) = self.decode_step_t_core_stream(
5091            e,
5092            tokens,
5093            pos0,
5094            cache,
5095            None,
5096            Some(&mut ck),
5097            None,
5098            None,
5099            None,
5100            None,
5101        )?;
5102        Ok((logits, DsparkVerifyCkpt(ck)))
5103    }
5104
5105    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
5106    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
5107    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
5108    pub(crate) fn dspark_commit_prefix(
5109        &self,
5110        e: &Engine,
5111        cache: &mut Cache,
5112        snap: &crate::cache::CacheSnapshot,
5113        ckpt: &DsparkVerifyCkpt,
5114        keep: usize,
5115    ) -> Result<(), Box<dyn std::error::Error>> {
5116        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
5117    }
5118
5119    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
5120    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
5121    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
5122    /// from the stash of column keep-1), slab-addressed and batched into two copy
5123    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
5124    pub(crate) fn dspark_commit_prefix_slab(
5125        &self,
5126        e: &Engine,
5127        cache: &mut Cache,
5128        snap: &crate::cache::CacheSnapshot,
5129        ctx: &DsparkVerifyGraphs,
5130        keep: usize,
5131    ) -> Result<(), Box<dyn std::error::Error>> {
5132        use cudarc::driver::DevicePtr;
5133        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
5134        let mut conv_src: Vec<u64> = Vec::new();
5135        let mut ssm_src: Vec<u64> = Vec::new();
5136        let mut conv_dst: Vec<u64> = Vec::new();
5137        let mut ssm_dst: Vec<u64> = Vec::new();
5138        for il in 0..self.layers.len() {
5139            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
5140                kvl.len = saved + keep;
5141                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5142            }
5143            if let Some(rl) = cache.recur[il].as_ref() {
5144                let (pc, ps, _cw, _sw) = ctx
5145                    .slab_row(e, il, keep - 1)
5146                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
5147                conv_src.push(pc);
5148                ssm_src.push(ps);
5149                let st = &e.gpu.stream();
5150                let (dc, _g0) = rl.conv_state.device_ptr(st);
5151                let (ds, _g1) = rl.ssm_state.device_ptr(st);
5152                conv_dst.push(dc as u64);
5153                ssm_dst.push(ds as u64);
5154            }
5155        }
5156        let n = conv_src.len();
5157        if n > 0 {
5158            if state_copy_batch_on() {
5159                let mut tt = vec![0u64; 2 * n];
5160                tt[..n].copy_from_slice(&conv_src);
5161                tt[n..].copy_from_slice(&conv_dst);
5162                let ct = e.htod_u64(&tt)?;
5163                tt[..n].copy_from_slice(&ssm_src);
5164                tt[n..].copy_from_slice(&ssm_dst);
5165                let st = e.htod_u64(&tt)?;
5166                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
5167                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
5168            } else {
5169                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
5170                let row = keep - 1;
5171                for il in 0..self.layers.len() {
5172                    let Some(rl) = cache.recur[il].as_mut() else {
5173                        continue;
5174                    };
5175                    let k = ctx.lin_pos[&il];
5176                    {
5177                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
5178                        let win = sv.slice(row * cw..(row + 1) * cw);
5179                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
5180                    }
5181                    {
5182                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
5183                        let win = sv.slice(row * sw..(row + 1) * sw);
5184                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
5185                    }
5186                }
5187            }
5188        }
5189        cache.pos = snap.pos + keep;
5190        Ok(())
5191    }
5192
5193    /// Qwen35-family verify trunk in the live serving numeric class.
5194    ///
5195    /// Serving intentionally keeps this architecture in the generic batched program even at
5196    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
5197    ///
5198    /// Two arms, one numeric class:
5199    /// - DENSE (`Arch::Qwen35`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
5200    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
5201    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
5202    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
5203    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
5204    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
5205    ///   program its isolated serving step would). One weight read per layer per round
5206    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
5207    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
5208    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
5209    ///   serving layer body, preserving single-session autoregressive cache order (the
5210    ///   correctness reference; also the rollback seam for the t-parallel arm).
5211    ///
5212    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
5213    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
5214    #[allow(clippy::too_many_arguments)]
5215    fn qwen35_verify_batch_layers(
5216        &self,
5217        e: &Engine,
5218        x: CudaSlice<f32>,
5219        lo: usize,
5220        hi: usize,
5221        pos0: usize,
5222        t: usize,
5223        cache: &mut Cache,
5224        ckpt: Option<&mut VerifyCkpt>,
5225        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5226        graphs: Option<&mut DsparkVerifyGraphs>,
5227    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5228        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
5229        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
5230        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
5231        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
5232        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
5233        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
5234        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
5235            || !matches!(
5236                self.cfg.arch,
5237                memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
5238            )
5239            || t > 16;
5240        if rowwise {
5241            if stream.is_some() {
5242                // rowwise replays per row with host cache.pos — irreconcilable with a
5243                // device position counter. Burst callers must keep t <= 16 and the
5244                // ROWWISE env unset; refusing beats silently mispositioned rows.
5245                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
5246                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
5247                    .into());
5248            }
5249            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
5250        } else {
5251            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
5252        }
5253    }
5254
5255    /// The per-row correctness reference: replay each verify row through the authoritative
5256    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
5257    #[allow(clippy::too_many_arguments)]
5258    fn qwen35_verify_rowwise(
5259        &self,
5260        e: &Engine,
5261        mut x: CudaSlice<f32>,
5262        lo: usize,
5263        hi: usize,
5264        pos0: usize,
5265        t: usize,
5266        cache: &mut Cache,
5267        mut ckpt: Option<&mut VerifyCkpt>,
5268    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5269        let n_embd = self.cfg.n_embd as usize;
5270        let saved_pos = cache.pos;
5271        let mut ph_last = std::time::Instant::now();
5272        for il in lo..hi {
5273            let mut next = e.uninit(t * n_embd)?;
5274            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5275                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5276                    Some(Vec::with_capacity(t - 1))
5277                } else {
5278                    None
5279                };
5280            for r in 0..t {
5281                cache.pos = pos0 + r;
5282                let mut row = e.uninit(n_embd)?;
5283                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5284                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5285                let mut one = [&mut *cache];
5286                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
5287                let out = match self.decode_batch_layers(
5288                    e,
5289                    row,
5290                    &mut one,
5291                    &ctx,
5292                    &row_pos,
5293                    &mut ph_last,
5294                ) {
5295                    Ok(out) => out,
5296                    Err(error) => {
5297                        cache.pos = saved_pos;
5298                        return Err(error);
5299                    }
5300                };
5301                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5302                if r + 1 < t {
5303                    if let Some(states) = col_states.as_mut() {
5304                        let recur = cache.recur[il]
5305                            .as_ref()
5306                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
5307                        states.push((
5308                            e.clone_dtod(&recur.conv_state)?,
5309                            e.clone_dtod(&recur.ssm_state)?,
5310                        ));
5311                    }
5312                }
5313            }
5314            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
5315                checkpoint.cols[il] = Some(states);
5316            }
5317            x = next;
5318        }
5319        cache.pos = saved_pos;
5320        Ok(x)
5321    }
5322
5323    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
5324    ///
5325    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
5326    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
5327    /// pins the serving batch tier already carries:
5328    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
5329    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
5330    ///     alone;
5331    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
5332    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
5333    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
5334    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
5335    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
5336    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
5337    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
5338    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
5339    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
5340    /// program its isolated B=1 serving step would.
5341    ///
5342    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
5343    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
5344    #[allow(clippy::too_many_arguments)]
5345    fn qwen35_verify_tparallel(
5346        &self,
5347        e: &Engine,
5348        mut x: CudaSlice<f32>,
5349        lo: usize,
5350        hi: usize,
5351        pos0: usize,
5352        t: usize,
5353        cache: &mut Cache,
5354        mut ckpt: Option<&mut VerifyCkpt>,
5355        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5356        mut graphs: Option<&mut DsparkVerifyGraphs>,
5357    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5358        let seqs_append =
5359            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
5360        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
5361
5362        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
5363        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
5364        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
5365        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
5366        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
5367        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
5368        // full-verify bodies).
5369        if stream.is_some() && graphs.is_some() {
5370            return Err(
5371                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
5372                        cannot arm together"
5373                    .into(),
5374            );
5375        }
5376        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
5377        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
5378        // moves the kv caches). Then:
5379        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
5380        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
5381        //    full-verify graph per (vt, rung) — linear layers through the shared
5382        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
5383        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
5384        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
5385        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
5386        //    the full-attention layers run eager (batched rows when eligible).
5387        if let Some(g) = graphs.as_deref_mut() {
5388            g.refresh_tables(e, cache)?;
5389            g.round_slab = false;
5390            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
5391                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
5392                // full capture past the ceiling falls through to the segment/eager arms.
5393                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
5394                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
5395                    g.round_slab = true;
5396                    return Ok(out);
5397                }
5398            }
5399            // Round-atomic ceiling check for the segment door: if any linear run in this
5400            // walk would need a NEW capture past the ceiling, the whole round runs the
5401            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
5402            // would corrupt the commit).
5403            if !g.segments_ready(self, lo, hi, t) {
5404                graphs = None;
5405            }
5406        }
5407        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
5408        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
5409        let pos_d = match stream {
5410            Some((_, ctr)) => {
5411                let mut p = e.alloc_uninit::<i32>(t)?;
5412                e.pos_iota(ctr, &mut p, t)?;
5413                p
5414            }
5415            None => {
5416                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
5417                e.htod_i32(&pos_host)?
5418            }
5419        };
5420        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
5421        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
5422        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
5423        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
5424        // rides the dc rows kernels and never reaches the fallback).
5425        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
5426        let mut il = lo;
5427        while il < hi {
5428            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5429                let mut end = il;
5430                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
5431                    end += 1;
5432                }
5433                let g = graphs.as_deref_mut().expect("checked above");
5434                x = g.run_segment(self, e, il, end, &x, t, cache)?;
5435                g.round_slab = true;
5436                il = end;
5437                continue;
5438            }
5439            let layer = &self.layers[il];
5440            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
5441                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
5442                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
5443                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
5444                x = self.qwen35_tparallel_linear_layer(
5445                    e,
5446                    il,
5447                    &x,
5448                    t,
5449                    cache,
5450                    ckpt.as_deref_mut(),
5451                    None,
5452                    None,
5453                )?;
5454                il += 1;
5455                continue;
5456            }
5457            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
5458            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
5459            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
5460            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
5461            // run (lane/draftcost-moe).
5462            x = self.qwen35_tparallel_fa_layer(
5463                e,
5464                il,
5465                &x,
5466                t,
5467                cache,
5468                FaLayerArgs {
5469                    pos_d: &pos_d,
5470                    pos_rows: &mut pos_rows,
5471                    pos0,
5472                    seqs_append,
5473                    batch_fa_on,
5474                    graph_cap: None,
5475                    stream,
5476                    ckpt: ckpt.as_deref_mut(),
5477                },
5478            )?;
5479            il += 1;
5480        }
5481        Ok(x)
5482    }
5483
5484    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
5485    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
5486    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
5487    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
5488    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
5489    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
5490    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
5491    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
5492    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
5493    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
5494    /// original singles chain, byte-for-byte.
5495    #[allow(clippy::too_many_arguments)]
5496    fn qwen35_tparallel_dense_ffn(
5497        &self,
5498        e: &Engine,
5499        ffn_gate: &crate::model::GpuTensor,
5500        ffn_up: &crate::model::GpuTensor,
5501        ffn_down: &crate::model::GpuTensor,
5502        zn: &CudaSlice<f32>,
5503        t: usize,
5504        n_embd: usize,
5505    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5506        let n_ff = ffn_gate.out_features();
5507        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
5508        if Engine::tk_ffn_dual_on() {
5509            if let Some(((g, gs), (u, us))) =
5510                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
5511            {
5512                if e.uses_q8_1_fast(ffn_down) {
5513                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
5514                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
5515                }
5516                let mut act = e.uninit(t * n_ff)?;
5517                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
5518                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5519                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
5520            }
5521        }
5522        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
5523        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
5524        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
5525        let mut act = e.uninit(t * n_ff)?;
5526        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
5527        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
5528        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
5529    }
5530
5531    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
5532    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
5533    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
5534    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
5535    ///
5536    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
5537    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
5538    ///   generation's cache lands at new addresses that only the per-verify table refresh
5539    ///   knows — the slice-3 baked-address lesson);
5540    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
5541    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
5542    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
5543    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
5544    ///   round whose rows all sit inside the rung;
5545    /// - the host len bump moves to the replay caller (captured host code does not
5546    ///   re-run at replay).
5547    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
5548    /// host-branches on t_kv and must never be captured.
5549    #[allow(clippy::too_many_arguments)]
5550    fn qwen35_tparallel_fa_layer(
5551        &self,
5552        e: &Engine,
5553        il: usize,
5554        x: &CudaSlice<f32>,
5555        t: usize,
5556        cache: &mut Cache,
5557        args: FaLayerArgs<'_>,
5558    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5559        use cudarc::driver::DevicePtr;
5560        let cfg = &self.cfg;
5561        let n_embd = cfg.n_embd as usize;
5562        let eps = cfg.rms_eps;
5563        let head_dim_global = cfg.head_dim_k as usize;
5564        let layer = &self.layers[il];
5565        let FaLayerArgs {
5566            pos_d,
5567            pos_rows,
5568            pos0,
5569            seqs_append,
5570            batch_fa_on,
5571            graph_cap,
5572            stream,
5573            mut ckpt,
5574        } = args;
5575
5576        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
5577        let anorm = layer.attn_norm.float_data();
5578        let mut xn = e.uninit(t * n_embd)?;
5579        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
5580        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
5581
5582        let mixed: CudaSlice<f32> = match &layer.mixer {
5583            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
5584            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
5585            // per-row serving-kernel chain cannot run (host state swaps keyed on host
5586            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
5587            // rebuild — the per-row chain only produces per-column clones). GDN rides
5588            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
5589            // and its one-scan recurrence is pinned bit-identical to T chained T=1
5590            // steps (its header + kernel-check). Position-independent, so no counter
5591            // plumbing is needed. Guards mirror the generic call site exactly.
5592            Mixer::Linear(la) if stream.is_some() => {
5593                if !(t >= 3 || (t == 2 && spec_m2()))
5594                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
5595                    || !e.uses_q8_1_fast(&la.ssm_out)
5596                {
5597                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
5598                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
5599                        .into());
5600                }
5601                let want = ckpt.is_some();
5602                let (out, stash) =
5603                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
5604                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
5605                    ck.gdn[il] = Some(st);
5606                }
5607                out
5608            }
5609            Mixer::Linear(_) => {
5610                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
5611            }
5612            Mixer::Full(fa) => {
5613                let geometry = cfg.full_attention_geometry_at(il as u32);
5614                let n_head = geometry.n_head as usize;
5615                let n_head_kv = geometry.n_head_kv as usize;
5616                let head_dim = geometry.head_dim_k as usize;
5617                let rope_dims = geometry.n_rot as usize;
5618                let rope_base = geometry.rope_base;
5619                let scale = geometry.attention_scale();
5620                // Batched projections: one weight read serves all T rows.
5621                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
5622                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
5623                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
5624                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
5625                    [&fa.wq, &fa.wk, &fa.wv],
5626                    &hq,
5627                    &hd,
5628                    t,
5629                )? {
5630                    Some(mut g3) => {
5631                        let v = g3.pop().unwrap();
5632                        let k = g3.pop().unwrap();
5633                        let qf = g3.pop().unwrap();
5634                        (qf, k, v)
5635                    }
5636                    None => (
5637                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
5638                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
5639                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
5640                    ),
5641                };
5642                let gated =
5643                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
5644                let (mut q, gate) = if gated {
5645                    let mut qs = e.uninit(t * n_head * head_dim)?;
5646                    let mut gs = e.uninit(t * n_head * head_dim)?;
5647                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
5648                    (qs, Some(gs))
5649                } else {
5650                    (qf, None)
5651                };
5652                let mut qn = e.uninit(t * n_head * head_dim)?;
5653                e.rms_norm(
5654                    &q,
5655                    fa.q_norm.float_data(),
5656                    &mut qn,
5657                    head_dim,
5658                    t * n_head,
5659                    eps,
5660                )?;
5661                q = qn;
5662                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
5663                e.rms_norm(
5664                    &k,
5665                    fa.k_norm.float_data(),
5666                    &mut kn,
5667                    head_dim,
5668                    t * n_head_kv,
5669                    eps,
5670                )?;
5671                k = kn;
5672                e.rope_neox(
5673                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
5674                )?;
5675                e.rope_neox(
5676                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5677                )?;
5678
5679                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
5680                // draft), each through the b_n=1 serving kernels at its own t_kv.
5681                let q_dim = n_head * head_dim;
5682                let kv_dim = n_head_kv * head_dim;
5683                let mut attn = e.uninit(t * q_dim)?;
5684                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
5685                    let kvl = cache.kv[il].as_ref().unwrap();
5686                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
5687                    // the batched twins; the per-row fallback reads pair 0 (same cache
5688                    // for every row of one layer). Graph mode reads the ctx table.
5689                    let local: Option<CudaSlice<u64>> = match graph_cap {
5690                        Some(_) => None,
5691                        None => {
5692                            let s = &e.gpu.stream();
5693                            let (pk, _g) = kvl.k.device_ptr(s);
5694                            let (pv, _g2) = kvl.v.device_ptr(s);
5695                            let mut tbl = Vec::with_capacity(2 * t);
5696                            for _ in 0..t {
5697                                tbl.push(pk as u64);
5698                                tbl.push(pv as u64);
5699                            }
5700                            Some(e.htod_u64(&tbl)?)
5701                        }
5702                    };
5703                    (
5704                        kvl.kv_dim_k,
5705                        kvl.kv_dim_v,
5706                        kvl.k_tok_bytes,
5707                        kvl.v_tok_bytes,
5708                        kvl.len,
5709                        local,
5710                    )
5711                };
5712                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
5713                    Some((tb, off, _)) => (tb, off),
5714                    None => (kv_local.as_ref().expect("built above"), 0),
5715                };
5716                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
5717                // section batches into the z-batched serving twins when every row of
5718                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
5719                // guards are evaluated at the round's FIRST and LAST t_kv — the
5720                // eligibility window (vec floor .. v4 max) and each split-ladder rung
5721                // are intervals in t_kv, so ends-inside means all-inside (the straddle
5722                // law). Appending all T rows before any attend is read-equivalent to
5723                // the interleaved order: row r's walk reads keys 0..len0+r only, and
5724                // rows > r land at slots it never touches; every written cache row is
5725                // the per-token appender's exact warp program (kernel-check pinned).
5726                let t_kv_first = len0 + 1;
5727                let t_kv_last = len0 + t;
5728                let rows_batched = t >= 2
5729                    && seqs_append
5730                    && batch_fa_on
5731                    && dspark_fa_rows_on()
5732                    // the z-batched twins read stacked rows at the CACHE's kv dims;
5733                    // the projection stack is [T, n_head_kv*head_dim] — they must be
5734                    // the same stride or row z misaligns (true for this family; the
5735                    // guard keeps any asymmetric-kv model on the per-row loop).
5736                    && kdk == kv_dim
5737                    && kdv == kv_dim
5738                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
5739                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
5740                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
5741                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
5742                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
5743                // grid only — bytes proven equal above). Capture-time invariants refuse
5744                // loudly rather than bake a divergent body.
5745                let (size_kv_max, sp) = match graph_cap {
5746                    Some((_, _, rung)) => {
5747                        if !rows_batched {
5748                            return Err(format!(
5749                                "fa graph capture: layer {il} round is not batchable \
5750                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
5751                                 must never be captured"
5752                            )
5753                            .into());
5754                        }
5755                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
5756                        if t_kv_last > rung
5757                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
5758                        {
5759                            return Err(format!(
5760                                "fa graph capture: rung {rung} does not cover round \
5761                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
5762                            )
5763                            .into());
5764                        }
5765                        (rung, sp_r)
5766                    }
5767                    None => (
5768                        t_kv_last,
5769                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
5770                    ),
5771                };
5772                if let Some((_, ctr)) = stream {
5773                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
5774                    // — the generic stream arm's exact shape (rows kernels are pinned
5775                    // byte-identical to the per-row programs by kernel-check). Host len
5776                    // stays a stale lower bound; the burst drain reconciles it.
5777                    let kvl = cache.kv[il].as_mut().unwrap();
5778                    e.append_kv_quantized_rows_dc(
5779                        &k,
5780                        &v,
5781                        &mut kvl.k,
5782                        &mut kvl.v,
5783                        ctr,
5784                        t,
5785                        kdk,
5786                        kdv,
5787                        ktb,
5788                        vtb,
5789                        Engine::kv_fp8_on(),
5790                    )?;
5791                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
5792                    let k_view = e.view_u8(&kvl.k, upper * ktb);
5793                    let v_view = e.view_u8(&kvl.v, upper * vtb);
5794                    e.fa_decode_rows_dc(
5795                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
5796                        t, scale, ktb, vtb, 0, false,
5797                    )?;
5798                } else if rows_batched {
5799                    e.append_kv_quantized_seqs(
5800                        &k,
5801                        &v,
5802                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
5803                        pos_d,
5804                        t,
5805                        kdk,
5806                        kdv,
5807                        ktb,
5808                        vtb,
5809                    )?;
5810                    if graph_cap.is_none() {
5811                        cache.kv[il].as_mut().unwrap().len += t;
5812                    }
5813                    e.fa_decode_batch_seqs_v4(
5814                        &q,
5815                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
5816                        pos_d,
5817                        &mut attn,
5818                        head_dim,
5819                        n_head,
5820                        n_head_kv,
5821                        t,
5822                        size_kv_max,
5823                        scale,
5824                        sp,
5825                        ktb,
5826                        vtb,
5827                    )?;
5828                } else {
5829                    if pos_rows.is_none() {
5830                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
5831                        // the dc rows kernels above and never reaches this fallback).
5832                        *pos_rows = Some(match stream {
5833                            Some((_, ctr)) => (0..t)
5834                                .map(|r| {
5835                                    let mut b = e.alloc_uninit::<i32>(1)?;
5836                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
5837                                    Ok(b)
5838                                })
5839                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
5840                            None => (0..t)
5841                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
5842                                .collect::<Result<_, _>>()?,
5843                        });
5844                    }
5845                    let pos_rows = pos_rows.as_ref().unwrap();
5846                    for r in 0..t {
5847                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
5848                        // whose row 0 is this row (arithmetic-free materialization copies,
5849                        // same as decode's per-seq fallback arm).
5850                        let mut k_row = e.uninit(kv_dim)?;
5851                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
5852                        let mut v_row = e.uninit(kv_dim)?;
5853                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
5854                        let pos_row = &pos_rows[r];
5855                        let kvl = cache.kv[il].as_mut().unwrap();
5856                        if seqs_append {
5857                            e.append_kv_quantized_seqs(
5858                                &k_row,
5859                                &v_row,
5860                                &kv_tbl.slice(kv_off..kv_off + 2),
5861                                pos_row,
5862                                1,
5863                                kdk,
5864                                kdv,
5865                                ktb,
5866                                vtb,
5867                            )?;
5868                            kvl.len += 1;
5869                        } else {
5870                            e.append_kv_quantized_view(
5871                                &k_row.slice(0..kv_dim),
5872                                &v_row.slice(0..kv_dim),
5873                                &mut kvl.k,
5874                                &mut kvl.v,
5875                                kvl.len,
5876                                kvl.kv_dim_k,
5877                                kvl.kv_dim_v,
5878                                kvl.k_tok_bytes,
5879                                kvl.v_tok_bytes,
5880                                Engine::kv_fp8_on(),
5881                            )?;
5882                            kvl.len += 1;
5883                        }
5884                        let t_kv = kvl.len;
5885                        let mut q_row = e.uninit(q_dim)?;
5886                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
5887                        let mut a_row = e.uninit(q_dim)?;
5888                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
5889                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
5890                            e.fa_decode_batch_seqs_v4(
5891                                &q_row,
5892                                &kv_tbl.slice(kv_off..kv_off + 2),
5893                                pos_row,
5894                                &mut a_row,
5895                                head_dim,
5896                                n_head,
5897                                n_head_kv,
5898                                1,
5899                                t_kv,
5900                                scale,
5901                                sp0_r,
5902                                ktb,
5903                                vtb,
5904                            )?;
5905                        } else {
5906                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
5907                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
5908                            let mut a_view = a_row.slice_mut(0..q_dim);
5909                            e.fa_decode_kvmod_view(
5910                                &q_row.slice(0..q_dim),
5911                                &k_view,
5912                                &v_view,
5913                                &mut a_view,
5914                                head_dim,
5915                                n_head,
5916                                n_head_kv,
5917                                t_kv,
5918                                scale,
5919                                kvl.k_tok_bytes,
5920                                kvl.v_tok_bytes,
5921                                Engine::kv_fp8_on(),
5922                            )?;
5923                        }
5924                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
5925                    }
5926                }
5927
5928                // Output gate (element-wise) + o-proj at m=T.
5929                let attn_g = match &gate {
5930                    Some(g) => {
5931                        let n = t * q_dim;
5932                        let mut gsig = e.uninit(n)?;
5933                        e.sigmoid(g, &mut gsig, n)?;
5934                        let mut ag = e.uninit(n)?;
5935                        e.mul(&attn, &gsig, &mut ag, n)?;
5936                        ag
5937                    }
5938                    None => attn,
5939                };
5940                e.matmul(&fa.wo, &attn_g, t)?
5941            }
5942        };
5943
5944        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
5945        let pnorm = layer.post_attn_norm.float_data();
5946        let mut x1 = e.uninit(t * n_embd)?;
5947        let mut zn = e.uninit(t * n_embd)?;
5948        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
5949        let ffn_out = match &layer.ffn {
5950            crate::hybrid::Ffn::Dense {
5951                ffn_gate,
5952                ffn_up,
5953                ffn_down,
5954            } => {
5955                assert!(
5956                    self.cfg.m3.is_none(),
5957                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
5958                );
5959                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
5960            }
5961            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
5962        };
5963        let mut x2 = e.uninit(t * n_embd)?;
5964        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
5965        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
5966        self.dflash_tap(e, cache, il, &x2, t)?;
5967        Ok(x2)
5968    }
5969
5970    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
5971    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
5972    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
5973    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
5974    /// bit-identical by construction:
5975    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
5976    ///   the device sequence is driven entirely by the 6-entry pointer table, which
5977    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
5978    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
5979    ///   legacy post-swap clone read.
5980    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
5981    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
5982    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
5983    /// None builds the per-verify table exactly as before.
5984    #[allow(clippy::too_many_arguments)]
5985    fn qwen35_tparallel_linear_layer(
5986        &self,
5987        e: &Engine,
5988        il: usize,
5989        x: &CudaSlice<f32>,
5990        t: usize,
5991        cache: &mut Cache,
5992        mut ckpt: Option<&mut VerifyCkpt>,
5993        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
5994        table_src: Option<(&CudaSlice<u64>, usize)>,
5995    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5996        use cudarc::driver::DevicePtr;
5997        let cfg = &self.cfg;
5998        let n_embd = cfg.n_embd as usize;
5999        let eps = cfg.rms_eps;
6000        let layer = &self.layers[il];
6001        let Mixer::Linear(la) = &layer.mixer else {
6002            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
6003        };
6004        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6005        let anorm = layer.attn_norm.float_data();
6006        let mut xn = e.uninit(t * n_embd)?;
6007        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6008        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6009
6010        let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
6011        let d_state = ssm.state_size as usize;
6012        let num_k = ssm.group_count as usize;
6013        let num_v = ssm.time_step_rank as usize;
6014        let d_conv = ssm.conv_kernel as usize;
6015        let key_dim = d_state * num_k;
6016        let value_dim = d_state * num_v;
6017        let conv_dim = key_dim * 2 + value_dim;
6018        let gdn_scale = 1.0 / (d_state as f32).sqrt();
6019
6020        // ---- batched projections: one weight read for all T rows ----
6021        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
6022        // per (tensor, token, row) to the four singles; refused (layout/tier) or
6023        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
6024        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
6025            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
6026            &hq,
6027            &hd,
6028            t,
6029        )? {
6030            Some(mut g4) => {
6031                let alpha = g4.pop().unwrap();
6032                let beta_raw = g4.pop().unwrap();
6033                let z = g4.pop().unwrap();
6034                let qkv_mixed = g4.pop().unwrap();
6035                (qkv_mixed, z, beta_raw, alpha)
6036            }
6037            None => (
6038                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
6039                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
6040                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
6041                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
6042            ),
6043        };
6044        let beta_w = la.ssm_beta.out_features();
6045        let alpha_w = la.ssm_alpha.out_features();
6046        let qkv_w = la.wqkv.out_features();
6047
6048        // ---- per-row state chain through the b_n=1 serving kernels ----
6049        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
6050        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
6051        let table_local: Option<CudaSlice<u64>> = match table_src {
6052            Some(_) => None,
6053            None => {
6054                let rl = cache.recur[il].as_ref().unwrap();
6055                let s = &e.gpu.stream();
6056                let (pc, _g0) = rl.conv_state.device_ptr(s);
6057                let (p0, _g1) = rl.ssm_state.device_ptr(s);
6058                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
6059                Some(e.htod_u64(&[
6060                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
6061                ])?)
6062            }
6063        };
6064        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
6065            Some((tb, off)) => (tb, off),
6066            None => (table_local.as_ref().unwrap(), 0),
6067        };
6068        let mut o_all = e.uninit(t * value_dim)?;
6069        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6070            if ckpt.is_some() && stash.is_none() && t >= 2 {
6071                Some(Vec::with_capacity(t - 1))
6072            } else {
6073                None
6074            };
6075        let mut stash = stash;
6076        // Per-row scratch reused across rows (uninit is cheap but not free at
6077        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
6078        // [T, ...] buffers — zero arithmetic-free copies in this loop.
6079        let mut conv_out = e.uninit(conv_dim)?;
6080        let mut q_l2 = e.uninit(value_dim)?;
6081        let mut k_l2 = e.uninit(value_dim)?;
6082        let mut v_gd = e.uninit(value_dim)?;
6083        let mut beta_b = e.uninit(num_v)?;
6084        let mut g_log = e.uninit(num_v)?;
6085        for r in 0..t {
6086            let base = toff + if r % 2 == 0 { 0 } else { 3 };
6087            let conv_view = table.slice(base..base + 1);
6088            let in_view = table.slice(base + 1..base + 2);
6089            let out_view = table.slice(base + 2..base + 3);
6090            e.ssm_conv1d_fused_decode_b_view(
6091                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
6092                &conv_view,
6093                la.ssm_conv1d.float_data(),
6094                &mut conv_out,
6095                conv_dim,
6096                d_conv,
6097                1,
6098            )?;
6099            e.gdn_prep_decode_b_view(
6100                &conv_out,
6101                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
6102                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
6103                la.ssm_dt.float_data(),
6104                la.ssm_a.float_data(),
6105                &mut q_l2,
6106                &mut k_l2,
6107                &mut v_gd,
6108                &mut beta_b,
6109                &mut g_log,
6110                d_state,
6111                num_v,
6112                num_k,
6113                key_dim,
6114                eps,
6115                conv_dim,
6116                1,
6117            )?;
6118            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
6119            e.gdn_scan_s128_batched_view(
6120                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
6121                gdn_scale,
6122            )?;
6123            if r + 1 < t {
6124                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
6125                // odd rows write s0 — the same physical state the legacy post-swap
6126                // canonical clone read.
6127                let rl = cache.recur[il]
6128                    .as_ref()
6129                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
6130                let ssm_src = if r % 2 == 0 {
6131                    &rl.ssm_state_alt
6132                } else {
6133                    &rl.ssm_state
6134                };
6135                match stash.as_mut() {
6136                    Some((conv_slab, ssm_slab)) => {
6137                        // BOTH stash reads go through the pointer table at run time: the
6138                        // ssm handles ping-pong between rounds, and the ctx (with its
6139                        // captured graphs) outlives the Cache — a fresh generation's
6140                        // conv/ssm buffers land at new addresses that only the per-round
6141                        // table refresh knows. A baked direct copy would read freed
6142                        // memory (parity was the slice-3 smoke divergence; cache
6143                        // lifetime is the cross-generation twin).
6144                        e.copy_indirect_src_f32(
6145                            &conv_view,
6146                            conv_slab,
6147                            r * conv_dim * (d_conv - 1),
6148                            conv_dim * (d_conv - 1),
6149                        )?;
6150                        // The ssm handles PING-PONG between rounds: a captured direct
6151                        // copy would bake the capture-time physical buffer and read the
6152                        // wrong parity after any odd-vt round (the slice-3 smoke
6153                        // divergence). Read the src address from row r's OUT table
6154                        // entry at run time — the same entry the scan just wrote.
6155                        e.copy_indirect_src_f32(
6156                            &out_view,
6157                            ssm_slab,
6158                            r * d_state * d_state * num_v,
6159                            d_state * d_state * num_v,
6160                        )?;
6161                    }
6162                    None => {
6163                        if let Some(states) = col_states.as_mut() {
6164                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
6165                        }
6166                    }
6167                }
6168            }
6169        }
6170        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
6171        // handle motion is identical and the device sequence never read the handles.
6172        if t % 2 == 1 {
6173            let rl = cache.recur[il].as_mut().unwrap();
6174            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
6175        }
6176        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6177            checkpoint.cols[il] = Some(states);
6178        }
6179
6180        // ---- batched gated norm + out-projection at m=T ----
6181        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
6182            let (gq, gd) = e.gated_rmsnorm_q8_1(
6183                &o_all,
6184                la.ssm_norm.float_data(),
6185                &z,
6186                d_state,
6187                t * num_v,
6188                eps,
6189            )?;
6190            let g0 = e.zeros(0)?;
6191            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
6192        } else {
6193            let mut gn = e.uninit(t * value_dim)?;
6194            e.gated_rmsnorm(
6195                &o_all,
6196                la.ssm_norm.float_data(),
6197                &z,
6198                &mut gn,
6199                d_state,
6200                t * num_v,
6201                eps,
6202            )?;
6203            e.matmul(&la.ssm_out, &gn, t)?
6204        };
6205
6206        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6207        let pnorm = layer.post_attn_norm.float_data();
6208        let mut x1 = e.uninit(t * n_embd)?;
6209        let mut zn = e.uninit(t * n_embd)?;
6210        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6211        let ffn_out = match &layer.ffn {
6212            crate::hybrid::Ffn::Dense {
6213                ffn_gate,
6214                ffn_up,
6215                ffn_down,
6216            } => {
6217                assert!(
6218                    self.cfg.m3.is_none(),
6219                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6220                );
6221                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6222            }
6223            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6224        };
6225        let mut x2 = e.uninit(t * n_embd)?;
6226        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6227        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6228        self.dflash_tap(e, cache, il, &x2, t)?;
6229        Ok(x2)
6230    }
6231
6232    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
6233    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
6234    /// carried in from outside the range) and exits with the range's final residual materialized
6235    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
6236    /// instead of one.
6237    ///
6238    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
6239    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
6240    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
6241    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
6242    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
6243    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
6244    /// code — there is no "split version" of the verify math.
6245    ///
6246    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
6247    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
6248    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
6249    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
6250    #[allow(clippy::too_many_arguments)]
6251    fn verify_layers(
6252        &self,
6253        e: &Engine,
6254        mut x: CudaSlice<f32>,
6255        lo: usize,
6256        hi: usize,
6257        pos_d: &CudaSlice<i32>,
6258        pos0: usize,
6259        t: usize,
6260        cache: &mut Cache,
6261        mut ckpt: Option<&mut VerifyCkpt>,
6262        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6263        graphs: Option<&mut DsparkVerifyGraphs>,
6264    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6265        if self.cfg.step35.is_some() {
6266            if stream.is_some() {
6267                return Err(
6268                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
6269                            cannot express the SWA offset KV view)"
6270                        .into(),
6271                );
6272            }
6273            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
6274        }
6275        if self.qwen35_serving_class() {
6276            return self.qwen35_verify_batch_layers(
6277                e,
6278                x,
6279                lo,
6280                hi,
6281                pos0,
6282                t,
6283                cache,
6284                ckpt.take(),
6285                stream,
6286                graphs,
6287            );
6288        }
6289        let n_embd = self.cfg.n_embd as usize;
6290        let eps = self.cfg.rms_eps;
6291        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
6292        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
6293        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
6294        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
6295        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
6296        // residual the next layer needs) as its `res` output. Falls back to the separate add
6297        // when the next layer is off the fused-q8 path.
6298        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
6299        for il in lo..hi {
6300            let layer = &self.layers[il];
6301            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
6302            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
6303            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
6304            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
6305            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
6306            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
6307            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
6308            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
6309            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
6310            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
6311            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
6312            // projections only; Linear mixer: the batched arm — the per-column fallback needs
6313            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
6314            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
6315            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
6316            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
6317            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
6318            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
6319            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
6320            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
6321            let lin_q8_only = match &layer.mixer {
6322                Mixer::Linear(la) => {
6323                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
6324                }
6325                Mixer::Full(_) if self.cfg.step35.is_some() => false,
6326                _ => true,
6327            };
6328            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
6329            // a non-fused layer still performs the residual add.
6330            let taken = pending.take();
6331            let (h, h_q8) = if norm_fused && lin_q8_only {
6332                let pair = match taken {
6333                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
6334                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
6335                    Some((x1p, f1p)) => {
6336                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
6337                        let p = e.add_rms_norm_q8_1(
6338                            &x1p,
6339                            &f1p,
6340                            layer.attn_norm.float_data(),
6341                            &mut x2,
6342                            n_embd,
6343                            t,
6344                            eps,
6345                        )?;
6346                        x = x2;
6347                        p
6348                    }
6349                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
6350                };
6351                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
6352            } else {
6353                if let Some((x1p, f1p)) = taken {
6354                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6355                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
6356                    x = x2;
6357                }
6358                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
6359                if norm_fused {
6360                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6361                } else {
6362                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6363                }
6364                (h, None)
6365            };
6366            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
6367
6368            let mixed = match &layer.mixer {
6369                Mixer::Full(fa) => self.full_attn_verify(
6370                    e,
6371                    fa,
6372                    &h,
6373                    h_q8_ref,
6374                    pos_d,
6375                    t,
6376                    cache,
6377                    il,
6378                    stream.map(|(_, c)| c),
6379                )?,
6380                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6381                Mixer::Linear(la) => {
6382                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
6383                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
6384                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
6385                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
6386                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
6387                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
6388                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
6389                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
6390                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
6391                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
6392                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
6393                    if (t >= 3 || (t == 2 && spec_m2()))
6394                        && mixer_fast
6395                        && e.uses_q8_1_fast(&la.ssm_out)
6396                    {
6397                        let want = ckpt.is_some();
6398                        let (out, stash) =
6399                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
6400                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6401                            ck.gdn[il] = Some(st);
6402                        }
6403                        out
6404                    } else {
6405                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
6406                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6407                            if ckpt.is_some() && t >= 2 {
6408                                Some(Vec::with_capacity(t - 1))
6409                            } else {
6410                                None
6411                            };
6412                        for col in 0..t {
6413                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
6414                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
6415                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
6416                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
6417                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
6418                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
6419                            // (pure dtod — cannot change any computed value). Last column skipped:
6420                            // rebuild targets are j <= t-1 columns.
6421                            if let Some(cs) = col_states.as_mut() {
6422                                if col + 1 < t {
6423                                    let rl = cache.recur[il].as_ref().unwrap();
6424                                    cs.push((
6425                                        e.clone_dtod(&rl.conv_state)?,
6426                                        e.clone_dtod(&rl.ssm_state)?,
6427                                    ));
6428                                }
6429                            }
6430                        }
6431                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
6432                            // ReplaySSM-assessment instrumentation (2026-07-30): the
6433                            // per-column clones are the only true state snapshots left in
6434                            // the verify (the batched path stashes INPUTS and replays).
6435                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
6436                                static ONCE: std::sync::Once = std::sync::Once::new();
6437                                let bytes: usize =
6438                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
6439                                ONCE.call_once(|| eprintln!(
6440                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
6441                                    cs.len(), bytes as f64 / 1e6));
6442                            }
6443                            ck.cols[il] = Some(cs);
6444                        }
6445                        out
6446                    }
6447                }
6448            };
6449
6450            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
6451            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
6452            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
6453            let ffn_fuse = match &layer.ffn {
6454                crate::hybrid::Ffn::Dense {
6455                    ffn_gate, ffn_up, ..
6456                } => {
6457                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
6458                        && e.uses_q8_1_fast(ffn_gate)
6459                        && e.uses_q8_1_fast(ffn_up)
6460                }
6461                crate::hybrid::Ffn::Moe(_) => false,
6462            };
6463            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
6464            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
6465            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
6466            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
6467            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
6468            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
6469            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
6470            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
6471            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
6472            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
6473            // mirror decode's dispatch or spec self-consistency fails.
6474            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
6475            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
6476            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
6477            let mut z = e.zeros(0)?; // replaced below on the unfused arms
6478            let z_q8 = if fuse_q8 {
6479                Some(e.add_rms_norm_q8_1(
6480                    &x,
6481                    &mixed,
6482                    layer.post_attn_norm.float_data(),
6483                    &mut x1,
6484                    n_embd,
6485                    t,
6486                    eps,
6487                )?)
6488            } else {
6489                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
6490                if ffn_fuse {
6491                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
6492                    e.rms_norm_decode(
6493                        &x1,
6494                        layer.post_attn_norm.float_data(),
6495                        &mut zf,
6496                        n_embd,
6497                        t,
6498                        eps,
6499                    )?;
6500                } else {
6501                    e.add_rms_norm(
6502                        &x,
6503                        &mixed,
6504                        layer.post_attn_norm.float_data(),
6505                        &mut x1,
6506                        &mut zf,
6507                        n_embd,
6508                        t,
6509                        eps,
6510                    )?;
6511                }
6512                z = zf;
6513                None
6514            };
6515            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
6516            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
6517            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
6518            let ffn_out = match &layer.ffn {
6519                crate::hybrid::Ffn::Dense {
6520                    ffn_gate,
6521                    ffn_up,
6522                    ffn_down,
6523                } => {
6524                    let n_ff = ffn_gate.out_features();
6525                    if let Some((zq, zd)) = z_q8.as_ref() {
6526                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
6527                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
6528                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
6529                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
6530                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
6531                        // structure at nrows=t.
6532                        let pair =
6533                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
6534                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
6535                                None => None,
6536                            };
6537                        let (gate, gs, up, us) = match pair {
6538                            Some(x4) => x4,
6539                            None => (
6540                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
6541                                1.0, // scale already applied inside _pre
6542                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
6543                                1.0,
6544                            ),
6545                        };
6546                        if e.uses_q8_1_fast(ffn_down) {
6547                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
6548                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
6549                        } else {
6550                            let mut act = vbuf(e, t * n_ff)?;
6551                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
6552                            e.matmul_decode_exact(ffn_down, &act, t)?
6553                        }
6554                    } else {
6555                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
6556                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
6557                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
6558                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
6559                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
6560                        let (gate, up) =
6561                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
6562                                Some(pair) => pair,
6563                                None => (
6564                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
6565                                    e.matmul_decode_exact(ffn_up, &z, t)?,
6566                                ),
6567                            };
6568                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
6569                        Self::ffn_act_lim(
6570                            e,
6571                            &self.cfg,
6572                            &gate,
6573                            &up,
6574                            1.0,
6575                            1.0,
6576                            dense_lim,
6577                            &mut act,
6578                            t * n_ff,
6579                        )?;
6580                        e.matmul_decode_exact(ffn_down, &act, t)?
6581                    }
6582                }
6583                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
6584            };
6585            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
6586            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
6587            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
6588            pending = Some((x1, ffn_out));
6589        }
6590        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
6591        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
6592        if let Some((x1p, f1p)) = pending.take() {
6593            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6594            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
6595            x = x2;
6596        }
6597        Ok(x)
6598    }
6599    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
6600    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
6601    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
6602    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
6603    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
6604    /// ssm state exactly like T sequential decode steps.
6605    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
6606    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
6607    #[allow(clippy::too_many_arguments)]
6608    fn linear_attn_verify_t(
6609        &self,
6610        e: &Engine,
6611        la: &LinearAttnLayer,
6612        h: &CudaSlice<f32>,
6613        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
6614        t: usize,
6615        cache: &mut Cache,
6616        il: usize,
6617        want_stash: bool,
6618    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
6619        let cfg = &self.cfg;
6620        let ssm = cfg.ssm.as_ref().unwrap();
6621        let d_state = ssm.state_size as usize;
6622        let num_k = ssm.group_count as usize;
6623        let num_v = ssm.time_step_rank as usize;
6624        let d_conv = ssm.conv_kernel as usize;
6625        let key_dim = d_state * num_k;
6626        let conv_dim = key_dim * 2 + d_state * num_v;
6627        let eps = cfg.rms_eps;
6628        let scale = 1.0 / (d_state as f32).sqrt();
6629
6630        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
6631        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
6632        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
6633        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
6634        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
6635        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
6636        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
6637        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
6638        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
6639        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
6640        // Bit-identical per (tensor,token,row) — see spec_fused_t().
6641        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
6642        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
6643        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
6644        // and feeds every projection; the caller guaranteed all four input projections are
6645        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
6646        let h_q8_t = if h_q8.is_none()
6647            && spec_fused_t()
6648            && (2..=4).contains(&t)
6649            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
6650                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
6651        {
6652            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
6653        } else {
6654            None
6655        };
6656        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
6657        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
6658            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
6659        let (qkv_mixed, z) = {
6660            let mut fused = None;
6661            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
6662                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
6663                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
6664            } else if let Some((hq, hd)) = hq8_any {
6665                if spec_fused_t() && (2..=4).contains(&t) {
6666                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
6667                }
6668            }
6669            match (fused, hq8_any) {
6670                (Some(pair), _) => pair,
6671                (None, Some((hq, hd))) if h_q8.is_some() => (
6672                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
6673                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
6674                ),
6675                (None, _) => (
6676                    e.matmul_decode_exact(&la.wqkv, h, t)?,
6677                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
6678                ),
6679            }
6680        };
6681        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
6682        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
6683        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
6684        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
6685        let (beta_raw, alpha) = if t == 1 {
6686            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
6687            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
6688                Some(((mut b, bs), (mut a, as_))) => {
6689                    if bs != 1.0 {
6690                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
6691                    }
6692                    if as_ != 1.0 {
6693                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
6694                    }
6695                    (b, a)
6696                }
6697                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
6698                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
6699                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
6700                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
6701                    Some((b, a)) => (b, a),
6702                    None => (
6703                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
6704                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
6705                    ),
6706                },
6707            }
6708        } else {
6709            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
6710            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
6711            let mut nvfp4_fused = None;
6712            let mut q8_fused = None;
6713            if let Some((hq, hd)) = hq8_any {
6714                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
6715                    nvfp4_fused =
6716                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
6717                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
6718                        static ONCE: std::sync::Once = std::sync::Once::new();
6719                        ONCE.call_once(|| {
6720                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
6721                        });
6722                    }
6723                }
6724                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
6725                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
6726                }
6727            }
6728            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
6729                if bs != 1.0 {
6730                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
6731                }
6732                if as_ != 1.0 {
6733                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
6734                }
6735                (b, a)
6736            } else if let Some(pair) = q8_fused {
6737                pair
6738            } else {
6739                match hq8_any {
6740                    Some((hq, hd)) if h_q8.is_some() => (
6741                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
6742                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
6743                    ),
6744                    _ => (
6745                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
6746                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
6747                    ),
6748                }
6749            }
6750        };
6751
6752        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
6753        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
6754        let rl = cache.recur[il].as_mut().unwrap();
6755        let mut conv_out = e.uninit(conv_dim * t)?;
6756        e.ssm_conv1d_tm_state(
6757            &qkv_mixed,
6758            &mut rl.conv_state,
6759            la.ssm_conv1d.float_data(),
6760            &mut conv_out,
6761            conv_dim,
6762            t,
6763            d_conv,
6764        )?;
6765
6766        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
6767        let mut q_g = e.uninit(d_state * num_v * t)?;
6768        let mut k_g = e.uninit(d_state * num_v * t)?;
6769        let mut v_g = e.uninit(d_state * num_v * t)?;
6770        e.qkv_to_gdn_repack(
6771            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
6772        )?;
6773        let mut q_l2 = e.uninit(d_state * num_v * t)?;
6774        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
6775        let mut k_l2 = e.uninit(d_state * num_v * t)?;
6776        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
6777        let mut beta = e.uninit(t * num_v)?;
6778        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
6779        let mut g_log = e.uninit(t * num_v)?;
6780        e.gdn_glog(
6781            &alpha,
6782            la.ssm_dt.float_data(),
6783            la.ssm_a.float_data(),
6784            &mut g_log,
6785            num_v,
6786            t,
6787        )?;
6788
6789        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
6790        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
6791        let mut o = e.uninit(d_state * num_v * t)?;
6792        {
6793            let crate::cache::RecurLayer {
6794                ssm_state,
6795                ssm_state_alt,
6796                ..
6797            } = rl;
6798            e.gdn_scan_s128(
6799                &q_l2,
6800                &k_l2,
6801                &v_g,
6802                &g_log,
6803                &beta,
6804                ssm_state,
6805                ssm_state_alt,
6806                &mut o,
6807                num_v,
6808                t,
6809                scale,
6810            )?;
6811        }
6812        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
6813
6814        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
6815        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
6816        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
6817        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
6818        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
6819        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
6820        let out = if e.uses_q8_1_fast(&la.ssm_out) {
6821            let (gq, gd) =
6822                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
6823            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
6824        } else {
6825            let mut gn = e.uninit(d_state * num_v * t)?;
6826            e.gated_rmsnorm(
6827                &o,
6828                la.ssm_norm.float_data(),
6829                &z,
6830                &mut gn,
6831                d_state,
6832                num_v * t,
6833                eps,
6834            )?;
6835            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
6836            // would fall to dp4a with a different FP reduction order — same class of bug as
6837            // the input projs).
6838            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
6839        };
6840        let stash = if want_stash {
6841            Some(GdnStash {
6842                qkv_mixed,
6843                q_l2,
6844                k_l2,
6845                v_g,
6846                g_log,
6847                beta,
6848            })
6849        } else {
6850            None
6851        };
6852        Ok((out, stash))
6853    }
6854
6855    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
6856    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
6857    /// - Full-attn KV: truncate len to snapshot + j. The verify's appended rows for those columns
6858    ///   are bit-identical to what an eager T=1 chain writes (the decode-exact contract the
6859    ///   verify-probe gates), so keeping them == replaying them.
6860    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
6861    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
6862    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
6863    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
6864    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
6865    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
6866    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
6867    fn commit_verified_prefix(
6868        &self,
6869        e: &Engine,
6870        cache: &mut Cache,
6871        snap: &crate::cache::CacheSnapshot,
6872        ckpt: &VerifyCkpt,
6873        j: usize,
6874        kv_lens_done: bool,
6875        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
6876    ) -> Result<(), Box<dyn std::error::Error>> {
6877        let cfg = &self.cfg;
6878        let ssm = cfg.ssm.as_ref().unwrap();
6879        let d_state = ssm.state_size as usize;
6880        let num_k = ssm.group_count as usize;
6881        let num_v = ssm.time_step_rank as usize;
6882        let d_conv = ssm.conv_kernel as usize;
6883        let conv_dim = d_state * num_k * 2 + d_state * num_v;
6884        let scale = 1.0 / (d_state as f32).sqrt();
6885        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
6886        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
6887        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
6888        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
6889        // buffers and stream order are identical to the per-layer memcpy sequence; the
6890        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
6891        let mut batched_cols = false;
6892        if state_copy_batch_on() && dev_j.is_none() {
6893            use cudarc::driver::DevicePtr;
6894            let s = &e.gpu.stream();
6895            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
6896            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
6897            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
6898            let mut uniform = true;
6899            for il in 0..self.layers.len() {
6900                let Some(rl) = cache.recur[il].as_ref() else {
6901                    continue;
6902                };
6903                if ckpt.gdn[il].is_some() {
6904                    continue; // kernel-rebuild arm restores below, per layer
6905                }
6906                let Some(cols) = &ckpt.cols[il] else {
6907                    continue; // missing-ckpt error surfaces in the main loop
6908                };
6909                let (c, st) = &cols[j - 1];
6910                if conv_pairs.is_empty() {
6911                    conv_words = c.len();
6912                    ssm_words = st.len();
6913                } else if c.len() != conv_words || st.len() != ssm_words {
6914                    uniform = false;
6915                    break;
6916                }
6917                let (pc, _g0) = c.device_ptr(s);
6918                let (dc, _g1) = rl.conv_state.device_ptr(s);
6919                let (ps, _g2) = st.device_ptr(s);
6920                let (ds, _g3) = rl.ssm_state.device_ptr(s);
6921                conv_pairs.push((pc as u64, dc as u64));
6922                ssm_pairs.push((ps as u64, ds as u64));
6923            }
6924            if uniform && !conv_pairs.is_empty() {
6925                let n = conv_pairs.len();
6926                let mut t = vec![0u64; 2 * n];
6927                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
6928                    t[k] = src;
6929                    t[n + k] = dst;
6930                }
6931                let conv_t = e.htod_u64(&t)?;
6932                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
6933                    t[k] = src;
6934                    t[n + k] = dst;
6935                }
6936                let ssm_t = e.htod_u64(&t)?;
6937                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
6938                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
6939                batched_cols = true;
6940            }
6941        }
6942        for il in 0..self.layers.len() {
6943            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6944                kvl.len = saved + j;
6945                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
6946                if !kv_lens_done {
6947                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6948                }
6949            }
6950            if let Some(rl) = cache.recur[il].as_mut() {
6951                if let Some(st) = &ckpt.gdn[il] {
6952                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
6953                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
6954                    if let Some((acc, base, t_v)) = dev_j {
6955                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
6956                        e.ssm_conv_ring_rebuild_dc(
6957                            &st.qkv_mixed,
6958                            ring_old,
6959                            &mut rl.conv_state,
6960                            conv_dim,
6961                            acc,
6962                            base,
6963                            t_v,
6964                            d_conv,
6965                        )?;
6966                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
6967                        e.gdn_scan_s128_dc(
6968                            &st.q_l2,
6969                            &st.k_l2,
6970                            &st.v_g,
6971                            &st.g_log,
6972                            &st.beta,
6973                            state_in,
6974                            &mut rl.ssm_state,
6975                            &mut o,
6976                            num_v,
6977                            acc,
6978                            base,
6979                            t_v,
6980                            scale,
6981                        )?;
6982                    } else {
6983                        e.ssm_conv_ring_rebuild(
6984                            &st.qkv_mixed,
6985                            ring_old,
6986                            &mut rl.conv_state,
6987                            conv_dim,
6988                            j,
6989                            d_conv,
6990                        )?;
6991                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
6992                        e.gdn_scan_s128(
6993                            &st.q_l2,
6994                            &st.k_l2,
6995                            &st.v_g,
6996                            &st.g_log,
6997                            &st.beta,
6998                            state_in,
6999                            &mut rl.ssm_state,
7000                            &mut o,
7001                            num_v,
7002                            j,
7003                            scale,
7004                        )?;
7005                    }
7006                } else if let Some(cols) = &ckpt.cols[il] {
7007                    if !batched_cols {
7008                        let (c, s) = &cols[j - 1];
7009                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
7010                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
7011                    }
7012                } else {
7013                    return Err(
7014                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
7015                    );
7016                }
7017            }
7018        }
7019        cache.pos = snap.pos + j;
7020        Ok(())
7021    }
7022
7023    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
7024    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
7025    fn commit_verified_prefix_stream(
7026        &self,
7027        e: &Engine,
7028        cache: &mut Cache,
7029        snap: &crate::cache::CacheSnapshot,
7030        ckpt: &VerifyCkpt,
7031        acc: &CudaSlice<u32>,
7032        base: usize,
7033        t_v: usize,
7034    ) -> Result<(), Box<dyn std::error::Error>> {
7035        let cfg = &self.cfg;
7036        let ssm = cfg.ssm.as_ref().unwrap();
7037        let d_state = ssm.state_size as usize;
7038        let num_k = ssm.group_count as usize;
7039        let num_v = ssm.time_step_rank as usize;
7040        let d_conv = ssm.conv_kernel as usize;
7041        let conv_dim = d_state * num_k * 2 + d_state * num_v;
7042        let scale = 1.0 / (d_state as f32).sqrt();
7043        for il in 0..self.layers.len() {
7044            if let Some(rl) = cache.recur[il].as_mut() {
7045                let st = ckpt.gdn[il]
7046                    .as_ref()
7047                    .ok_or("stream restore: batched-linear stash missing")?;
7048                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7049                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7050                e.ssm_conv_ring_rebuild_dc(
7051                    &st.qkv_mixed,
7052                    ring_old,
7053                    &mut rl.conv_state,
7054                    conv_dim,
7055                    acc,
7056                    base,
7057                    t_v,
7058                    d_conv,
7059                )?;
7060                let mut o = e.uninit(d_state * num_v * t_v)?;
7061                e.gdn_scan_s128_dc(
7062                    &st.q_l2,
7063                    &st.k_l2,
7064                    &st.v_g,
7065                    &st.g_log,
7066                    &st.beta,
7067                    state_in,
7068                    &mut rl.ssm_state,
7069                    &mut o,
7070                    num_v,
7071                    acc,
7072                    base,
7073                    t_v,
7074                    scale,
7075                )?;
7076            }
7077        }
7078        Ok(())
7079    }
7080
7081    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
7082    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
7083    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
7084    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
7085    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
7086    pub fn decode_step_t_aux2(
7087        &self,
7088        e: &Engine,
7089        tokens: &[u32],
7090        pos0: usize,
7091        cache: &mut Cache,
7092        aux_layers: &[usize],
7093        pred_col: Option<usize>,
7094    ) -> Result<
7095        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
7096        Box<dyn std::error::Error>,
7097    > {
7098        let cfg = &self.cfg;
7099        let n_embd = cfg.n_embd as usize;
7100        let eps = cfg.rms_eps;
7101        let t = tokens.len();
7102        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7103        let pos_d = e.htod_i32(&pos_vec)?;
7104        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7105        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
7106        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
7107        let want_pred = pred_col.is_some();
7108
7109        for (il, layer) in self.layers.iter().enumerate() {
7110            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
7111            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7112            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7113            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7114            if norm_fused {
7115                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7116            } else {
7117                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7118            }
7119            let mixed = match &layer.mixer {
7120                Mixer::Full(fa) => {
7121                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
7122                }
7123                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7124                Mixer::Linear(la) => {
7125                    let mut out = e.zeros(t * n_embd)?;
7126                    for col in 0..t {
7127                        let mut h_col = e.zeros(n_embd)?;
7128                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
7129                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7130                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7131                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7132                    }
7133                    out
7134                }
7135            };
7136            let ffn_fuse = match &layer.ffn {
7137                crate::hybrid::Ffn::Dense {
7138                    ffn_gate, ffn_up, ..
7139                } => {
7140                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7141                        && e.uses_q8_1_fast(ffn_gate)
7142                        && e.uses_q8_1_fast(ffn_up)
7143                }
7144                crate::hybrid::Ffn::Moe(_) => false,
7145            };
7146            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
7147            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7148            if ffn_fuse {
7149                e.add(&x, &mixed, &mut x1, t * n_embd)?;
7150                e.rms_norm_decode(
7151                    &x1,
7152                    layer.post_attn_norm.float_data(),
7153                    &mut z,
7154                    n_embd,
7155                    t,
7156                    eps,
7157                )?;
7158            } else {
7159                e.add_rms_norm(
7160                    &x,
7161                    &mixed,
7162                    layer.post_attn_norm.float_data(),
7163                    &mut x1,
7164                    &mut z,
7165                    n_embd,
7166                    t,
7167                    eps,
7168                )?;
7169            }
7170            let ffn_out = match &layer.ffn {
7171                crate::hybrid::Ffn::Dense {
7172                    ffn_gate,
7173                    ffn_up,
7174                    ffn_down,
7175                } => {
7176                    let n_ff = ffn_gate.out_features();
7177                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
7178                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
7179                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7180                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
7181                    Self::ffn_act_lim(
7182                        e,
7183                        &self.cfg,
7184                        &gate,
7185                        &up,
7186                        1.0,
7187                        1.0,
7188                        self.cfg.clamp_shexp_at(il as u32),
7189                        &mut act,
7190                        t * n_ff,
7191                    )?;
7192                    e.matmul_decode_exact(ffn_down, &act, t)?
7193                }
7194                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7195            };
7196            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7197            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7198            if aux_layers.contains(&il) {
7199                let mut a = e.zeros(n_embd)?;
7200                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
7201                aux_last.push(a);
7202                if let Some(pc) = pred_col {
7203                    let mut ap = e.zeros(n_embd)?;
7204                    e.copy_view_into(
7205                        &mut ap,
7206                        0,
7207                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
7208                        n_embd,
7209                    )?;
7210                    aux_pred.push(ap);
7211                }
7212            }
7213            x = x2;
7214        }
7215        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
7216        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7217        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
7218        let host = e.dtoh(&logits)?;
7219        cache.pos += t;
7220        Ok((
7221            host,
7222            aux_last,
7223            if want_pred { Some(aux_pred) } else { None },
7224        ))
7225    }
7226
7227    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
7228    /// `step35_decode_attn`.
7229    ///
7230    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
7231    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
7232    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
7233    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
7234    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
7235    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
7236    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
7237    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
7238    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
7239    /// position of each query row. A batched twin would have to reproduce all of that AND the
7240    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
7241    /// take one `base_len`, not a per-row offset).
7242    ///
7243    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
7244    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
7245    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
7246    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
7247    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
7248    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
7249    /// step35 twin is a perf lane's job and must be gated against this arm.
7250    ///
7251    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
7252    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
7253    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
7254    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
7255    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
7256    #[allow(clippy::too_many_arguments)]
7257    fn step35_verify(
7258        &self,
7259        e: &Engine,
7260        fa: &FullAttnLayer,
7261        h: &CudaSlice<f32>,
7262        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7263        t: usize,
7264        cache: &mut Cache,
7265        il: usize,
7266    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7267        let n_embd = self.cfg.n_embd as usize;
7268        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
7269        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
7270        // site by the `Mixer::Full(_) if self.cfg.step35.is_some() => false` arm of
7271        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
7272        // cannot regress it into silently reading an empty buffer.
7273        assert_eq!(
7274            h.len(),
7275            t * n_embd,
7276            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
7277             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
7278            h_q8.is_some()
7279        );
7280        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
7281        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
7282        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
7283        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
7284        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
7285        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
7286        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
7287        for r in 0..t {
7288            // Absolute position of this query row. `cache.pos` is the committed length at round
7289            // start and every row before r has already been appended by this loop, so the r-th
7290            // verify token sits at cache.pos + r — the same position eager decode would give it.
7291            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
7292            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
7293            e.copy_view_into(
7294                &mut h_row,
7295                0,
7296                &h.slice(r * n_embd..(r + 1) * n_embd),
7297                n_embd,
7298            )?;
7299            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
7300            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
7301            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
7302            debug_assert_eq!(
7303                o.len(),
7304                n_embd,
7305                "step35_decode_attn returns post-wo [n_embd]"
7306            );
7307            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
7308        }
7309        Ok(out)
7310    }
7311
7312    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
7313    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
7314    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
7315    #[allow(clippy::too_many_arguments)]
7316    fn full_attn_verify(
7317        &self,
7318        e: &Engine,
7319        fa: &FullAttnLayer,
7320        h: &CudaSlice<f32>,
7321        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7322        pos_d: &CudaSlice<i32>,
7323        t: usize,
7324        cache: &mut Cache,
7325        il: usize,
7326        stream_ctr: Option<&CudaSlice<i32>>,
7327    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7328        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
7329        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
7330        // its own arm. A verify that silently computes different attention than decode defeats the
7331        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
7332        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
7333        // shape and not laziness.
7334        if self.cfg.step35.is_some() {
7335            if stream_ctr.is_some() {
7336                return Err(
7337                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7338                            cannot express the SWA offset KV view; same root cause as the dc \
7339                            decode refusal) — run spec without the stream arm"
7340                        .into(),
7341                );
7342            }
7343            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
7344        }
7345        let cfg = &self.cfg;
7346        let geometry = cfg.full_attention_geometry_at(il as u32);
7347        let n_head = geometry.n_head as usize;
7348        let n_head_kv = geometry.n_head_kv as usize;
7349        let head_dim = geometry.head_dim_k as usize;
7350        let eps = cfg.rms_eps;
7351        let scale = geometry.attention_scale();
7352        let n_embd = cfg.n_embd as usize;
7353
7354        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
7355        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
7356        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
7357        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
7358        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
7359        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
7360        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
7361        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
7362        let (qf, mut k, v) = {
7363            let mut fused = None;
7364            let qkv_fast =
7365                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
7366            if t == 1 && qkv_fast {
7367                let (hq_o, hd_o);
7368                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7369                    Some(p) => p,
7370                    None => {
7371                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
7372                        (&hq_o, &hd_o)
7373                    }
7374                };
7375                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
7376            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
7377                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
7378                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
7379                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
7380                let (hq_o, hd_o);
7381                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7382                    Some(p) => p,
7383                    None => {
7384                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
7385                        (&hq_o, &hd_o)
7386                    }
7387                };
7388                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
7389            }
7390            match (fused, h_q8) {
7391                (Some(triple), _) => triple,
7392                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
7393                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
7394                (None, Some((hq, hd))) if qkv_fast => (
7395                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
7396                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
7397                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
7398                ),
7399                (None, _) => (
7400                    e.matmul_decode_exact(&fa.wq, h, t)?,
7401                    e.matmul_decode_exact(&fa.wk, h, t)?,
7402                    e.matmul_decode_exact(&fa.wv, h, t)?,
7403                ),
7404            }
7405        };
7406        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
7407        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7408        let (mut q, gate) = if gated {
7409            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7410            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7411            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7412            (q, Some(gate))
7413        } else {
7414            (qf, None)
7415        };
7416
7417        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
7418        e.rms_norm(
7419            &q,
7420            fa.q_norm.float_data(),
7421            &mut qn,
7422            head_dim,
7423            n_head * t,
7424            eps,
7425        )?;
7426        q = qn;
7427        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
7428        e.rms_norm(
7429            &k,
7430            fa.k_norm.float_data(),
7431            &mut kn,
7432            head_dim,
7433            n_head_kv * t,
7434            eps,
7435        )?;
7436        k = kn;
7437        let rope_dims = geometry.n_rot as usize;
7438        e.rope_neox(
7439            &mut q,
7440            pos_d,
7441            head_dim,
7442            rope_dims,
7443            n_head,
7444            t,
7445            geometry.rope_base,
7446            1.0,
7447        )?;
7448        e.rope_neox(
7449            &mut k,
7450            pos_d,
7451            head_dim,
7452            rope_dims,
7453            n_head_kv,
7454            t,
7455            geometry.rope_base,
7456            1.0,
7457        )?;
7458
7459        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
7460        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
7461        let kvl = cache.kv[il].as_mut().unwrap();
7462        let (kv_dim_k, kv_dim_v, ktb, vtb) =
7463            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
7464        if let Some(ctr) = stream_ctr {
7465            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
7466            // math on a (block, token) grid, documented byte-identical); host len is a stale
7467            // LOWER BOUND under pre-issue (drain reconciles it).
7468            e.append_kv_quantized_rows_dc(
7469                &k,
7470                &v,
7471                &mut kvl.k,
7472                &mut kvl.v,
7473                ctr,
7474                t,
7475                kv_dim_k,
7476                kv_dim_v,
7477                ktb,
7478                vtb,
7479                crate::Engine::kv_fp8_on(),
7480            )?;
7481        } else {
7482            for i in 0..t {
7483                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
7484                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
7485                e.append_kv_quantized_view(
7486                    &k_row,
7487                    &v_row,
7488                    &mut kvl.k,
7489                    &mut kvl.v,
7490                    kvl.len + i,
7491                    kv_dim_k,
7492                    kv_dim_v,
7493                    ktb,
7494                    vtb,
7495                    crate::Engine::kv_fp8_on(),
7496                )?;
7497            }
7498            kvl.len += t;
7499        }
7500
7501        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
7502        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
7503        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
7504        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
7505        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
7506        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
7507        // keys. The verify appends all T tokens first but bounds the key range per row.
7508        //
7509        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
7510        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
7511        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
7512        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
7513        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
7514        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
7515        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
7516        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
7517        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
7518        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
7519        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
7520        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
7521        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
7522        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
7523        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
7524        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
7525        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
7526        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
7527        if let Some(ctr) = stream_ctr {
7528            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
7529            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
7530            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
7531            let upper = kvl.len + t + 64;
7532            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
7533            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
7534            e.fa_decode_rows_dc(
7535                &q,
7536                &k_view,
7537                &v_view,
7538                &mut attn,
7539                head_dim,
7540                n_head,
7541                n_head_kv,
7542                ctr,
7543                upper.min(cache.max_ctx),
7544                t,
7545                scale,
7546                ktb,
7547                vtb,
7548                0,
7549                false,
7550            )?;
7551        } else if spec_lean() && t == 1 {
7552            let t_kv = base_len + 1;
7553            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
7554            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
7555            e.fa_decode_kvmod(
7556                &q,
7557                &k_view,
7558                &v_view,
7559                &mut attn,
7560                head_dim,
7561                n_head,
7562                n_head_kv,
7563                t_kv,
7564                scale,
7565                ktb,
7566                vtb,
7567                crate::Engine::kv_fp8_on(),
7568            )?;
7569        } else if e.fa_rows_eligible(base_len, head_dim) {
7570            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
7571            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
7572            e.fa_decode_rows(
7573                &q,
7574                &k_view,
7575                &v_view,
7576                &mut attn,
7577                head_dim,
7578                n_head,
7579                n_head_kv,
7580                base_len,
7581                t,
7582                scale,
7583                ktb,
7584                vtb,
7585                None,
7586                false,
7587                crate::Engine::kv_fp8_on(),
7588                None,
7589            )?;
7590        } else {
7591            for r in 0..t {
7592                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
7593                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
7594                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
7595                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
7596                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
7597                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
7598                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
7599                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
7600                e.fa_decode_kvmod(
7601                    &q_row,
7602                    &k_view_r,
7603                    &v_view_r,
7604                    &mut attn_row,
7605                    head_dim,
7606                    n_head,
7607                    n_head_kv,
7608                    t_kv_r,
7609                    scale,
7610                    ktb,
7611                    vtb,
7612                    crate::Engine::kv_fp8_on(),
7613                )?;
7614                e.copy_into(
7615                    &mut attn,
7616                    r * n_head * head_dim,
7617                    &attn_row,
7618                    n_head * head_dim,
7619                )?;
7620            }
7621        }
7622
7623        let attn_g = match &gate {
7624            Some(gate) => {
7625                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
7626                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
7627                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
7628                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
7629                ag
7630            }
7631            None => attn,
7632        };
7633        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
7634        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
7635        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
7636    }
7637
7638    /// Context-linear bytes for a plain serving session's trunk cache.
7639    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
7640        crate::cache::cache_bytes_per_token(&self.cfg)
7641    }
7642
7643    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
7644    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
7645        (
7646            self.plain_session_kv_bytes_per_token(),
7647            crate::cache::cache_ring_bytes_per_token(&self.cfg),
7648            crate::cache::cache_ring_row_cap(&self.cfg),
7649        )
7650    }
7651
7652    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
7653    /// scratch. With no MTP head this equals the plain coefficient.
7654    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
7655        let scratch = self
7656            .mtp
7657            .as_ref()
7658            .map(|mtp| {
7659                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
7660                k + v
7661            })
7662            .unwrap_or(0);
7663        self.plain_session_kv_bytes_per_token()
7664            .saturating_add(scratch)
7665    }
7666
7667    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
7668    /// capped by the same SWA ring rows as the trunk.
7669    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
7670        let total = self.spec_session_kv_bytes_per_token();
7671        let (_, mut ring, rows) = self.plain_session_kv_shape();
7672        if rows > 0 {
7673            ring = ring.saturating_add(
7674                self.mtp
7675                    .as_ref()
7676                    .map(|mtp| {
7677                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
7678                        k + v
7679                    })
7680                    .unwrap_or(0),
7681            );
7682        }
7683        (total, ring, rows)
7684    }
7685
7686    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
7687    /// the NextN head to draft K tokens then verifies them in one batched target forward.
7688    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
7689    /// acceptance rate. `k` = draft length per round.
7690    ///
7691    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
7692    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
7693    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
7694    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
7695    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
7696    /// captured graph references is event-free; the spec loop is strictly single-stream.
7697    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
7698    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
7699    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
7700    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
7701    /// generate_spec_inner2.
7702    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
7703    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
7704    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
7705    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
7706    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
7707    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
7708    pub fn new_session(
7709        &self,
7710        e: &Engine,
7711        max_ctx: usize,
7712    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
7713        Ok(SpecSession {
7714            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
7715            // is the SERVING spec-session path, and with the ppN door open across two cards a
7716            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
7717            // round — the wrong-card class already fixed on the two batched serving paths
7718            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
7719            // branch, same allocations), so single-device behavior is byte-unchanged.
7720            cache: crate::pp::new_cache(e, &self.cfg, max_ctx)?,
7721            scratch: MtpScratch::new(
7722                e,
7723                &self.cfg,
7724                max_ctx,
7725                self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7726            )?,
7727            committed: Vec::new(),
7728            last_h: None,
7729            next_pred: None,
7730            sctr: 0,
7731            uctr: 0,
7732            draft_ctx: None,
7733            pending_tok: None,
7734            turn_ckpt: None,
7735            telem: SpecTelemetryCounters::default(),
7736            capture_at: None,
7737            boundary_captures: Vec::new(),
7738            ckpt_at: None,
7739        })
7740    }
7741
7742    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
7743    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
7744    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
7745    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
7746    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
7747    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
7748    /// worker always receives a fully-warm continuation session (committed = whole
7749    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
7750    /// boundary logits on the empty-suffix shape).
7751    ///
7752    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
7753    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
7754    /// request, and plain feeds a carried suffix via eager `decode_step` below
7755    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
7756    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
7757    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
7758    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
7759    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
7760    /// burst prime.
7761    ///
7762    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
7763    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
7764    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
7765    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
7766    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
7767    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
7768    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
7769    /// cold session draws from the identical row at counter 0 and then runs its rounds from
7770    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
7771    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
7772    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
7773    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
7774    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
7775    ///
7776    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
7777    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
7778    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
7779    /// and are never routed here.
7780    ///
7781    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
7782    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
7783    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
7784    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
7785    /// entry stays published for the next request.
7786    #[allow(clippy::too_many_arguments)]
7787    pub fn spec_session_from_restored(
7788        &self,
7789        e: &Engine,
7790        mut cache: Cache,
7791        prefix: Vec<u32>,
7792        suffix: &[u32],
7793        draft_k: &CudaSlice<u8>,
7794        draft_v: &CudaSlice<u8>,
7795        draft_k_tok_bytes: usize,
7796        draft_v_tok_bytes: usize,
7797        draft_len: usize,
7798        last_h: &[f32],
7799        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
7800        // when a suffix follows — the feed's own logits are the boundary then.
7801        boundary_logits: &[f32],
7802        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
7803        // ONE place instead of being half-applied by the worker.
7804        sampling: Option<SpecSampling>,
7805        require_anchor: bool,
7806        max_ctx: usize,
7807        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
7808        // prompt position to split the suffix feed at and capture the extended-entry
7809        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
7810        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
7811        // WHY: the prompt-end capture below includes the template's live generation header
7812        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
7813        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
7814        // diverged from every future prompt and the hit boundary FROZE at the first
7815        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
7816        republish_at: Option<usize>,
7817    ) -> Result<SpecSession, (Option<Cache>, String)> {
7818        let pos = prefix.len();
7819        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
7820            Err((Some(cache), msg))
7821        };
7822        if self.mtp.is_none() {
7823            return fail(cache, "no MTP head attached (nothing to draft with)".into());
7824        }
7825        if pos == 0 {
7826            return fail(cache, "empty committed prefix".into());
7827        }
7828        if cache.pos != pos {
7829            let msg = format!(
7830                "restored cache pos {} != restored prefix len {pos}",
7831                cache.pos
7832            );
7833            return fail(cache, msg);
7834        }
7835        if draft_len != pos {
7836            return fail(
7837                cache,
7838                format!("draft plane len {draft_len} != restored prefix len {pos}"),
7839            );
7840        }
7841        if pos + suffix.len() >= max_ctx {
7842            return fail(
7843                cache,
7844                format!(
7845                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
7846                    pos + suffix.len(),
7847                ),
7848            );
7849        }
7850        let mut scratch = match MtpScratch::new(
7851            e,
7852            &self.cfg,
7853            max_ctx,
7854            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
7855        ) {
7856            Ok(s) => s,
7857            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
7858        };
7859        if scratch.kv.ring.is_some() {
7860            return fail(
7861                cache,
7862                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
7863            );
7864        }
7865        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
7866            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
7867        {
7868            return fail(
7869                cache,
7870                format!(
7871                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
7872                     {}/{} bytes/token (stale entry across a format change)",
7873                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
7874                ),
7875            );
7876        }
7877        if pos > scratch.cap {
7878            return fail(
7879                cache,
7880                format!(
7881                    "draft plane rows {pos} exceed scratch capacity {}",
7882                    scratch.cap
7883                ),
7884            );
7885        }
7886        let kb = pos * draft_k_tok_bytes;
7887        let vb = pos * draft_v_tok_bytes;
7888        if draft_k.len() < kb || draft_v.len() < vb {
7889            return fail(
7890                cache,
7891                format!(
7892                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
7893                    draft_k.len(),
7894                    draft_v.len(),
7895                ),
7896            );
7897        }
7898        if kb > 0 {
7899            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
7900                return fail(cache, format!("draft K restore copy failed: {err}"));
7901            }
7902        }
7903        if vb > 0 {
7904            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
7905                return fail(cache, format!("draft V restore copy failed: {err}"));
7906            }
7907        }
7908        if let Err(err) = scratch.set_len(e, pos) {
7909            return fail(cache, format!("draft scratch len set failed: {err}"));
7910        }
7911        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
7912            // anchor upload failure is acceptance-only when a suffix feed follows (fill
7913            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
7914            // burst entry asserts committed + last_h + next_pred) — the caller says which.
7915            e.htod(last_h).ok()
7916        } else {
7917            None
7918        };
7919        if require_anchor && last_h_dev.is_none() {
7920            return fail(
7921                cache,
7922                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
7923            );
7924        }
7925        let mut committed = prefix;
7926        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
7927        // what the empty-suffix continuation assert in the burst entry requires.
7928        let next_pred;
7929        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
7930        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
7931        // drawing its own first token from the same row.
7932        let mut sctr = 0u32;
7933        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
7934        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
7935        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
7936        // after the suffix joins `committed` below.
7937        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
7938        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
7939        if !suffix.is_empty() {
7940            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
7941            // From here on the trunk cache mutates: failures return Err((None, _)) and
7942            // the worker serves the request cold-plain instead of reusing the carrier.
7943            let dirty =
7944                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
7945            let n_embd = self.cfg.n_embd as usize;
7946            let t = suffix.len();
7947            let mut h_rows = match e.uninit(t * n_embd) {
7948                Ok(b) => b,
7949                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
7950            };
7951            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
7952            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
7953            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
7954            let b_rel = republish_at
7955                .and_then(|abs| abs.checked_sub(pos))
7956                .filter(|&r| r > 0 && r < t);
7957            let mut feed_logits = Vec::new();
7958            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
7959                || e.frozen_cpu_experts_prefer_tokenwise_prime();
7960            let mut fed = 0usize;
7961            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
7962                if seg_end <= fed {
7963                    continue;
7964                }
7965                let seg = &suffix[fed..seg_end];
7966                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
7967                if batched {
7968                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
7969                    // queued after this segment ride `queued_after` so Step35 arm selection
7970                    // stays keyed to the request's end (tick-seg law).
7971                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
7972                        Ok((l, _h_seed, hiddens)) => {
7973                            if let Err(err) =
7974                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
7975                            {
7976                                return dirty(format!("suffix hidden copy: {err}"));
7977                            }
7978                            feed_logits = l;
7979                        }
7980                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
7981                    }
7982                } else {
7983                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
7984                    for (i, &tok) in seg.iter().enumerate() {
7985                        match self.decode_step_h(e, tok, &mut cache) {
7986                            Ok((l, h)) => {
7987                                if let Err(err) =
7988                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
7989                                {
7990                                    return dirty(format!("suffix hidden copy: {err}"));
7991                                }
7992                                feed_logits = l;
7993                            }
7994                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
7995                        }
7996                    }
7997                }
7998                fed = seg_end;
7999                if Some(seg_end) == b_rel {
8000                    // The stable pre-generation boundary: capture the extended-entry
8001                    // publication AND this session's own turn checkpoint here instead of at
8002                    // prompt-end (both would otherwise carry the volatile live-header tail
8003                    // the next re-render replaces). Failure silent, turn_ckpt convention.
8004                    debug_assert_eq!(
8005                        cache.pos,
8006                        pos + seg_end,
8007                        "stable-boundary capture off the feed split"
8008                    );
8009                    if spec_restore_republish_on() {
8010                        if let Ok(snap) = cache.snapshot(e) {
8011                            boundary_captures.push(SpecBoundaryCapture {
8012                                snap,
8013                                pos: pos + seg_end,
8014                                logits: feed_logits.clone(),
8015                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
8016                            });
8017                        }
8018                    }
8019                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8020                        e.uninit(n_embd).and_then(|mut a| {
8021                            e.copy_view_into(
8022                                &mut a,
8023                                0,
8024                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
8025                                n_embd,
8026                            )?;
8027                            Ok(a)
8028                        });
8029                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
8030                        restored_turn_ckpt = Some(SpecCheckpoint {
8031                            snap,
8032                            pos: pos + seg_end,
8033                            last_h,
8034                        });
8035                    }
8036                }
8037            }
8038            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
8039            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
8040            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
8041            // with T). Fill failures are acceptance-only — truncate to the restored rows
8042            // and continue; the burst's own set_len keeps the invariant.
8043            let mtp = self.mtp.as_ref().expect("mtp checked above");
8044            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8045            let embd_gpu = if spec_host_embd() {
8046                None
8047            } else {
8048                Some(
8049                    self.embd_gpu
8050                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8051                )
8052            };
8053            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8054            let fill_chunk = 4096usize;
8055            let mut filled = true;
8056            let mut start = 0usize;
8057            'fill: while start < t {
8058                let end = (start + fill_chunk).min(t);
8059                let tc = end - start;
8060                let Ok(mut phs) = e.zeros(tc * n_embd) else {
8061                    filled = false;
8062                    break 'fill;
8063                };
8064                let (src_lo, dst_off, n_copy) = if start == 0 {
8065                    (0, n_embd, (tc - 1) * n_embd)
8066                } else {
8067                    ((start - 1) * n_embd, 0, tc * n_embd)
8068                };
8069                if start == 0 {
8070                    if let Some(lh) = last_h_dev.as_ref() {
8071                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
8072                            filled = false;
8073                            break 'fill;
8074                        }
8075                    }
8076                }
8077                if n_copy > 0
8078                    && e.copy_view_into(
8079                        &mut phs,
8080                        dst_off,
8081                        &h_rows.slice(src_lo..src_lo + n_copy),
8082                        n_copy,
8083                    )
8084                    .is_err()
8085                {
8086                    filled = false;
8087                    break 'fill;
8088                }
8089                if self
8090                    .mtp_kv_fill(
8091                        e,
8092                        mtp,
8093                        &suffix[start..end],
8094                        &phs,
8095                        pos + start,
8096                        &mut scratch,
8097                        embd_dev,
8098                    )
8099                    .is_err()
8100                {
8101                    filled = false;
8102                    break 'fill;
8103                }
8104                start = end;
8105            }
8106            if !filled {
8107                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
8108                // so keep only the restored rows resident and let verify arbitrate.
8109                if let Err(err) = scratch.set_len(e, pos) {
8110                    return dirty(format!("scratch truncation after failed fill: {err}"));
8111                }
8112            }
8113            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
8114            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
8115            // finding (d)). Pre-lane, publication was armed only for COLD sessions
8116            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
8117            // non-continuation burst — but a converted hit's first burst IS a continuation,
8118            // so a growing conversation learned exactly ONE boundary and turn 3 could never
8119            // hit a longer prefix than turn 2 did.
8120            //
8121            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
8122            // line — the trunk is primed over the whole prompt, nothing is generated, and the
8123            // draft plane rows [0..prompt) are filled just above. That is a complete
8124            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
8125            // publishes; the worker's existing publication sweep picks it up because it is
8126            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
8127            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
8128            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
8129            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
8130            // publication is an optimization, never a correctness dependency.
8131            //
8132            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
8133            // entry's tail is the live generation header the next re-render replaces, so on a
8134            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
8135            // the stable-boundary capture above IS this publication, minus the poisoned tail.
8136            if spec_restore_republish_on() && boundary_captures.is_empty() {
8137                debug_assert_eq!(
8138                    cache.pos,
8139                    pos + t,
8140                    "extended-entry capture must sit at the restored session's prompt end",
8141                );
8142                if let Ok(snap) = cache.snapshot(e) {
8143                    boundary_captures.push(SpecBoundaryCapture {
8144                        snap,
8145                        pos: pos + t,
8146                        logits: feed_logits.clone(),
8147                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
8148                    });
8149                }
8150            }
8151            // continuation seed: the feed's boundary logits ARE the plain path's boundary
8152            // logits (same program), so greedy's argmax here is plain's first emitted token,
8153            // and the sampled draw is the cold sampled session's own first token.
8154            next_pred = Some(if sampled {
8155                let sp = sampling.expect("sampled implies a sampler");
8156                // `committed` is still the restored prefix here; the suffix joins it below —
8157                // so this is the last-N window over the WHOLE prompt, exactly the cold
8158                // session's own window at its first token.
8159                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
8160                match sample_boundary_token(
8161                    e,
8162                    &feed_logits,
8163                    &sp,
8164                    &hist,
8165                    &mut sctr,
8166                    "restore-suffix-feed",
8167                ) {
8168                    Ok(t) => t,
8169                    // the trunk is already fed: hand nothing back, the worker serves the
8170                    // request cold-plain. Never fall back to an argmax — that would put a
8171                    // greedy token in a sampled stream to save a slow path.
8172                    Err(err) => {
8173                        return dirty(format!("boundary token draw failed: {err}"));
8174                    }
8175                }
8176            } else {
8177                argmax(&feed_logits) as u32
8178            });
8179            let mut lh = match e.uninit(n_embd) {
8180                Ok(b) => b,
8181                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
8182            };
8183            if let Err(err) = e.copy_view_into(
8184                &mut lh,
8185                0,
8186                &h_rows.slice((t - 1) * n_embd..t * n_embd),
8187                n_embd,
8188            ) {
8189                return dirty(format!("boundary hidden copy: {err}"));
8190            }
8191            last_h_dev = Some(lh);
8192            committed.extend_from_slice(suffix);
8193        } else {
8194            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
8195            // ENTRY's boundary logits are the boundary row, and this is the token the cold
8196            // session emits from that same row. Owned here rather than in the worker so the
8197            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
8198            if boundary_logits.is_empty() {
8199                return fail(
8200                    cache,
8201                    "full-cover restore without the entry's boundary logits".into(),
8202                );
8203            }
8204            next_pred = Some(if sampled {
8205                let sp = sampling.expect("sampled implies a sampler");
8206                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
8207                match sample_boundary_token(
8208                    e,
8209                    boundary_logits,
8210                    &sp,
8211                    &hist,
8212                    &mut sctr,
8213                    "restore-full-cover",
8214                ) {
8215                    Ok(t) => t,
8216                    // nothing has been mutated on this shape — hand the carrier back and let
8217                    // the hit serve PLAIN (the banked pre-lane path).
8218                    Err(err) => {
8219                        return fail(cache, format!("boundary token draw failed: {err}"));
8220                    }
8221                }
8222            } else {
8223                argmax(boundary_logits) as u32
8224            });
8225        }
8226        Ok(SpecSession {
8227            cache,
8228            scratch,
8229            committed,
8230            last_h: last_h_dev,
8231            next_pred,
8232            sctr,
8233            uctr: 0,
8234            draft_ctx: None,
8235            pending_tok: None,
8236            // Stable-boundary capture from the split feed above (None on the legacy shape):
8237            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
8238            // affinity probe declined ("no turn checkpoint retained") and the conversation
8239            // fell back to the frozen prefix entry forever.
8240            turn_ckpt: restored_turn_ckpt,
8241            telem: SpecTelemetryCounters::default(),
8242            capture_at: None,
8243            boundary_captures,
8244            ckpt_at: None,
8245        })
8246    }
8247
8248    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
8249    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
8250    /// snapshot, or draft-KV row that only corrupts the following round.
8251    pub fn optipipe_compare_session_state(
8252        &self,
8253        e: &Engine,
8254        reference: &SpecSession,
8255        candidate: &SpecSession,
8256    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
8257        fn fail(what: &str) -> Box<dyn std::error::Error> {
8258            format!("optipipe state mismatch: {what}").into()
8259        }
8260        fn same_f32(a: &[f32], b: &[f32]) -> bool {
8261            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
8262        }
8263        fn compare_layers(
8264            es: &Engine,
8265            range: std::ops::Range<usize>,
8266            reference: &SpecSession,
8267            candidate: &SpecSession,
8268            report: &mut OptiForkStateIdentity,
8269        ) -> Result<(), Box<dyn std::error::Error>> {
8270            for il in range {
8271                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
8272                    (Some(a), Some(b)) => {
8273                        if a.len != b.len {
8274                            return Err(fail(&format!(
8275                                "layer {il} host KV len {} != {}",
8276                                a.len, b.len
8277                            )));
8278                        }
8279                        let ad = es.dtoh_i32(&a.len_d)?;
8280                        let bd = es.dtoh_i32(&b.len_d)?;
8281                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
8282                            return Err(fail(&format!(
8283                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
8284                                a.len,
8285                            )));
8286                        }
8287                        let kb = a.len * a.k_tok_bytes;
8288                        let vb = a.len * a.v_tok_bytes;
8289                        if kb > 0 {
8290                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
8291                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
8292                            if ak != bk {
8293                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
8294                                return Err(fail(&format!(
8295                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
8296                                    at / a.k_tok_bytes,
8297                                    at % a.k_tok_bytes,
8298                                    ak[at],
8299                                    bk[at],
8300                                )));
8301                            }
8302                        }
8303                        if vb > 0 {
8304                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
8305                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
8306                            if av != bv {
8307                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
8308                                return Err(fail(&format!(
8309                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
8310                                    at / a.v_tok_bytes,
8311                                    at % a.v_tok_bytes,
8312                                    av[at],
8313                                    bv[at],
8314                                )));
8315                            }
8316                        }
8317                        report.trunk_kv_bytes += kb + vb;
8318                    }
8319                    (None, None) => {}
8320                    _ => return Err(fail(&format!("layer {il} KV presence"))),
8321                }
8322                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
8323                    (Some(a), Some(b)) => {
8324                        let ac = es.dtoh(&a.conv_state)?;
8325                        let bc = es.dtoh(&b.conv_state)?;
8326                        if !same_f32(&ac, &bc) {
8327                            return Err(fail(&format!("layer {il} conv state")));
8328                        }
8329                        let as_ = es.dtoh(&a.ssm_state)?;
8330                        let bs = es.dtoh(&b.ssm_state)?;
8331                        if !same_f32(&as_, &bs) {
8332                            return Err(fail(&format!("layer {il} SSM state")));
8333                        }
8334                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
8335                    }
8336                    (None, None) => {}
8337                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
8338                }
8339            }
8340            Ok(())
8341        }
8342
8343        if reference.committed != candidate.committed {
8344            return Err(fail("committed token ids"));
8345        }
8346        if reference.cache.pos != candidate.cache.pos
8347            || reference.cache.max_ctx != candidate.cache.max_ctx
8348        {
8349            return Err(fail("cache pos/capacity"));
8350        }
8351        if reference.pending_tok != candidate.pending_tok
8352            || reference.next_pred != candidate.next_pred
8353            || reference.sctr != candidate.sctr
8354            || reference.uctr != candidate.uctr
8355        {
8356            return Err(fail("pending/prediction/counter tail"));
8357        }
8358
8359        let mut report = OptiForkStateIdentity::default();
8360        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
8361            let rt = crate::pp::PpNRt::get(e)?;
8362            for stage in 0..rt.n_stages() {
8363                let _scope = rt.enter(stage);
8364                compare_layers(
8365                    rt.engine(stage, e),
8366                    fence[stage]..fence[stage + 1],
8367                    reference,
8368                    candidate,
8369                    &mut report,
8370                )?;
8371            }
8372        } else {
8373            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
8374        }
8375
8376        let (a, b) = (&reference.scratch.kv, &candidate.scratch.kv);
8377        if a.len != b.len || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)? {
8378            return Err(fail("draft scratch length"));
8379        }
8380        let kb = a.len * a.k_tok_bytes;
8381        let vb = a.len * a.v_tok_bytes;
8382        if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
8383            return Err(fail("draft scratch K bytes"));
8384        }
8385        if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
8386            return Err(fail("draft scratch V bytes"));
8387        }
8388        report.scratch_kv_bytes = kb + vb;
8389
8390        match (&reference.last_h, &candidate.last_h) {
8391            (Some(a), Some(b)) => {
8392                let ah = e.dtoh(a)?;
8393                let bh = e.dtoh(b)?;
8394                if !same_f32(&ah, &bh) {
8395                    return Err(fail("last hidden/seed bytes"));
8396                }
8397                report.hidden_bytes = ah.len() * 4;
8398            }
8399            (None, None) => {}
8400            _ => return Err(fail("last hidden/seed presence")),
8401        }
8402        Ok(report)
8403    }
8404
8405    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
8406    /// retained prompt-end checkpoint, so a request whose prompt matches
8407    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
8408    ///
8409    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
8410    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
8411    /// restored from the device copy taken there, draft scratch length reset, `committed`
8412    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
8413    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
8414    /// every burst after it are identical to a cold run of the same token stream — the
8415    /// committed-tokens-authoritative contract.
8416    ///
8417    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
8418    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
8419    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
8420    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
8421    /// (the scratch KV, the resident embedding), none of which the rewind moves.
8422    ///
8423    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
8424    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
8425    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
8426    pub fn spec_rewind_to_checkpoint(
8427        &self,
8428        e: &Engine,
8429        sess: &mut SpecSession,
8430    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
8431        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
8432            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
8433        }) {
8434            return Err(
8435                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
8436            );
8437        }
8438        let Some(ckpt) = sess.turn_ckpt.take() else {
8439            return Ok(None);
8440        };
8441        assert!(
8442            ckpt.pos <= sess.committed.len(),
8443            "checkpoint past committed ({} > {})",
8444            ckpt.pos,
8445            sess.committed.len()
8446        );
8447        // Restore through each layer's owning engine. A single primary-engine rollback is not
8448        // sufficient when the serving cache is stage-owned under cross-device PP.
8449        crate::pp::restore_cache_checkpoint(e, &self.cfg, None, &mut sess.cache, &ckpt.snap)?;
8450        debug_assert_eq!(
8451            sess.cache.pos, ckpt.pos,
8452            "rollback landed off the checkpoint"
8453        );
8454        sess.scratch.set_len(e, ckpt.pos)?;
8455        sess.committed.truncate(ckpt.pos);
8456        sess.last_h = Some(ckpt.last_h);
8457        sess.next_pred = None;
8458        sess.pending_tok = None;
8459        Ok(Some(ckpt.pos))
8460    }
8461
8462    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
8463    /// checkpoint without re-priming the checkpoint prefix.
8464    ///
8465    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
8466    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
8467    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
8468    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
8469    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
8470    ///
8471    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
8472    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
8473    pub fn spec_grow_and_rewind_to_checkpoint(
8474        &self,
8475        e: &Engine,
8476        sess: &mut SpecSession,
8477        target_cap: usize,
8478    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
8479        if target_cap <= sess.cache.max_ctx {
8480            return self.spec_rewind_to_checkpoint(e, sess);
8481        }
8482        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
8483            return Ok(None);
8484        };
8485        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
8486            return Err(format!(
8487                "checkpoint pos {} outside committed length {}",
8488                ckpt.pos,
8489                sess.committed.len(),
8490            )
8491            .into());
8492        }
8493        if ckpt.pos > target_cap {
8494            return Err(format!(
8495                "checkpoint pos {} exceeds grown capacity {target_cap}",
8496                ckpt.pos,
8497            )
8498            .into());
8499        }
8500
8501        let mut grown_cache = crate::pp::new_cache(e, &self.cfg, target_cap)?;
8502        let mut grown_scratch = MtpScratch::new(
8503            e,
8504            &self.cfg,
8505            target_cap,
8506            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8507        )?;
8508        crate::pp::restore_cache_checkpoint(
8509            e,
8510            &self.cfg,
8511            Some(&sess.cache),
8512            &mut grown_cache,
8513            &ckpt.snap,
8514        )?;
8515
8516        let src = &sess.scratch.kv;
8517        let dst = &mut grown_scratch.kv;
8518        if ckpt.pos > src.len
8519            || src.kv_dim_k != dst.kv_dim_k
8520            || src.kv_dim_v != dst.kv_dim_v
8521            || src.k_tok_bytes != dst.k_tok_bytes
8522            || src.v_tok_bytes != dst.v_tok_bytes
8523        {
8524            return Err(format!(
8525                "checkpoint draft layout mismatch (pos {}, source len {})",
8526                ckpt.pos, src.len,
8527            )
8528            .into());
8529        }
8530        let kb = ckpt.pos * src.k_tok_bytes;
8531        let vb = ckpt.pos * src.v_tok_bytes;
8532        if kb > 0 {
8533            e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
8534        }
8535        if vb > 0 {
8536            e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
8537        }
8538        grown_scratch.set_len(e, ckpt.pos)?;
8539        // The old scratch is dropped immediately after publication below. Bound its D2D reads
8540        // first; growth happens once per rewritten turn, outside the decode hot loop.
8541        e.stream().synchronize()?;
8542
8543        let ckpt = sess
8544            .turn_ckpt
8545            .take()
8546            .expect("checkpoint remained present through transactional grow");
8547        let pos = ckpt.pos;
8548        sess.cache = grown_cache;
8549        sess.scratch = grown_scratch;
8550        sess.committed.truncate(pos);
8551        sess.last_h = Some(ckpt.last_h);
8552        sess.next_pred = None;
8553        sess.pending_tok = None;
8554        sess.draft_ctx = None;
8555        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
8556        debug_assert_eq!(
8557            sess.scratch.kv.len, pos,
8558            "grown draft rewind landed off checkpoint"
8559        );
8560        Ok(Some(pos))
8561    }
8562
8563    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
8564    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
8565    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
8566    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
8567    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
8568    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
8569    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
8570    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
8571    /// park-time flush is a future request whose sampler is not knowable here (residual
8572    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
8573    pub fn spec_flush_pending(
8574        &self,
8575        e: &Engine,
8576        sess: &mut SpecSession,
8577        sampling: Option<SpecSampling>,
8578    ) -> Result<(), Box<dyn std::error::Error>> {
8579        let Some(b) = sess.pending_tok.take() else {
8580            return Ok(());
8581        };
8582        let mtp = self
8583            .mtp
8584            .as_ref()
8585            .expect("pending carry requires an MTP head");
8586        let n_embd = self.cfg.n_embd as usize;
8587        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8588        let embd_gpu = if spec_host_embd() {
8589            None
8590        } else {
8591            Some(
8592                self.embd_gpu
8593                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8594            )
8595        };
8596        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8597        let pos_b = sess.cache.pos;
8598        sess.scratch.set_len(e, pos_b)?;
8599        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
8600        sess.next_pred = Some(match sampling {
8601            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
8602                // window includes `b` itself: it is committed by this pass, and the pre-lane
8603                // code never counted a boundary token in the penalty history at all.
8604                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
8605                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
8606            }
8607            _ => argmax(&lg_b) as u32,
8608        });
8609        let anchor = sess
8610            .last_h
8611            .as_ref()
8612            .expect("pending carry requires last_h (the predecessor-row anchor)");
8613        self.mtp_kv_fill(e, mtp, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
8614        sess.last_h = Some(hb);
8615        sess.committed.push(b);
8616        Ok(())
8617    }
8618
8619    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
8620    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
8621    /// rounds through that same graph. Other model families keep their eager T=1 contract.
8622    fn spec_target_step_h(
8623        &self,
8624        e: &Engine,
8625        token: u32,
8626        cache: &mut Cache,
8627    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
8628        if self.cfg.step35.is_none() && !self.qwen35_serving_class() {
8629            return self.decode_step_h(e, token, cache);
8630        }
8631        let pos0 = cache.pos;
8632        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
8633        Ok((e.dtoh(&logits)?, hidden))
8634    }
8635
8636    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
8637    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
8638    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
8639    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
8640    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
8641    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
8642    /// dispatch sites cannot drift apart again.
8643    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
8644    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
8645    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
8646    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
8647    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
8648    /// eligibility sites so they cannot drift (the qwen35_serving_class lesson).
8649    fn mtp_graph_capturable(&self) -> bool {
8650        self.mtp
8651            .as_ref()
8652            .map(|m| match &m.ffn {
8653                crate::hybrid::Ffn::Dense { .. } => true,
8654                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
8655            })
8656            .unwrap_or(false)
8657    }
8658
8659    fn qwen35_serving_class(&self) -> bool {
8660        matches!(
8661            self.cfg.arch,
8662            memra_gguf::config::Arch::Qwen35 | memra_gguf::config::Arch::Qwen35Moe
8663        )
8664    }
8665
8666    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
8667    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
8668    /// session already exist.
8669    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
8670        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
8671            || !spec_devacc()
8672            || spec_replay_env_enabled()
8673            || spec_stream()
8674            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
8675            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
8676            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
8677            || std::env::var("MEMRA_SPEC_PMIN")
8678                .ok()
8679                .and_then(|v| v.parse::<f32>().ok())
8680                .unwrap_or(0.0)
8681                > 0.0
8682            || self.is_gemma4_e4b()
8683            || self.cfg.gemma4.is_some()
8684            || self.mtp.is_none()
8685        {
8686            return false;
8687        }
8688        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
8689            return false;
8690        };
8691        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
8692            return false;
8693        }
8694        crate::pp::PpNRt::get(e)
8695            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
8696            .unwrap_or(false)
8697    }
8698
8699    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
8700    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
8701    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
8702    #[allow(clippy::too_many_arguments)]
8703    pub fn generate_spec_session_pair(
8704        &self,
8705        e: &Engine,
8706        sess_a: &mut SpecSession,
8707        max_new_a: usize,
8708        k_a: usize,
8709        sess_b: &mut SpecSession,
8710        max_new_b: usize,
8711        k_b: usize,
8712    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
8713    {
8714        if !self.spec_pipe_available(e) {
8715            return Err("two-session speculative pipeline is outside its reduced matrix".into());
8716        }
8717        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
8718            return Err(
8719                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
8720            );
8721        }
8722        for sess in [&*sess_a, &*sess_b] {
8723            if sess.committed.is_empty()
8724                || sess.last_h.is_none()
8725                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
8726            {
8727                return Err("two-session speculative pipeline requires warm continuations".into());
8728            }
8729        }
8730
8731        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8732            && !spec_host_embd()
8733            && self.mtp_graph_capturable()
8734            && !crate::model::full_prec_enabled();
8735        let graph_a = graph_ok && k_a + 2 < 96;
8736        let graph_b = graph_ok && k_b + 2 < 96;
8737        let was_tracking = e.ctx().is_event_tracking();
8738        if (graph_a || graph_b) && was_tracking {
8739            unsafe {
8740                e.ctx().disable_event_tracking();
8741            }
8742        }
8743
8744        static LOGGED: std::sync::Once = std::sync::Once::new();
8745        LOGGED.call_once(|| {
8746            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
8747        });
8748        let sync = std::sync::Arc::new(SpecPipeSync::new());
8749        let lane_a = SpecPipeLane {
8750            sync: sync.clone(),
8751            lane: 0,
8752        };
8753        let lane_b = SpecPipeLane { sync, lane: 1 };
8754        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
8755        let (result_a, result_b) = std::thread::scope(|scope| {
8756            let b = scope.spawn(move || {
8757                let mut finish = SpecPipeFinish::new(&lane_b);
8758                let sess_b = unsafe { sess_b_ptr.get_mut() };
8759                let result = e
8760                    .ctx()
8761                    .bind_to_thread()
8762                    .map_err(|err| err.to_string())
8763                    .and_then(|_| {
8764                        self.generate_spec_inner2(
8765                            e,
8766                            &[],
8767                            max_new_b,
8768                            k_b,
8769                            graph_b,
8770                            Some(sess_b),
8771                            None,
8772                            None,
8773                            None,
8774                            None,
8775                            Some(&lane_b),
8776                        )
8777                        .map_err(|err| err.to_string())
8778                    });
8779                finish.close(result.is_err());
8780                result
8781            });
8782            let mut finish = SpecPipeFinish::new(&lane_a);
8783            let result_a = self.generate_spec_inner2(
8784                e,
8785                &[],
8786                max_new_a,
8787                k_a,
8788                graph_a,
8789                Some(sess_a),
8790                None,
8791                None,
8792                None,
8793                None,
8794                Some(&lane_a),
8795            );
8796            finish.close(result_a.is_err());
8797            let result_b = b
8798                .join()
8799                .map_err(|_| "paired speculative session B panicked".to_string())
8800                .and_then(|r| r);
8801            (result_a, result_b)
8802        });
8803
8804        if (graph_a || graph_b) && was_tracking {
8805            unsafe {
8806                e.ctx().enable_event_tracking();
8807            }
8808        }
8809        let result_a = result_a?;
8810        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
8811        Ok((result_a, result_b))
8812    }
8813
8814    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
8815    /// message rendered through the chat template continuation). Returns (new tokens emitted,
8816    /// drafted, accepted); session.committed grows by suffix + emitted.
8817    pub fn generate_spec_session(
8818        &self,
8819        e: &Engine,
8820        sess: &mut SpecSession,
8821        suffix: &[u32],
8822        max_new: usize,
8823        k: usize,
8824    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8825        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
8826    }
8827
8828    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
8829    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
8830    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
8831    /// for the filtered target (feat/filtered-spec).
8832    ///
8833    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
8834    /// output — once right after the prime's first token, then once per round commit — so a
8835    /// streaming caller can flush text at round cadence instead of once per burst. The slices
8836    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
8837    /// timing only: token bytes, session state, and exactness are untouched.
8838    ///
8839    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
8840    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
8841    /// the caller's scheduler regains control without waiting the burst out. Burst size is
8842    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
8843    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
8844    /// drains and the defensive tail flush can land with nothing new committed).
8845    #[allow(clippy::too_many_arguments)]
8846    pub fn generate_spec_session_sampled(
8847        &self,
8848        e: &Engine,
8849        sess: &mut SpecSession,
8850        suffix: &[u32],
8851        max_new: usize,
8852        k: usize,
8853        sampling: Option<SpecSampling>,
8854        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8855    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8856        self.generate_spec_session_sampled_prime_split(
8857            e, sess, suffix, max_new, k, sampling, None, on_commit,
8858        )
8859    }
8860
8861    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
8862    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
8863    /// pass `None` and stay on the existing zero-prime path.
8864    #[allow(clippy::too_many_arguments)]
8865    pub fn generate_spec_session_sampled_prime_split(
8866        &self,
8867        e: &Engine,
8868        sess: &mut SpecSession,
8869        suffix: &[u32],
8870        max_new: usize,
8871        k: usize,
8872        sampling: Option<SpecSampling>,
8873        prime_split: Option<usize>,
8874        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8875    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8876        self.generate_spec_session_constrained_prime_split(
8877            e,
8878            sess,
8879            suffix,
8880            max_new,
8881            k,
8882            sampling,
8883            None,
8884            prime_split,
8885            on_commit,
8886        )
8887    }
8888
8889    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
8890    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
8891    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
8892    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
8893    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
8894    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
8895    /// may drop (drafter is unconstrained); that is measured, not hidden.
8896    #[allow(clippy::too_many_arguments)]
8897    pub fn generate_spec_session_constrained(
8898        &self,
8899        e: &Engine,
8900        sess: &mut SpecSession,
8901        suffix: &[u32],
8902        max_new: usize,
8903        k: usize,
8904        sampling: Option<SpecSampling>,
8905        constraint: Option<&mut dyn SpecConstraint>,
8906        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8907    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8908        self.generate_spec_session_constrained_prime_split(
8909            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
8910        )
8911    }
8912
8913    #[allow(clippy::too_many_arguments)]
8914    pub fn generate_spec_session_constrained_prime_split(
8915        &self,
8916        e: &Engine,
8917        sess: &mut SpecSession,
8918        suffix: &[u32],
8919        max_new: usize,
8920        k: usize,
8921        sampling: Option<SpecSampling>,
8922        constraint: Option<&mut dyn SpecConstraint>,
8923        prime_split: Option<usize>,
8924        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
8925    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8926        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
8927            return Err(
8928                "constrained spec decode is greedy-only (worker routes sampled \
8929                        constrained to plain decode)"
8930                    .into(),
8931            );
8932        }
8933        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
8934        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
8935        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
8936        // serve continuation case — consume the carry in-loop with zero solo passes.
8937        if sess.pending_tok.is_some()
8938            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
8939        {
8940            self.spec_flush_pending(e, sess, sampling)?;
8941        }
8942
8943        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
8944        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
8945        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
8946        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8947            && !spec_host_embd()
8948            && self.mtp_graph_capturable()
8949            && k + 2 < 96
8950            && !crate::model::full_prec_enabled();
8951        let was_tracking = e.ctx().is_event_tracking();
8952        if graph_draft && was_tracking {
8953            unsafe {
8954                e.ctx().disable_event_tracking();
8955            }
8956        }
8957        let r = self.generate_spec_inner2(
8958            e,
8959            suffix,
8960            max_new,
8961            k,
8962            graph_draft,
8963            Some(sess),
8964            sampling,
8965            constraint,
8966            on_commit,
8967            prime_split,
8968            None,
8969        );
8970        if graph_draft && was_tracking {
8971            unsafe {
8972                e.ctx().enable_event_tracking();
8973            }
8974        }
8975        let (out, d, a) = r?;
8976        Ok((out, d, a))
8977    }
8978
8979    pub fn generate_spec(
8980        &self,
8981        e: &Engine,
8982        prompt: &[u32],
8983        max_new: usize,
8984        k: usize,
8985    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
8986        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
8987        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
8988        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
8989            && !spec_host_embd()
8990            && self.mtp_graph_capturable()
8991            && k + 2 < 96
8992            && !crate::model::full_prec_enabled();
8993        if !graph_draft {
8994            return self.generate_spec_inner2(
8995                e, prompt, max_new, k, false, None, None, None, None, None, None,
8996            );
8997        }
8998        let was_tracking = e.ctx().is_event_tracking();
8999        if was_tracking {
9000            unsafe {
9001                e.ctx().disable_event_tracking();
9002            }
9003        }
9004        let r = self.generate_spec_inner2(
9005            e, prompt, max_new, k, true, None, None, None, None, None, None,
9006        );
9007        if was_tracking {
9008            unsafe {
9009                e.ctx().enable_event_tracking();
9010            }
9011        }
9012        r
9013    }
9014
9015    fn generate_spec_inner2(
9016        &self,
9017        e: &Engine,
9018        prompt: &[u32],
9019        max_new: usize,
9020        k: usize,
9021        graph_draft: bool,
9022        mut sess: Option<&mut SpecSession>,
9023        sampling: Option<SpecSampling>,
9024        mut constraint: Option<&mut dyn SpecConstraint>,
9025        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9026        prime_split: Option<usize>,
9027        pipe: Option<&SpecPipeLane>,
9028    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9029        assert!(k >= 1, "k must be >= 1");
9030        if let Some(p) = pipe {
9031            p.setup_begin()?;
9032        }
9033        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9034        let mut flushed = 0usize;
9035        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9036        // at the next round boundary (same exit as max_new reached — the session tail runs).
9037        // Initialized by the unconditional post-prime flush below.
9038        let mut keep_going;
9039        let mtp = self
9040            .mtp
9041            .as_ref()
9042            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9043        let n_vocab = self.output.out_features();
9044        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9045        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9046        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9047        let d_vocab = mtp
9048            .shared_head_head
9049            .as_ref()
9050            .unwrap_or(&self.output)
9051            .out_features();
9052        let n_embd = self.cfg.n_embd as usize;
9053        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
9054        // already committed (their state is in the caches); 0 = fresh single-shot call.
9055        let session_mode = sess.is_some();
9056        let max_ctx = match sess.as_ref() {
9057            Some(s) => s.cache.max_ctx,
9058            None => prompt.len() + max_new + k + 8,
9059        };
9060        let mut own_cache;
9061        let mut own_scratch;
9062        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
9063        // (requested split, destination list). Single-shot per burst; fresh calls have none.
9064        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
9065        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
9066        // committed-length position; consumed one-shot like `capture_at`. None = legacy
9067        // prompt-end capture below.
9068        let mut ckpt_req: Option<usize> = None;
9069        let (
9070            cache,
9071            scratch,
9072            mut sess_tail,
9073            mut sess_draft_slot,
9074            mut sess_pending_slot,
9075            sess_ckpt_slot,
9076            sess_telem,
9077        ): (
9078            &mut Cache,
9079            &mut MtpScratch,
9080            Option<(
9081                &mut Vec<u32>,
9082                &mut Option<CudaSlice<f32>>,
9083                &mut Option<u32>,
9084                &mut u32,
9085                &mut u32,
9086            )>,
9087            Option<&mut Option<DraftGraphCtx>>,
9088            Option<&mut Option<u32>>,
9089            Option<&mut Option<SpecCheckpoint>>,
9090            Option<&SpecTelemetryCounters>,
9091        ) = match sess.take() {
9092            Some(sr) => {
9093                let SpecSession {
9094                    cache,
9095                    scratch,
9096                    committed,
9097                    last_h,
9098                    next_pred,
9099                    sctr: s_sctr,
9100                    uctr: s_uctr,
9101                    draft_ctx,
9102                    pending_tok,
9103                    turn_ckpt,
9104                    telem,
9105                    capture_at,
9106                    boundary_captures,
9107                    ckpt_at,
9108                } = sr;
9109                sess_capture = Some((capture_at.take(), boundary_captures));
9110                ckpt_req = ckpt_at.take();
9111                (
9112                    cache,
9113                    scratch,
9114                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
9115                    Some(draft_ctx),
9116                    Some(pending_tok),
9117                    Some(turn_ckpt),
9118                    Some(telem),
9119                )
9120            }
9121            None => {
9122                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
9123                // `Cache::new` verbatim.
9124                own_cache = crate::pp::new_cache(e, &self.cfg, max_ctx)?;
9125                // Persistent scratch = max_ctx rows (~2KB/token quantized).
9126                own_scratch = MtpScratch::new(
9127                    e,
9128                    &self.cfg,
9129                    max_ctx,
9130                    self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9131                )?;
9132                (
9133                    &mut own_cache,
9134                    &mut own_scratch,
9135                    None,
9136                    None,
9137                    None,
9138                    None,
9139                    None,
9140                )
9141            }
9142        };
9143        let base = cache.pos;
9144        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
9145        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
9146        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
9147        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
9148        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
9149        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
9150        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
9151        // acceptance-only — exactness is verify's job either way).
9152        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
9153        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
9154        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
9155        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
9156        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
9157        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
9158        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
9159        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
9160        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
9161        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
9162        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
9163        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
9164        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
9165        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
9166        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
9167        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
9168        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
9169        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
9170        // + fallback seam).
9171        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
9172        // bar — the retained verify-state commit proven equivalent to sequential serving —
9173        // was waiting on this arch running the serving batched verify class, which the
9174        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
9175        // replay-free commit consumes is now produced by the SAME serving-class verify that
9176        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
9177        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
9178        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
9179        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
9180        // rollback + A/B seam.
9181        let spec_replay = spec_replay_env_enabled();
9182        if constraint.is_some() && spec_replay {
9183            return Err(
9184                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
9185                        (legacy replay commits an unmasked bonus)"
9186                    .into(),
9187            );
9188        }
9189        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
9190        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
9191        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
9192        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
9193
9194        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
9195        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
9196        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
9197        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
9198        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
9199        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
9200        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
9201        // generation exactly where the last turn stopped — no prime at all. The stashed
9202        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
9203        // committed.last() by the same rule this entry applies to a cold prime's last row —
9204        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
9205        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
9206        // where the sampler and the session's Philox counters were live). `last_h` seeds the
9207        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
9208        let continuation = prompt.is_empty();
9209        if continuation {
9210            assert!(session_mode, "empty prompt requires a session");
9211            assert!(
9212                sess_tail
9213                    .as_ref()
9214                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
9215                        && lh.is_some()
9216                        && (np.is_some() || carried_pending.is_some())),
9217                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
9218            );
9219        }
9220        let mut prime_logits;
9221        let mut prompt_h: Option<CudaSlice<f32>> = None;
9222        let t_prime = std::time::Instant::now();
9223        let batched_prime = !continuation
9224            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
9225            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9226            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
9227        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
9228        if prime_split.is_some() && continuation {
9229            return Err("spec prime split requires a non-empty prime".into());
9230        }
9231        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
9232        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
9233        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
9234        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
9235        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
9236        // cannot honor (outside this prime's range) silently drops the capture — the
9237        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
9238        let ckpt_rel = if continuation {
9239            None
9240        } else {
9241            ckpt_req
9242                .and_then(|abs| abs.checked_sub(base))
9243                .filter(|&r| r > 0 && r < prompt.len())
9244        };
9245        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
9246        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
9247        // the legacy single-split program, byte-for-byte.
9248        let mut stops: Vec<usize> = Vec::new();
9249        for b in [prime_split, ckpt_rel].into_iter().flatten() {
9250            if !stops.contains(&b) {
9251                stops.push(b);
9252            }
9253        }
9254        stops.sort_unstable();
9255        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
9256        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
9257        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
9258        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
9259        if continuation {
9260            prime_logits = Vec::new();
9261        } else if !stops.is_empty() {
9262            if let Some(&first) = stops.first() {
9263                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
9264                    return Err(format!(
9265                        "spec prime split {first} is below PRIME_MIN_T {}",
9266                        crate::hybrid_forward::PRIME_MIN_T,
9267                    )
9268                    .into());
9269                }
9270            }
9271            // Mirror the plain worker's boundary stops exactly. Each segment is a
9272            // request-level prime (`queued_after` keeps Step35 arm selection independent of
9273            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
9274            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
9275            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
9276            // coherent prompt.
9277            let mut h_all = e.uninit(prompt.len() * n_embd)?;
9278            prime_logits = Vec::new();
9279            let mut prev = 0usize;
9280            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
9281                if seg_end <= prev {
9282                    continue;
9283                }
9284                let seg = &prompt[prev..seg_end];
9285                let is_final = seg_end == prompt.len();
9286                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
9287                    && (!is_final
9288                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9289                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
9290                if batched_seg {
9291                    let (l, _, h_seg) =
9292                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
9293                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
9294                    prime_logits = l;
9295                } else {
9296                    for (i, &tok) in seg.iter().enumerate() {
9297                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
9298                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
9299                        prime_logits = l;
9300                    }
9301                }
9302                prev = seg_end;
9303                if is_final {
9304                    break;
9305                }
9306                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
9307                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
9308                // states are about to be advanced in place by the next segment, so this is
9309                // the ONLY moment the boundary's recurrent state exists. Capture iff the
9310                // worker requested exactly this stop (cold sessions only — `capture_at` is
9311                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
9312                // publication is an optimization, never a correctness dependency.
9313                if base == 0 {
9314                    if let Some((requested, slot)) = sess_capture.as_mut() {
9315                        // Publish at the requested miss-LCP stop (the shared-prefix class)
9316                        // AND at the stable-boundary stop (the next-turn re-render class,
9317                        // lane/frspec-multiturn-cache) — the same boundary set the plain
9318                        // prefill tick learns. Without the second entry, the turn after a
9319                        // cold re-park could only hit the OLDER lcp entry (the measured
9320                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
9321                        // rewound to 15222). Dedupe is the worker sweep's has_key.
9322                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
9323                            if let Ok(snap) = cache.snapshot(e) {
9324                                slot.push(SpecBoundaryCapture {
9325                                    snap,
9326                                    pos: seg_end,
9327                                    logits: prime_logits.clone(),
9328                                    // rows [0..seg_end) of h_all are primed — the following
9329                                    // segments append, never overwrite.
9330                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
9331                                });
9332                            }
9333                        }
9334                    }
9335                }
9336                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
9337                // same snapshot mechanics, installed post-prime in place of the prompt-end
9338                // capture the re-render class always diverged below.
9339                if ckpt_rel == Some(seg_end) {
9340                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9341                        e.uninit(n_embd).and_then(|mut a| {
9342                            e.copy_view_into(
9343                                &mut a,
9344                                0,
9345                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9346                                n_embd,
9347                            )?;
9348                            Ok(a)
9349                        });
9350                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
9351                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
9352                            snap,
9353                            pos: base + seg_end,
9354                            last_h,
9355                        }),
9356                        _ => None,
9357                    });
9358                }
9359            }
9360            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
9361                eprintln!(
9362                    "[spec-prime] stops={stops:?} tail={}",
9363                    prompt.len() - stops.last().copied().unwrap_or(0)
9364                );
9365            }
9366            prompt_h = Some(h_all);
9367        } else if batched_prime {
9368            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
9369            prime_logits = l;
9370            prompt_h = Some(hiddens);
9371        } else {
9372            prime_logits = Vec::new();
9373            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
9374            for (i, &tok) in prompt.iter().enumerate() {
9375                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
9376                if let Some(ph) = prompt_h.as_mut() {
9377                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
9378                }
9379                prime_logits = l;
9380            }
9381        }
9382        e.stream().synchronize()?;
9383        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
9384        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
9385        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
9386        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
9387        // prime_split. The mid-prompt capture above already consumed the request if it matched.
9388        if !continuation && base == 0 {
9389            if let Some((requested, slot)) = sess_capture.as_mut() {
9390                if *requested == Some(prompt.len()) && slot.is_empty() {
9391                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
9392                    if let Ok(snap) = cache.snapshot(e) {
9393                        slot.push(SpecBoundaryCapture {
9394                            snap,
9395                            pos: prompt.len(),
9396                            logits: prime_logits.clone(),
9397                            last_h: prompt_h
9398                                .as_ref()
9399                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
9400                                .unwrap_or_default(),
9401                        });
9402                    }
9403                }
9404            }
9405        }
9406        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
9407        // prime-subtraction hack.
9408        crate::PRIME_NANOS.store(
9409            t_prime.elapsed().as_nanos() as u64,
9410            std::sync::atomic::Ordering::Relaxed,
9411        );
9412
9413        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9414        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
9415        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
9416        let host_embd = spec_host_embd();
9417        let embd_gpu = if host_embd {
9418            None
9419        } else {
9420            Some(
9421                self.embd_gpu
9422                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9423            )
9424        };
9425        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9426        if host_embd {
9427            eprintln!(
9428                "[spec] host-row embedding: {} bytes kept off HBM",
9429                self.embd.raw.len()
9430            );
9431        }
9432        let mut out: Vec<u32> = Vec::with_capacity(max_new);
9433        let mut total_drafted = 0usize;
9434        let mut total_accepted = 0usize;
9435
9436        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
9437        // The sampler config, the session's Philox counters and the penalty window are parsed
9438        // HERE, above the boundary-token selection, because the boundary token must be drawn
9439        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
9440        // selection, which is the whole mechanical reason the boundary token was an argmax:
9441        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
9442        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
9443        // below takes the argmax path it always took).
9444        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
9445        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
9446        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
9447        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
9448        let sp = sampling.unwrap_or_else(|| SpecSampling {
9449            temp: std::env::var("MEMRA_SPEC_TEMP")
9450                .ok()
9451                .and_then(|v| v.parse().ok())
9452                .unwrap_or(0.0),
9453            seed: std::env::var("MEMRA_SEED")
9454                .ok()
9455                .and_then(|v| v.parse().ok())
9456                .unwrap_or(42),
9457            top_k: std::env::var("MEMRA_TOP_K")
9458                .ok()
9459                .and_then(|v| v.parse().ok())
9460                .unwrap_or(0),
9461            top_p: std::env::var("MEMRA_TOP_P")
9462                .ok()
9463                .and_then(|v| v.parse().ok())
9464                .unwrap_or(1.0),
9465            min_p: std::env::var("MEMRA_MIN_P")
9466                .ok()
9467                .and_then(|v| v.parse().ok())
9468                .unwrap_or(0.0),
9469            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
9470                .ok()
9471                .and_then(|v| v.parse().ok())
9472                .unwrap_or(0),
9473            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
9474                .ok()
9475                .and_then(|v| v.parse().ok())
9476                .unwrap_or(1.0),
9477            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
9478                .ok()
9479                .and_then(|v| v.parse().ok())
9480                .unwrap_or(0.0),
9481            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
9482                .ok()
9483                .and_then(|v| v.parse().ok())
9484                .unwrap_or(0.0),
9485        });
9486        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
9487        let sampled = sp_temp > 0.0;
9488        // Counters resume from the session (burst continuity: randomness must never repeat
9489        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
9490        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
9491        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
9492        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
9493        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
9494        // for the penalized+filtered target). History = generated tokens, host-tracked window.
9495        let pen_on = sampled
9496            && sp.penalty_last_n > 0
9497            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
9498        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
9499        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
9500        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
9501        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
9502        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
9503        // which is what the API contract says and what the plain sampler's own `history` does.
9504        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
9505        let mut pen_hist: Vec<u32> = if pen_on {
9506            let sess_hist: &[u32] = if spec_pen_session_on() {
9507                sess_tail
9508                    .as_ref()
9509                    .map(|(c, ..)| c.as_slice())
9510                    .unwrap_or(&[])
9511            } else {
9512                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
9513            };
9514            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
9515        } else {
9516            Vec::new()
9517        };
9518        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
9519        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
9520        // request's own filtered/penalized target through the session's Philox stream
9521        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
9522        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
9523        // Emit it, then FEED it to establish the loop invariant below.
9524        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
9525        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
9526        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
9527        // prompt's last logits (plain constrained-greedy identity); a continuation without
9528        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
9529        // worker never resumes constrained sessions from the pool, so this cannot fire).
9530        if let Some(c) = constraint.as_deref_mut() {
9531            if continuation && carried_pending.is_none() {
9532                return Err("constrained spec continuation requires a carried pending \
9533                            (pool resume is unconstrained-only)"
9534                    .into());
9535            }
9536            if !continuation {
9537                c.mask_logits(&mut prime_logits)
9538                    .map_err(|e2| format!("constraint: {e2}"))?;
9539            }
9540        }
9541        let mut last_token = if let Some(b) = carried_pending {
9542            b
9543        } else if continuation {
9544            // A continuation's boundary token was DRAWN by the burst that stashed it (the
9545            // session tail below), or by `spec_session_from_restored` for a converted
9546            // prefix-cache hit — in both cases from the correct logits row with this same
9547            // session's Philox stream, which is why it can be consumed here as-is.
9548            sess_tail.as_ref().unwrap().2.unwrap()
9549        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
9550            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
9551        } else {
9552            // greedy (byte contract), the rollback door, or constrained (masked-argmax
9553            // identity — the worker routes sampled+constrained to the plain path, and this
9554            // function refuses the combination outright above).
9555            argmax(&prime_logits) as u32
9556        };
9557        if pen_on {
9558            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
9559            // emitted token into its penalty history, and pre-lane the burst's first token
9560            // was invisible to penalties forever (never pushed, and never in `committed`
9561            // until this burst's tail). Covers the carry/continuation seeds too — neither is
9562            // in `committed` yet.
9563            pen_hist.push(last_token);
9564        }
9565        if carried_pending.is_none() {
9566            out.push(last_token);
9567            // grammar advances with every emitted token (carried pendings were consumed
9568            // by the burst that emitted them).
9569            if let Some(c) = constraint.as_deref_mut() {
9570                c.consume(last_token)
9571                    .map_err(|e2| format!("constraint: {e2}"))?;
9572            }
9573        }
9574        if continuation {
9575            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
9576            // overhang so the chain's first append lands at slot base (== committed.len()).
9577            scratch.set_len(e, base)?;
9578        }
9579        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
9580        // concatenating to the full `out`). Called after the prime's first token and after each
9581        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
9582        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
9583        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
9584        fn flush_commit(
9585            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
9586            out: &[u32],
9587            flushed: &mut usize,
9588        ) -> bool {
9589            if let Some(f) = cb.as_mut() {
9590                let keep = f(&out[*flushed..]);
9591                *flushed = out.len();
9592                keep
9593            } else {
9594                true
9595            }
9596        }
9597        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
9598        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
9599        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
9600        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
9601        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
9602        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
9603        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
9604        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
9605        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
9606        // those, so their residual mass is p(x), correct by construction).
9607        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
9608            match &mtp.d2t {
9609                Some(map) => Some(e.htod_u32_v(map)?),
9610                None => None,
9611            }
9612        } else {
9613            None
9614        };
9615        let mut q_full_buf: Option<CudaSlice<f32>> = None;
9616        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
9617        // dspark sampled-admission walk); byte-identical to the closure it replaces.
9618        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
9619        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
9620        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
9621        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
9622        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
9623        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
9624        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
9625        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
9626        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
9627        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
9628        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
9629        let t_ent = std::time::Instant::now();
9630
9631        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
9632        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
9633        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
9634        // the one that matters (a history-rewriting client mutates what the session GENERATED,
9635        // so the next turn's prompt agrees with this one up to exactly here).
9636        //
9637        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
9638        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
9639        // hold exactly `base + prompt.len()` rows and nothing generated.
9640        //
9641        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
9642        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
9643        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
9644        // `<think>` block the client strips, so every later turn's diff diverged exactly one
9645        // token below the checkpoint and affinity declined 100% of the time. Measured on the
9646        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
9647        // whole mechanism inert while looking, from the outside, like a working
9648        // correctness-declines-safely path — hence the decline log carries the offsets.
9649        //
9650        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
9651        // state (the reason a spec session could not rewind before). The draft scratch needs no
9652        // copy: rows below the boundary are rewritten by the next turn's own fill.
9653        //
9654        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
9655        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
9656        // checkpoint rather than replacing it with a strictly worse one.
9657        //
9658        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
9659        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
9660        // fail the burst that is already running — so the error is swallowed, loud only under
9661        // MEMRA_DEBUG_SPEC.
9662        //
9663        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
9664        // posture above was DISPROVED for the think-posture template class — the prompt's own
9665        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
9666        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
9667        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
9668        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
9669        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
9670        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
9671        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
9672        if let Some(slot) = sess_ckpt_slot {
9673            if let Some(early) = ckpt_early {
9674                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
9675                    eprintln!(
9676                        "[spec] stable-boundary turn checkpoint skipped; \
9677                               next turn re-primes in full"
9678                    );
9679                }
9680                *slot = early;
9681            } else if !continuation {
9682                let pos = cache.pos;
9683                debug_assert_eq!(
9684                    pos,
9685                    base + prompt.len(),
9686                    "turn checkpoint must sit at the prompt end, before the init feed"
9687                );
9688                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9689                    if let Some(ph) = &prompt_h {
9690                        // hidden of the LAST primed row = the predecessor anchor at this
9691                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
9692                        // last_h, and what the next prime's fill reads for its first row).
9693                        let np = prompt.len();
9694                        e.uninit(n_embd).and_then(|mut a| {
9695                            e.copy_view_into(
9696                                &mut a,
9697                                0,
9698                                &ph.slice((np - 1) * n_embd..np * n_embd),
9699                                n_embd,
9700                            )?;
9701                            Ok(a)
9702                        })
9703                    } else {
9704                        Err("no prompt hiddens".into())
9705                    };
9706                match (cache.snapshot(e), anchor) {
9707                    (Ok(snap), Ok(last_h)) => {
9708                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
9709                    }
9710                    (s, a) => {
9711                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
9712                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
9713                            let err = s
9714                                .err()
9715                                .map(|e| e.to_string())
9716                                .or_else(|| a.err().map(|e| e.to_string()))
9717                                .unwrap_or_default();
9718                            eprintln!(
9719                                "[spec] turn checkpoint skipped ({err}); \
9720                                       next turn re-primes in full"
9721                            );
9722                        }
9723                    }
9724                }
9725            }
9726        }
9727        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
9728        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
9729        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
9730        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
9731        let mut last_pred = 0u32;
9732        let mut last_col_logits: Option<CudaSlice<f32>> = None;
9733        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
9734        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
9735        let mut init_logits_host: Option<Vec<f32>> = None;
9736        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
9737            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
9738            last_pred = argmax(&init_logits) as u32;
9739            if constraint.is_some() {
9740                init_logits_host = Some(init_logits.clone());
9741            }
9742            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
9743            if sampled {
9744                last_col_logits = Some(e.htod(&init_logits)?);
9745            }
9746            h
9747        } else {
9748            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
9749            let lh = sess_tail
9750                .as_ref()
9751                .unwrap()
9752                .1
9753                .as_ref()
9754                .expect("pending carry requires last_h");
9755            e.clone_dtod(lh)?
9756        };
9757        let t_init = t_ent.elapsed();
9758        let mut last_col_stats: Option<(f32, f32, f32)> = None;
9759        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
9760        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
9761        // stable pointer for the graph-draft round-start copy.
9762        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
9763        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
9764        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
9765        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
9766        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
9767        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
9768        // overwritten below).
9769        let mut fill_prev = e.clone_dtod(&h_seed0)?;
9770        {
9771            if let Some(ph) = &prompt_h {
9772                let np = prompt.len();
9773                e.copy_view_into(
9774                    &mut h_seed_buf,
9775                    0,
9776                    &ph.slice((np - 1) * n_embd..np * n_embd),
9777                    n_embd,
9778                )?;
9779            } else if continuation {
9780                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
9781                    if let Some(lh) = lh.as_ref() {
9782                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
9783                    }
9784                }
9785            }
9786        }
9787        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
9788        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
9789
9790        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
9791        let fork_mode = OptiForkGateMode::configured();
9792        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
9793        // the end. Metric normalization vs the reference engine: BOTH engines count
9794        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
9795        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
9796        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
9797        let mut st_drafted = vec![0usize; k];
9798        let mut st_accepted = vec![0usize; k];
9799        let mut st_len_hist = vec![0usize; k + 1];
9800        let mut st_full = 0usize;
9801        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
9802        // stop the draft chain early when the head's softmax confidence in its own pick drops
9803        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
9804        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
9805        let p_min = *PMIN.get_or_init(|| {
9806            std::env::var("MEMRA_SPEC_PMIN")
9807                .ok()
9808                .and_then(|v| v.parse().ok())
9809                .unwrap_or(0.0)
9810        });
9811        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
9812        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
9813        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
9814        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
9815        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
9816        // verify batch is not); the j==0 exemption stays for pending-less rounds.
9817        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
9818            .map(|v| v == "1")
9819            .unwrap_or(false);
9820
9821        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
9822        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
9823        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
9824        // cuBLAS path in an exotic head) falls back to the eager draft chain.
9825        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
9826        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
9827        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
9828        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
9829        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
9830        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
9831        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
9832        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
9833        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
9834            Some(c) => c,
9835            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
9836        };
9837        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
9838        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
9839        if sampled && dctx.g_q.len() < d_vocab {
9840            dctx.g_q = e.zeros(d_vocab)?;
9841            dctx.g_perturb = e.zeros(d_vocab)?;
9842        }
9843        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
9844        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
9845        // truncation (the correctness backstop) stops cutting every tight-schema round.
9846        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
9847        // shape, so a parked graph of the other shape is dropped and recaptured.
9848        let dmask_on = constraint
9849            .as_deref()
9850            .is_some_and(|c| c.draft_mask_enabled());
9851        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
9852        if dmask_on && dctx.g_dmask.len() < dmask_words {
9853            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
9854            dctx.graph = None; // the old capture baked the old (or no) mask pointer
9855            dctx.failed.clear_greedy();
9856            dctx.keeper.clear();
9857        }
9858        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
9859            dctx.graph = None;
9860            dctx.failed.clear_greedy();
9861            dctx.keeper.clear();
9862        }
9863        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
9864            let DraftGraphCtx {
9865                g_tok,
9866                g_pos,
9867                g_seed,
9868                g_p,
9869                g_dmask,
9870                ..
9871            } = &mut dctx;
9872            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
9873            // host uploads the position's real words, so the warmups stay grammar-free.
9874            if dmask_on {
9875                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
9876            }
9877            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
9878            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
9879            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
9880            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
9881            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
9882            // passes (and, in serve, other sessions) recycle those addresses and the replay then
9883            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
9884            let cap_res = e.capture_graph_retained(|e| {
9885                self.mtp_head_forward_cap(
9886                    e,
9887                    mtp,
9888                    g_tok,
9889                    g_pos,
9890                    g_seed,
9891                    g_p,
9892                    &mut *scratch,
9893                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
9894                    true,
9895                    embd_gpu.expect("graph draft requires resident embedding"),
9896                    embd_qt,
9897                    embd_rb,
9898                    d_vocab,
9899                    None,
9900                    None,
9901                    if dmask_on {
9902                        Some((g_dmask_ro, dmask_words))
9903                    } else {
9904                        None
9905                    },
9906                )
9907            });
9908            match cap_res {
9909                Ok((g, keep)) => {
9910                    scratch.set_len(e, base)?;
9911                    dctx.graph = Some(g);
9912                    dctx.graph_masked = dmask_on;
9913                    dctx.keeper = keep;
9914                }
9915                Err(err) => {
9916                    scratch.set_len(e, base)?;
9917                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
9918                    // silent. Once per flip — mark returns None on an already-failed ctx.
9919                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
9920                        eprintln!("{line}");
9921                    }
9922                }
9923            }
9924        }
9925        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
9926        // graph object, built only when sampled && graph-eligible — the greedy capture above is
9927        // untouched (and skipped when sampled: its graph would never be launched). Same head
9928        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
9929        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
9930        // once per round); the raw head logits land in the persistent g_q for the host's
9931        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
9932        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
9933        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
9934        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
9935        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
9936        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
9937        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
9938        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
9939        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
9940        // this compare misses at most ONCE per resumed request — the first burst recaptures
9941        // and every later burst in that request replays. A client that wants the parked graph
9942        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
9943        // stable across its whole conversation.
9944        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
9945        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
9946        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
9947        // force the eager draft (which computes stats/penalties per row).
9948        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
9949        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
9950        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
9951        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
9952        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
9953        // the request shape the vendor-default flip makes the majority).
9954        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
9955        let pure_temp = s_key.pure_temp();
9956        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
9957            dctx.graph_s = None;
9958            dctx.failed.clear_sampled();
9959            dctx.s_key = None;
9960            dctx.q_slots.clear();
9961            dctx.keeper_s.clear();
9962        }
9963        if graph_draft
9964            && sampled
9965            && pure_temp
9966            && dctx.graph_s.is_none()
9967            && !dctx.failed.sampled_failed()
9968        {
9969            let DraftGraphCtx {
9970                g_tok,
9971                g_pos,
9972                g_seed,
9973                g_p,
9974                g_ctr,
9975                g_perturb,
9976                g_q,
9977                ..
9978            } = &mut dctx;
9979            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
9980            let cap_res = e.capture_graph_retained(|e| {
9981                self.mtp_head_forward_cap(
9982                    e,
9983                    mtp,
9984                    g_tok,
9985                    g_pos,
9986                    g_seed,
9987                    g_p,
9988                    &mut *scratch,
9989                    p_min > 0.0,
9990                    true,
9991                    embd_gpu.expect("graph draft requires resident embedding"),
9992                    embd_qt,
9993                    embd_rb,
9994                    d_vocab,
9995                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
9996                    None,
9997                    None, // constrained spec is greedy-only — sampled never carries a hook
9998                )
9999            });
10000            match cap_res {
10001                Ok((g, keep)) => {
10002                    scratch.set_len(e, base)?;
10003                    for _ in 0..k {
10004                        dctx.q_slots.push(e.zeros(d_vocab)?);
10005                    }
10006                    dctx.graph_s = Some(g);
10007                    dctx.s_key = Some(s_key);
10008                    dctx.keeper_s = keep;
10009                }
10010                Err(err) => {
10011                    scratch.set_len(e, base)?;
10012                    // LOUD flip (audit Q2): same contract as the greedy capture above.
10013                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
10014                        eprintln!("{line}");
10015                    }
10016                }
10017            }
10018        }
10019        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
10020        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
10021        // captured under this request's exact regime, and capture requires `pure_temp` — so a
10022        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
10023        // the graph arm, so it is asserted here rather than assumed: a future change that widens
10024        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
10025        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
10026        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
10027        // rather than launching it; the launch site re-tests `pure_temp` independently.
10028        if sampled && !pure_temp && dctx.graph_s.is_some() {
10029            debug_assert!(
10030                false,
10031                "sampled draft graph parked under {:?} survived into a FILTERED request \
10032                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
10033                 softmax, so the verify's filtered q would test a distribution the draft was \
10034                 never sampled from",
10035                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10036            );
10037            eprintln!(
10038                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
10039                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
10040                 EAGER — the key must carry every field that shapes q",
10041                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10042            );
10043            dctx.graph_s = None;
10044            dctx.s_key = None;
10045            dctx.q_slots.clear();
10046            dctx.keeper_s.clear();
10047        }
10048        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
10049        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
10050        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
10051        // arms below print which chain actually ran, so the probe never restates the condition.
10052        if skey_probe() {
10053            eprintln!(
10054                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
10055                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
10056                sampled as u8,
10057                pure_temp as u8,
10058                sp_temp,
10059                sp.top_k,
10060                sp.top_p,
10061                sp.min_p,
10062                pen_on as u8,
10063                k,
10064                graph_draft as u8,
10065                dctx.graph_s.is_some() as u8,
10066                dctx.s_key,
10067            );
10068        }
10069        let t_cap = t_ent.elapsed();
10070        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
10071        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
10072        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
10073        // fill: the first chain step processes it and appends its entry at slot prompt.len().
10074        if let Some(ph) = &prompt_h {
10075            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
10076            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
10077            // global positions [base..base+tp). Fresh call: base==0, identical to before.
10078            scratch.set_len(e, base)?;
10079            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
10080            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
10081            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
10082            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
10083            let tp = prompt.len();
10084            let fill_chunk: usize = if crate::cache::swa_ring_on() {
10085                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
10086            } else {
10087                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
10088                // meaning one monolithic fill.
10089                std::env::var("MEMRA_PRIME_CHUNK")
10090                    .ok()
10091                    .and_then(|v| v.parse().ok())
10092                    .unwrap_or(4096)
10093            };
10094            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
10095            let mut start = 0usize;
10096            while start < tp {
10097                let end = (start + fill_chunk).min(tp);
10098                let tc = end - start;
10099                {
10100                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
10101                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
10102                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
10103                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
10104                    let mut phs = e.zeros(tc * n_embd)?;
10105                    let (src_lo, dst_off) = if start == 0 {
10106                        (0, n_embd)
10107                    } else {
10108                        ((start - 1) * n_embd, 0)
10109                    };
10110                    let n_copy = if start == 0 {
10111                        (tc - 1) * n_embd
10112                    } else {
10113                        tc * n_embd
10114                    };
10115                    if start == 0 {
10116                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10117                            if let Some(lh) = lh.as_ref() {
10118                                e.copy_into(&mut phs, 0, lh, n_embd)?;
10119                            }
10120                        }
10121                    }
10122                    if n_copy > 0 {
10123                        e.copy_view_into(
10124                            &mut phs,
10125                            dst_off,
10126                            &ph.slice(src_lo..src_lo + n_copy),
10127                            n_copy,
10128                        )?;
10129                    }
10130                    self.mtp_kv_fill(
10131                        e,
10132                        mtp,
10133                        &prompt[start..end],
10134                        &phs,
10135                        base + start,
10136                        &mut *scratch,
10137                        embd_dev,
10138                    )?;
10139                }
10140                start = end;
10141            }
10142        }
10143        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
10144        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
10145        // (=1 brackets the whole call in run_spec.rs, prime included.)
10146        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
10147            unsafe extern "C" {
10148                fn cudaProfilerStart() -> i32;
10149            }
10150            unsafe {
10151                cudaProfilerStart();
10152            }
10153        }
10154        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
10155        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
10156        // consume each other's device outputs; the host drains the ring every M rounds. v1
10157        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
10158        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
10159        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
10160        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
10161        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
10162        let stream_on = crate::spec::spec_stream()
10163            && !sampled
10164            && !spec_replay
10165            && constraint.is_none()
10166            && !session_mode
10167            && embd_gpu.is_some()
10168            && !crate::model::full_prec_enabled()
10169            && k + 2 < 96;
10170        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
10171        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
10172        if stream_on {
10173            let cap = e.capture_graph(|e| {
10174                for j in 0..k.max(1) {
10175                    self.mtp_head_forward_cap(
10176                        e,
10177                        mtp,
10178                        &mut dctx.g_tok,
10179                        &mut dctx.g_pos,
10180                        &mut dctx.g_seed,
10181                        &mut dctx.g_p,
10182                        &mut *scratch,
10183                        true,
10184                        true,
10185                        embd_gpu.expect("round stream requires resident embedding"),
10186                        embd_qt,
10187                        embd_rb,
10188                        d_vocab,
10189                        None,
10190                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
10191                        None, // round-stream requires constraint.is_none() (see stream_on)
10192                    )?;
10193                }
10194                Ok(())
10195            });
10196            match cap {
10197                Ok(g) => {
10198                    scratch.set_len(e, 0)?;
10199                    stream_graph = Some(g);
10200                }
10201                Err(err) => {
10202                    scratch.set_len(e, 0)?;
10203                    if debug_spec {
10204                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
10205                    }
10206                }
10207            }
10208        }
10209        let stream_active = stream_on && stream_graph.is_some();
10210        if debug_spec {
10211            eprintln!(
10212                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
10213                crate::spec::spec_stream(),
10214                dctx.graph.is_some(),
10215                stream_graph.is_some()
10216            );
10217        }
10218        let t_v_s = k + 1;
10219        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
10220        // module (extracted 2026-07-12; the gemma burst reuses them).
10221        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
10222        let crate::round_stream::StreamBufs {
10223            mut vtok_d,
10224            mut brk_d,
10225            mut pend_d,
10226            last_pred_d,
10227            mut pos_ctr,
10228            mut pos_start_d,
10229            mut ring_d,
10230            acc_d: mut stream_acc,
10231            m_rounds,
10232            k: _,
10233        } = sb;
10234        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
10235            Some(crate::round_stream::kv_len_ptr_table(
10236                e,
10237                cache,
10238                Some(&pos_ctr),
10239            )?)
10240        } else {
10241            None
10242        };
10243
10244        let t_fill = t_ent.elapsed();
10245        let mut round = 0usize;
10246        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
10247        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
10248        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
10249        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
10250        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
10251        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
10252        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
10253        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
10254        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
10255        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
10256        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
10257        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
10258        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
10259        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
10260        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
10261        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
10262        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
10263        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
10264        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
10265        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
10266        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
10267        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
10268        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
10269        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
10270        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
10271        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
10272        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
10273        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
10274        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
10275        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
10276            .ok()
10277            .and_then(|v| v.parse().ok());
10278        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
10279            4
10280        } else if self.cfg.n_embd as usize >= 2500 {
10281            2
10282        } else {
10283            1
10284        };
10285        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
10286        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
10287        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
10288        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
10289        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
10290            .ok()
10291            .and_then(|v| v.parse().ok())
10292            .unwrap_or(1024);
10293        let floor_at = |pos: usize| -> usize {
10294            if adapt_floor_env.is_some() || pos < floor_ctx {
10295                adapt_floor
10296            } else if adapt_floor >= 4 {
10297                1
10298            } else {
10299                adapt_floor
10300            }
10301        };
10302        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
10303        // fixed-K default path is untouched by this whole block.
10304        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
10305            .ok()
10306            .and_then(|v| v.parse().ok())
10307            .unwrap_or(7);
10308        let k_cap = k.min(cap_max).max(1);
10309        let mut kc = k_cap;
10310        let mut opti_fork: Option<OptiForkState> = None;
10311        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
10312        if fork_mode != OptiForkGateMode::Disabled {
10313            let fence = crate::pp::pp_cuts(self.layers.len());
10314            let refusal = if !session_mode {
10315                Some("not-session")
10316            } else if k != 1 || adapt {
10317                Some("requires-fixed-k1")
10318            } else if sampled || constraint.is_some() || spec_replay {
10319                Some("sampled-constrained-or-replay")
10320            } else if pipe.is_some() {
10321                Some("two-session-pipeline")
10322            } else if !spec_devacc() {
10323                Some("requires-device-accept")
10324            } else if stream_active || crate::spec::spec_stream() {
10325                Some("round-stream")
10326            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
10327                Some("swa-ring")
10328            } else if crate::pp::pp_host_bounce_active() {
10329                Some("host-bounce")
10330            } else if fork_mode == OptiForkGateMode::Controller
10331                && cache.recur.iter().any(Option::is_some)
10332            {
10333                Some("controller-requires-zero-recurrent-state")
10334            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
10335                Some("requires-pp2")
10336            } else {
10337                None
10338            };
10339            if let Some(reason) = refusal {
10340                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10341                eprintln!("[opti-fork] refused reason={reason}");
10342            } else {
10343                let fence = fence.expect("validated PP-2 fence");
10344                let rt = crate::pp::PpNRt::get(e)?;
10345                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
10346                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
10347                let primary_supported =
10348                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
10349                if !rt.cross_device() || !primary_supported {
10350                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10351                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
10352                } else {
10353                    // Both recurrent snapshots and both seed generations are allocated before
10354                    // the first fork, each through its owning PP stage. Allocation failure
10355                    // therefore happens before any optimistic state mutation can occur.
10356                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10357                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10358                    let fork = OptiForkState::new(
10359                        e,
10360                        cache,
10361                        fork_mode,
10362                        alternate_snapshot,
10363                        &h_seed_buf,
10364                        &fill_prev,
10365                        rt,
10366                        fence[1],
10367                        self.layers.len(),
10368                    )?;
10369                    eprintln!(
10370                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
10371                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
10372                        fence[1],
10373                        fork.logical_payload_bytes[0],
10374                        fork.logical_payload_bytes[1],
10375                        fork.controller.map_or(0.0, |policy| policy.threshold),
10376                    );
10377                    fork_snapshot = Some(current_snapshot);
10378                    opti_fork = Some(fork);
10379                }
10380            }
10381        }
10382        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
10383        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
10384        let mut snap = match fork_snapshot {
10385            Some(snapshot) => snapshot,
10386            None => cache.snapshot(e)?,
10387        };
10388        let mut carried_opti: Option<OptiControllerTicket> = None;
10389        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
10390        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
10391        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
10392            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
10393        } else {
10394            None
10395        };
10396        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
10397        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
10398        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
10399        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
10400        // pass of any kind). Verify still
10401        // checks every emitted token against the target -> exactness holds by construction; only
10402        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
10403        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
10404        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
10405        let mut pending: Option<u32> = carried_pending;
10406        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
10407        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
10408        // the verify accept readback). Printed once at loop end via spec-stats.
10409        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
10410        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
10411        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
10412        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
10413        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
10414        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
10415        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
10416        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
10417        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
10418        let mut ph_wait = 0f64;
10419        let mut ph_commit = 0f64;
10420        let mut ph_t = std::time::Instant::now();
10421        let mut ph_mark = |acc: &mut f64, on: bool| {
10422            if on {
10423                let now = std::time::Instant::now();
10424                *acc += (now - ph_t).as_secs_f64();
10425                ph_t = now;
10426            }
10427        };
10428        if let Some(p) = pipe {
10429            p.setup_end();
10430        }
10431        while keep_going && out.len() < max_new {
10432            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
10433            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
10434            if let (true, Some(sg), Some(ptrs)) = (
10435                stream_active && round >= 1 && pending.is_some(),
10436                &stream_graph,
10437                &stream_ptrs,
10438            ) {
10439                if debug_spec {
10440                    static ONCE: std::sync::Once = std::sync::Once::new();
10441                    ONCE.call_once(|| {
10442                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
10443                    });
10444                }
10445                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
10446                e.set_u32_one(&mut pend_d, pending.unwrap())?;
10447                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
10448                for _mi in 0..m_rounds {
10449                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
10450                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
10451                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
10452                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
10453                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
10454                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
10455                    sg.launch()?;
10456                    e.spec_assemble_verify(
10457                        &g_tokp2k,
10458                        &pend_d,
10459                        d2t_dev.as_ref(),
10460                        &mut vtok_d,
10461                        &mut brk_d,
10462                        p_min,
10463                        k,
10464                        pmin0,
10465                    )?;
10466                    let mut ck = VerifyCkpt::new(self.layers.len());
10467                    let dummy = vec![0u32; t_v_s];
10468                    let (tl_d, vx) = self.decode_step_t_core_stream(
10469                        e,
10470                        &dummy,
10471                        0,
10472                        &mut *cache,
10473                        embd_dev,
10474                        Some(&mut ck),
10475                        Some((&vtok_d, &pos_ctr)),
10476                        None,
10477                        None,
10478                        None,
10479                    )?;
10480                    for j in 0..t_v_s {
10481                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
10482                    }
10483                    e.spec_accept_greedy_dc(
10484                        &preds_d,
10485                        &vtok_d,
10486                        &last_pred_d,
10487                        &brk_d,
10488                        &mut stream_acc,
10489                    )?;
10490                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
10491                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
10492                    self.commit_verified_prefix_stream(
10493                        e,
10494                        &mut *cache,
10495                        &snap,
10496                        &ck,
10497                        &stream_acc,
10498                        1,
10499                        t_v_s,
10500                    )?;
10501                    e.spec_rollback_stream(
10502                        ptrs,
10503                        &pos_start_d,
10504                        &stream_acc,
10505                        1,
10506                        self.layers.len() + 1,
10507                    )?;
10508                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
10509                }
10510                e.stream().synchronize()?;
10511                let ring_h = e.dtoh_u32(&ring_d)?;
10512                let cnt = ring_h[0] as usize;
10513                for i in 0..cnt {
10514                    if out.len() < max_new {
10515                        out.push(ring_h[1 + i]);
10516                    }
10517                }
10518                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
10519                for il in 0..self.layers.len() {
10520                    if let Some(kvl) = cache.kv[il].as_mut() {
10521                        kvl.len = pos_h;
10522                    }
10523                }
10524                cache.pos = pos_h;
10525                scratch.kv.len = pos_h;
10526                pending = Some(ring_h[cnt]); // last drained token = the live bonus
10527                last_token = ring_h[cnt];
10528                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
10529                total_accepted += cnt.saturating_sub(m_rounds);
10530                if let Some(t) = sess_telem {
10531                    // totals only — the burst's per-round accept counts stayed on device
10532                    // (that is the point of the round-stream arm). pos_* untouched.
10533                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
10534                }
10535                round += m_rounds;
10536                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
10537                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10538                continue;
10539            }
10540            let pipe_draft = match pipe {
10541                Some(p) => Some(p.draft_begin(round)?),
10542                None => None,
10543            };
10544            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
10545            let mut current_opti = carried_opti.take();
10546            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
10547                match opti_fork.as_mut() {
10548                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
10549                    None => None,
10550                    Some(_) => None,
10551                }
10552            } else {
10553                None
10554            };
10555            if current_opti.is_none() {
10556                if let Some(fork) = opti_fork.as_ref() {
10557                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
10558                } else {
10559                    cache.snapshot_into(e, &mut snap)?;
10560                }
10561            } else if snap.pos != pos {
10562                return Err(format!(
10563                    "optipipe carried snapshot pos {} != current pos {pos}",
10564                    snap.pos
10565                )
10566                .into());
10567            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
10568            ph_mark(&mut ph_rest, phase_on);
10569
10570            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
10571            // p-min semantics (both paths): stop the chain early when the head's confidence in
10572            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
10573            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
10574            let base0 = if pending.is_some() { 1usize } else { 0usize };
10575            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
10576            // accepted run + 1 (the gemma law — see the setup block above the loop).
10577            let k_this = if adapt { kc } else { k };
10578            let mut draft: Vec<u32> = Vec::with_capacity(k);
10579            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
10580            let mut controller_draft_prob: Option<f32> = None;
10581            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
10582            if let Some(ticket) = current_opti.as_mut() {
10583                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
10584                if ticket.verify_tokens[0] != carried_pending {
10585                    return Err(format!(
10586                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
10587                        ticket.verify_tokens[0],
10588                    )
10589                    .into());
10590                }
10591                draft.push(ticket.verify_tokens[1]);
10592                controller_draft_prob = Some(ticket.draft_prob);
10593                controller_eager_state = ticket
10594                    .take_eager_seed()
10595                    .map(|seed| (ticket.verify_tokens[1], seed));
10596            } else {
10597                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
10598                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
10599                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
10600                // rejected drafts and p-min extras via the len mechanism).
10601                scratch.set_len(e, pos + base0 - 1)?;
10602                if pen_on {
10603                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
10604                    // device dedup: `penalty_last_n` is usize::MAX for any serve request with
10605                    // a penalty, so without the cap this grew with the whole session.
10606                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
10607                    let w0 = pen_hist.len().saturating_sub(win);
10608                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
10609                }
10610                if sampled {
10611                    draft_logits.clear();
10612                    draft_stats.clear();
10613                }
10614                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
10615                // position's mask is computed on that clone and advanced by the PROPOSED token. The
10616                // real state moves only on emission (verify's job), so the emitted stream is
10617                // unchanged — the mask only removes tokens the verify would have truncated anyway.
10618                let mut dmask_live = dmask_on;
10619                if dmask_live {
10620                    let t_c = std::time::Instant::now();
10621                    constraint
10622                        .as_deref_mut()
10623                        .unwrap()
10624                        .draft_begin()
10625                        .map_err(|e2| format!("constraint: {e2}"))?;
10626                    dm_clone_ns += t_c.elapsed().as_nanos();
10627                    dm_rounds += 1;
10628                }
10629                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
10630                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
10631                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
10632                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
10633                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
10634                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
10635                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
10636                    for j in 0..k_this {
10637                        // per-position mask upload (contents only — the graph's baked pointer is
10638                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
10639                        // mask node degrades to a no-op ban instead of needing a second graph.
10640                        if dmask_live
10641                            && !upload_draft_mask(
10642                                e,
10643                                constraint.as_deref_mut().unwrap(),
10644                                &mut dctx.g_dmask,
10645                                mtp.d2t.as_ref(),
10646                                d_vocab,
10647                                dmask_words,
10648                            )?
10649                        {
10650                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
10651                            // genuinely miss the legal set): neutralize the captured mask node and
10652                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
10653                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
10654                            dmask_live = false;
10655                        }
10656                        gr.launch()?;
10657                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
10658                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
10659                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
10660                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
10661                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
10662                        // replay's embed node, and the MMU fault kills the CUDA context for the
10663                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
10664                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
10665                        // buffer (g_seed = the verify-side handoff vs head-side compute).
10666                        if (idx as usize) >= d_vocab {
10667                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
10668                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
10669                            // seed, untouched since the round-start copy — the pair discriminates
10670                            // "seed arrived poisoned" from "head forward produced NaN".
10671                            let seed_h = e.dtoh(&dctx.g_seed)?;
10672                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
10673                            let in_h = e.dtoh(&h_seed_buf)?;
10674                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
10675                            return Err(format!(
10676                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
10677                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
10678                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
10679                             the embed row (#87 trap)"
10680                            )
10681                            .into());
10682                        }
10683                        // trimmed draft vocab -> target token id (identity when no d2t map)
10684                        let d = match &mtp.d2t {
10685                            Some(map) => map[idx as usize],
10686                            None => idx,
10687                        };
10688                        let draft_p = if p_min > 0.0
10689                            || opti_fork
10690                                .as_ref()
10691                                .is_some_and(|fork| fork.controller.is_some())
10692                        {
10693                            Some(e.dtoh(&dctx.g_p)?[0])
10694                        } else {
10695                            None
10696                        };
10697                        if j == 0 {
10698                            controller_draft_prob = draft_p;
10699                        }
10700                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
10701                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
10702                                break;
10703                            }
10704                        }
10705                        draft.push(d);
10706                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
10707                        // index the argmax wrote — patch the persistent token buffer (4B htod).
10708                        if d != idx {
10709                            e.set_u32_one(&mut dctx.g_tok, d)?;
10710                        }
10711                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
10712                        // unmasked drafting for the remaining positions (verify still arbitrates).
10713                        // speculative advance; a chain the grammar can no longer follow (EOS
10714                        // proposed) ends here. The captured mask node always runs, so a dead chain
10715                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
10716                        if dmask_live
10717                            && !constraint
10718                                .as_deref_mut()
10719                                .unwrap()
10720                                .draft_advance(d)
10721                                .map_err(|e2| format!("constraint: {e2}"))?
10722                        {
10723                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
10724                            break;
10725                        }
10726                    }
10727                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
10728                // legal ONLY in the regime it was captured in. The condition used to read
10729                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
10730                // which it could not, because the key omitted the filters. Both halves are now
10731                // enforced: the key drops a stale graph, and this site refuses to launch one.
10732                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
10733                    if skey_probe() {
10734                        eprintln!(
10735                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
10736                             top_p={} min_p={} s_key_parked={:?}",
10737                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
10738                        );
10739                    }
10740                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
10741                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
10742                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
10743                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
10744                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
10745                    // stream. Host sctr advances in lockstep (computed, no readback needed).
10746                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
10747                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
10748                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
10749                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
10750                    for j in 0..k_this {
10751                        gr.launch()?;
10752                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
10753                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
10754                        // counts the p-min-discarded token too)
10755                        // q retention: ONE async D2D of the persistent head-logits buffer into this
10756                        // round's slot j (stream-ordered after the replay, before the next one).
10757                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
10758                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
10759                        // #87 SENTINEL TRAP (see the greedy graph arm above).
10760                        if (idx as usize) >= d_vocab {
10761                            let seed_h = e.dtoh(&dctx.g_seed)?;
10762                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
10763                            return Err(format!(
10764                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
10765                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
10766                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
10767                             (#87 trap)"
10768                            )
10769                            .into());
10770                        }
10771                        let d = match &mtp.d2t {
10772                            Some(map) => map[idx as usize],
10773                            None => idx,
10774                        };
10775                        draft_idx.push(idx);
10776                        if p_min > 0.0 {
10777                            let p = e.dtoh(&dctx.g_p)?[0];
10778                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
10779                                break;
10780                            }
10781                        }
10782                        draft.push(d);
10783                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
10784                        if d != idx {
10785                            e.set_u32_one(&mut dctx.g_tok, d)?;
10786                        }
10787                    }
10788                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
10789                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
10790                    for j in 0..draft.len().max(draft_idx.len()) {
10791                        let rows0 = e.htod_i32(&[0])?;
10792                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10793                        e.filter_stats(
10794                            &dctx.q_slots[j],
10795                            d_vocab,
10796                            &rows0,
10797                            &mut th_d,
10798                            &mut z_d,
10799                            &mut mx_d,
10800                            d_vocab,
10801                            1,
10802                            sp_temp,
10803                            sp.top_k,
10804                            sp.top_p,
10805                            sp.min_p,
10806                        )?;
10807                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
10808                    }
10809                } else {
10810                    if skey_probe() && sampled {
10811                        eprintln!(
10812                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
10813                             top_p={} min_p={} s_key_parked={:?}",
10814                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
10815                        );
10816                    }
10817                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
10818                    let mut e_tok = last_token;
10819                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
10820                    for j in 0..k_this {
10821                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
10822                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
10823                        let mtp_pos = pos + base0 + j;
10824                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
10825                        // A position with no legal draft-vocab row drops to unmasked drafting for
10826                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
10827                        if dmask_live {
10828                            dmask_live = upload_draft_mask(
10829                                e,
10830                                constraint.as_deref_mut().unwrap(),
10831                                &mut dctx.g_dmask,
10832                                mtp.d2t.as_ref(),
10833                                d_vocab,
10834                                dmask_words,
10835                            )?;
10836                        }
10837                        let (dl_d, h_nextn) = self.mtp_head_forward_dev(
10838                            e,
10839                            mtp,
10840                            e_tok,
10841                            &d_seed,
10842                            &mut *scratch,
10843                            mtp_pos,
10844                            embd_dev,
10845                            if dmask_live {
10846                                Some((&dctx.g_dmask, dmask_words))
10847                            } else {
10848                                None
10849                            },
10850                        )?;
10851                        let tok_d = if sampled {
10852                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
10853                            // the filtered softmax (filters off => th=0, exact v1 semantics).
10854                            if perturb_buf.is_none() {
10855                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
10856                            }
10857                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
10858                            if pen_on {
10859                                let h = pen_hist_d.as_ref().unwrap();
10860                                let nh = h.len();
10861                                e.penalize_logits(
10862                                    &mut q_row,
10863                                    h,
10864                                    nh,
10865                                    sp.penalty_repeat,
10866                                    sp.penalty_freq,
10867                                    sp.penalty_present,
10868                                    d_vocab,
10869                                )?;
10870                            }
10871                            let rows0 = e.htod_i32(&[0])?;
10872                            let (mut th_d, mut z_d, mut mx_d) =
10873                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
10874                            e.filter_stats(
10875                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
10876                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
10877                            )?;
10878                            let (th, z, mx) =
10879                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
10880                            let pb = perturb_buf.as_mut().unwrap();
10881                            e.gumbel_perturb_filtered(
10882                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
10883                            )?;
10884                            sctr += 1;
10885                            draft_logits.push(q_row);
10886                            draft_stats.push((mx, th, z));
10887                            e.argmax_token_device(pb, d_vocab)?
10888                        } else {
10889                            e.argmax_token_device(&dl_d, d_vocab)?
10890                        };
10891                        let idx = e.dtoh_u32_one(&tok_d)?;
10892                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
10893                        // here because the eager chain's operands are all readable: dl_d (the head
10894                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
10895                        if (idx as usize) >= d_vocab {
10896                            let dl_h = e.dtoh(&dl_d)?;
10897                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
10898                            let seed_h = e.dtoh(&d_seed)?;
10899                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
10900                            return Err(format!(
10901                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
10902                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
10903                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
10904                             embed row (#87 trap)"
10905                            )
10906                            .into());
10907                        }
10908                        let d = match &mtp.d2t {
10909                            Some(map) => map[idx as usize],
10910                            None => idx,
10911                        };
10912                        if sampled {
10913                            draft_idx.push(idx);
10914                        }
10915                        let draft_p = if p_min > 0.0
10916                            || opti_fork
10917                                .as_ref()
10918                                .is_some_and(|fork| fork.controller.is_some())
10919                        {
10920                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
10921                            Some(e.dtoh(&p_d)?[0])
10922                        } else {
10923                            None
10924                        };
10925                        if j == 0 {
10926                            controller_draft_prob = draft_p;
10927                        }
10928                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
10929                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
10930                                break;
10931                            }
10932                        }
10933                        draft.push(d);
10934                        e_tok = d;
10935                        d_seed = h_nextn;
10936                        // speculative advance; a chain the grammar can no longer follow (EOS
10937                        // proposed) ends here — the prefix already proposed still rides verify.
10938                        if dmask_live
10939                            && !constraint
10940                                .as_deref_mut()
10941                                .unwrap()
10942                                .draft_advance(d)
10943                                .map_err(|e2| format!("constraint: {e2}"))?
10944                        {
10945                            break;
10946                        }
10947                    }
10948                    if opti_fork
10949                        .as_ref()
10950                        .is_some_and(|fork| fork.controller.is_some())
10951                    {
10952                        controller_eager_state = Some((e_tok, d_seed));
10953                    }
10954                }
10955            }
10956            let k_round = draft.len();
10957            if let Some(p) = pipe {
10958                p.draft_end(round);
10959            }
10960            drop(pipe_draft);
10961
10962            ph_mark(&mut ph_draft, phase_on);
10963            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
10964            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
10965            let verify_tokens: Vec<u32> = match pending {
10966                Some(b) => {
10967                    let mut v = Vec::with_capacity(k_round + 1);
10968                    v.push(b);
10969                    v.extend_from_slice(&draft);
10970                    v
10971                }
10972                None => draft.clone(),
10973            };
10974            let base = if pending.is_some() { 1 } else { 0 };
10975            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
10976            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
10977            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
10978                Some(ticket.take_ckpt())
10979            } else if spec_replay {
10980                None
10981            } else {
10982                Some(VerifyCkpt::new(self.layers.len()))
10983            };
10984            let controller_can_probe = base == 1
10985                && k_round == 1
10986                && out.len().saturating_add(2) < max_new
10987                && controller_draft_prob.is_some()
10988                && opti_fork
10989                    .as_ref()
10990                    .and_then(|fork| fork.controller.as_ref())
10991                    .is_some_and(|policy| !policy.breaker_tripped);
10992            let mut successor_attempt: Option<OptiControllerTicket> = None;
10993            let mut rejected_probe: Option<(f32, u32)> = None;
10994            let mut controller_prepared: Option<OptiControllerPrepared> = None;
10995            if controller_can_probe {
10996                // Prepare d2/q and, on admission, d3 before either current verify half is
10997                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
10998                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
10999                // the primary stream after N stage 1 would serialize the supposed pipeline.
11000                let eager_pos = scratch.kv.len + 1;
11001                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
11002                    e,
11003                    mtp,
11004                    &mut dctx,
11005                    &mut *scratch,
11006                    d_vocab,
11007                    &mut controller_eager_state,
11008                    eager_pos,
11009                    embd_dev,
11010                )?;
11011                let first_probability = controller_draft_prob
11012                    .ok_or("optipipe controller probe lost first-token probability")?;
11013                let q_proxy = first_probability * pending_probability;
11014                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11015                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11016                let admitted = opti_fork
11017                    .as_ref()
11018                    .and_then(|fork| fork.controller.as_ref())
11019                    .ok_or("optipipe controller policy disappeared")?
11020                    .admit(q_proxy);
11021                if admitted {
11022                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11023                    let eager_pos = scratch.kv.len + 1;
11024                    let (optimistic_draft, optimistic_draft_probability) = self
11025                        .opti_controller_draft_step(
11026                            e,
11027                            mtp,
11028                            &mut dctx,
11029                            &mut *scratch,
11030                            d_vocab,
11031                            &mut controller_eager_state,
11032                            eager_pos,
11033                            embd_dev,
11034                        )?;
11035                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11036                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
11037                        debug_assert_eq!(token, optimistic_draft);
11038                        seed
11039                    });
11040                    controller_prepared = Some(OptiControllerPrepared {
11041                        verify_tokens: [optimistic_pending, optimistic_draft],
11042                        draft_prob: optimistic_draft_probability,
11043                        eager_seed,
11044                        q_proxy,
11045                        scratch_len: scratch.kv.len,
11046                    });
11047                } else {
11048                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11049                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11050                    rejected_probe = Some((q_proxy, optimistic_pending));
11051                    eprintln!(
11052                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
11053                        opti_fork
11054                            .as_ref()
11055                            .and_then(|fork| fork.controller.as_ref())
11056                            .expect("controller policy")
11057                            .threshold,
11058                    );
11059                }
11060            }
11061            let fork_attempt = match fork_generation.take() {
11062                Some(generation) if base == 1 && k_round == 1 => Some(generation),
11063                Some(generation) => {
11064                    opti_fork
11065                        .as_mut()
11066                        .expect("fork generation without fork state")
11067                        .retire(generation)?;
11068                    None
11069                }
11070                None => None,
11071            };
11072            let (tlogits_d, vx) = if let Some(p) = pipe {
11073                self.decode_step_t_core_pipelined(
11074                    e,
11075                    &verify_tokens,
11076                    pos,
11077                    &mut *cache,
11078                    embd_dev,
11079                    ckpt.as_mut(),
11080                    p,
11081                    round,
11082                )?
11083            } else if controller_can_probe {
11084                let fence = opti_fork
11085                    .as_ref()
11086                    .ok_or("optipipe controller probe lost fork state")?
11087                    .fence;
11088                let boundary = match current_opti.as_mut() {
11089                    Some(ticket) => ticket.take_boundary(),
11090                    None => self.verify_stage0_issue(
11091                        e,
11092                        &verify_tokens,
11093                        pos,
11094                        &mut *cache,
11095                        embd_dev,
11096                        ckpt.as_mut(),
11097                        None,
11098                        &fence,
11099                        Some(true),
11100                        None,
11101                    )?,
11102                };
11103                if let Some(prepared) = controller_prepared.take() {
11104                    let generation = {
11105                        let fork = opti_fork
11106                            .as_mut()
11107                            .ok_or("optipipe controller admission lost fork state")?;
11108                        let generation = fork.reserve_successor()?;
11109                        let rt = fork.rt;
11110                        let snapshot_fence = fork.fence;
11111                        opti_snapshot_one_stage_owned_into(
11112                            e,
11113                            cache,
11114                            rt,
11115                            &snapshot_fence,
11116                            0,
11117                            fork.successor_snapshot_mut(),
11118                        )?;
11119                        generation
11120                    };
11121                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
11122                    let successor_boundary = self.verify_stage0_issue(
11123                        e,
11124                        &prepared.verify_tokens,
11125                        pos + verify_tokens.len(),
11126                        &mut *cache,
11127                        embd_dev,
11128                        Some(&mut successor_ckpt),
11129                        None,
11130                        &fence,
11131                        Some(false),
11132                        None,
11133                    )?;
11134                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11135                    let fork = opti_fork
11136                        .as_ref()
11137                        .ok_or("optipipe controller ticket lost fork state")?;
11138                    successor_attempt = Some(fork.controller_ticket(
11139                        generation,
11140                        successor_boundary,
11141                        successor_ckpt,
11142                        prepared.verify_tokens,
11143                        prepared.draft_prob,
11144                        prepared.eager_seed,
11145                        prepared.q_proxy,
11146                        prepared.scratch_len,
11147                    ));
11148                    eprintln!(
11149                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
11150                         verify={:?}",
11151                        generation.id,
11152                        prepared.q_proxy,
11153                        fork.controller.expect("controller policy").threshold,
11154                        prepared.verify_tokens,
11155                    );
11156                }
11157                let result = self.verify_stage1_finish(
11158                    e,
11159                    boundary,
11160                    &mut *cache,
11161                    ckpt.as_mut(),
11162                    None,
11163                    &fence,
11164                    successor_attempt.is_none(),
11165                )?;
11166                if let Some(ticket) = current_opti.as_mut() {
11167                    ticket.settle();
11168                }
11169                if successor_attempt.is_some() {
11170                    let fork = opti_fork
11171                        .as_mut()
11172                        .ok_or("optipipe successor snapshot lost fork state")?;
11173                    let rt = fork.rt;
11174                    let snapshot_fence = fork.fence;
11175                    opti_snapshot_one_stage_owned_into(
11176                        e,
11177                        cache,
11178                        rt,
11179                        &snapshot_fence,
11180                        1,
11181                        fork.successor_snapshot_mut(),
11182                    )?;
11183                    // Publish N only after both independent successor-state queues are complete.
11184                    fork.rt.publish_to(1, &e.stream())?;
11185                }
11186                result
11187            } else if let Some(ticket) = current_opti.as_mut() {
11188                let fork = opti_fork
11189                    .as_mut()
11190                    .ok_or("optipipe carried controller ticket lost fork state")?;
11191                let boundary = ticket.take_boundary();
11192                let result = self.verify_stage1_finish(
11193                    e,
11194                    boundary,
11195                    &mut *cache,
11196                    ckpt.as_mut(),
11197                    None,
11198                    &fork.fence,
11199                    true,
11200                )?;
11201                ticket.settle();
11202                result
11203            } else if let Some(generation) = fork_attempt {
11204                let fork = opti_fork
11205                    .as_mut()
11206                    .expect("fork generation without fork state");
11207                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
11208                let action = fork.mode.action(generation.id);
11209                let boundary = self.verify_stage0_issue(
11210                    e,
11211                    &verify_tokens,
11212                    pos,
11213                    &mut *cache,
11214                    embd_dev,
11215                    ckpt.as_mut(),
11216                    None,
11217                    &fork.fence,
11218                    Some(true),
11219                    None,
11220                )?;
11221                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11222                let mut ticket = fork.ticket(generation, boundary);
11223                if action == OptiForkAction::Abort {
11224                    return Err(format!(
11225                        "optipipe forced abort with generation {} stage0 in flight",
11226                        generation.id,
11227                    )
11228                    .into());
11229                }
11230                fork.reconcile(
11231                    e,
11232                    &mut *cache,
11233                    &mut *scratch,
11234                    &snap,
11235                    &mut h_seed_buf,
11236                    &mut fill_prev,
11237                    generation,
11238                    action,
11239                    verify_tokens[0],
11240                )?;
11241                let result = if action == OptiForkAction::Hit {
11242                    let boundary = ticket.take_boundary();
11243                    self.verify_stage1_finish(
11244                        e,
11245                        boundary,
11246                        &mut *cache,
11247                        ckpt.as_mut(),
11248                        None,
11249                        &fork.fence,
11250                        true,
11251                    )?
11252                } else {
11253                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
11254                    // verify only after E_restart published the restored stage-0 state.
11255                    self.decode_step_t_core(
11256                        e,
11257                        &verify_tokens,
11258                        pos,
11259                        &mut *cache,
11260                        embd_dev,
11261                        ckpt.as_mut(),
11262                    )?
11263                };
11264                ticket.settle();
11265                debug_assert_eq!(ticket.generation, generation);
11266                fork.retire(generation)?;
11267                result
11268            } else {
11269                self.decode_step_t_core(
11270                    e,
11271                    &verify_tokens,
11272                    pos,
11273                    &mut *cache,
11274                    embd_dev,
11275                    ckpt.as_mut(),
11276                )?
11277            };
11278            let pipe_accept = match pipe {
11279                Some(p) => Some(p.accept_begin(round)?),
11280                None => None,
11281            };
11282
11283            ph_mark(&mut ph_verify, phase_on);
11284            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
11285            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
11286            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
11287            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
11288            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
11289            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
11290            // (== the bonus), so every index shifts by `base` and last_pred is unused.
11291            let t_v = verify_tokens.len();
11292            let mut preds: Vec<u32> = Vec::new();
11293            if !sampled {
11294                for j in 0..t_v {
11295                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
11296                }
11297                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
11298                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
11299                // next round's last_token = the next chain's embed lookup. Catch it at the
11300                // source with the column named — an all-NaN VERIFY column implicates the
11301                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
11302                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
11303                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
11304                    let mut probe = e.zeros(n_vocab)?;
11305                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
11306                    let col_h = e.dtoh(&probe)?;
11307                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
11308                    return Err(format!(
11309                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
11310                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
11311                         — the stage-split verify produced a poisoned column (#87 trap)",
11312                        preds[bad]
11313                    )
11314                    .into());
11315                }
11316            }
11317            ph_mark(&mut ph_wait, phase_on);
11318            let t_pred = |j: usize| -> u32 {
11319                if j == 0 && base == 0 {
11320                    last_pred
11321                } else {
11322                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
11323                    // used to call this from the sampled arm and panicked the worker; it now goes
11324                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
11325                    // out-of-range pred is a real bug, not something to paper over.
11326                    debug_assert!(
11327                        !sampled,
11328                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
11329                    );
11330                    preds[base + j - 1]
11331                }
11332            };
11333            let mut devacc_seeded = false;
11334            let mut devacc_acc: Option<CudaSlice<u32>> = None;
11335            let (n_acc, bonus) = if !sampled {
11336                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
11337                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
11338                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
11339                // gated on token identity vs the host walk (the arms below are bit-equal rules).
11340                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
11341                {
11342                    let draft_d = e.htod_u32_v(&draft)?;
11343                    let mut acc_out = e.alloc_u32_zeroed(2)?;
11344                    e.spec_accept_greedy(
11345                        &preds_d,
11346                        &draft_d,
11347                        last_pred,
11348                        base,
11349                        k_round,
11350                        &mut acc_out,
11351                    )?;
11352                    devacc_acc = Some(acc_out.clone());
11353                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
11354                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
11355                    // non-replay commit arms skip their host-offset seed copies (guarded below);
11356                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
11357                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
11358                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
11359                    // the update lands after the arms (devacc_seeded guard below).
11360                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
11361                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
11362                    // unified rule; full accept rewrites the verify-left value). Host mirrors
11363                    // update after the readback; commit_verified_prefix skips its len_d writes.
11364                    if let Some(successor) = successor_attempt.as_ref() {
11365                        opti_fork
11366                            .as_mut()
11367                            .ok_or("optipipe successor reconcile lost fork state")?
11368                            .queue_actual_reconcile(
11369                                e,
11370                                &snap,
11371                                &acc_out,
11372                                successor.verify_tokens[0],
11373                                base,
11374                            )?;
11375                    } else if let Some(ptrs) = &kv_len_ptrs {
11376                        let saved: Vec<i32> = (0..self.layers.len())
11377                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
11378                            .collect();
11379                        let saved_d = e.htod_i32(&saved)?;
11380                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
11381                    }
11382                    devacc_seeded = true;
11383                    let ab = e.dtoh_u32(&acc_out)?;
11384                    (ab[0] as usize, ab[1])
11385                } else {
11386                    let mut n_acc = 0usize;
11387                    for j in 0..k_round {
11388                        if t_pred(j) == draft[j] {
11389                            n_acc += 1;
11390                        } else {
11391                            break;
11392                        }
11393                    }
11394                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
11395                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
11396                    (n_acc, t_pred(n_acc))
11397                }
11398            } else {
11399                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
11400                if col_buf.is_none() {
11401                    col_buf = Some(e.zeros(n_vocab)?);
11402                }
11403                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
11404                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
11405                let mut pj = vec![0f32; k_round.max(1)];
11406                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
11407                if k_round > 0 {
11408                    let mut ids: Vec<u32> = Vec::new();
11409                    let mut rows: Vec<i32> = Vec::new();
11410                    for j in 0..k_round {
11411                        if j > 0 || base == 1 {
11412                            ids.push(draft[j]);
11413                            rows.push((base + j) as i32 - 1);
11414                        }
11415                    }
11416                    if !ids.is_empty() {
11417                        let nr = rows.len();
11418                        // penalties: materialize the used columns into one contiguous penalized
11419                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
11420                        // penalties: materialize used columns contiguously, penalize all rows in
11421                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
11422                        let p_rows: Vec<i32> = if pen_on {
11423                            (0..nr as i32).collect()
11424                        } else {
11425                            rows.clone()
11426                        };
11427                        if pen_on {
11428                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
11429                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
11430                            }
11431                            let pc = pcol_buf.as_mut().unwrap();
11432                            for (i2, &r) in rows.iter().enumerate() {
11433                                let c = r as usize;
11434                                e.copy_view_into(
11435                                    pc,
11436                                    i2 * n_vocab,
11437                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
11438                                    n_vocab,
11439                                )?;
11440                            }
11441                            let h = pen_hist_d.as_ref().unwrap();
11442                            let nh = h.len();
11443                            e.penalize_logits_rows(
11444                                pc,
11445                                h,
11446                                nh,
11447                                sp.penalty_repeat,
11448                                sp.penalty_freq,
11449                                sp.penalty_present,
11450                                n_vocab,
11451                                nr,
11452                            )?;
11453                        }
11454                        let p_src: &CudaSlice<f32> = if pen_on {
11455                            pcol_buf.as_ref().unwrap()
11456                        } else {
11457                            &tlogits_d
11458                        };
11459                        let rowsd = e.htod_i32(&p_rows)?;
11460                        let (mut th_d, mut z_d, mut mx_d) =
11461                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
11462                        e.filter_stats(
11463                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
11464                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
11465                        )?;
11466                        let idsd = e.htod_u32_v(&ids)?;
11467                        let mut outd = e.zeros(nr)?;
11468                        e.softmax_gather_filtered(
11469                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
11470                            sp_temp,
11471                        )?;
11472                        let outv = e.dtoh(&outd)?;
11473                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
11474                        let mut oi = 0usize;
11475                        for j in 0..k_round {
11476                            if j > 0 || base == 1 {
11477                                pj[j] = outv[oi];
11478                                oi += 1;
11479                            }
11480                        }
11481                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
11482                    }
11483                    if base == 0 {
11484                        let lc: &CudaSlice<f32> = if pen_on {
11485                            if col_buf.is_none() {
11486                                col_buf = Some(e.zeros(n_vocab)?);
11487                            }
11488                            let cb = col_buf.as_mut().unwrap();
11489                            e.copy_into(
11490                                cb,
11491                                0,
11492                                last_col_logits
11493                                    .as_ref()
11494                                    .expect("sampled: last_col_logits unset"),
11495                                n_vocab,
11496                            )?;
11497                            let h = pen_hist_d.as_ref().unwrap();
11498                            let nh = h.len();
11499                            e.penalize_logits(
11500                                cb,
11501                                h,
11502                                nh,
11503                                sp.penalty_repeat,
11504                                sp.penalty_freq,
11505                                sp.penalty_present,
11506                                n_vocab,
11507                            )?;
11508                            col_buf.as_ref().unwrap()
11509                        } else {
11510                            last_col_logits
11511                                .as_ref()
11512                                .expect("sampled: last_col_logits unset")
11513                        };
11514                        let rows0 = e.htod_i32(&[0])?;
11515                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11516                        e.filter_stats(
11517                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
11518                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
11519                        )?;
11520                        let idsd = e.htod_u32_v(&[draft[0]])?;
11521                        let mut outd = e.zeros(1)?;
11522                        e.softmax_gather_filtered(
11523                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
11524                        )?;
11525                        pj[0] = e.dtoh(&outd)?[0];
11526                        last_col_stats =
11527                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
11528                    }
11529                }
11530                // q source: the graph arm retained the head logits in the persistent q_slots;
11531                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
11532                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
11533                // computes them post-replay — graph engages only filter/penalty-free, so the
11534                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
11535                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
11536                    &dctx.q_slots
11537                } else {
11538                    &draft_logits
11539                };
11540                let mut n_acc = 0usize;
11541                for j in 0..k_round {
11542                    let (qmx, qth, qz) = draft_stats[j];
11543                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
11544                    let rowsd = e.htod_i32(&[0])?;
11545                    let thd = e.htod(&[qth])?;
11546                    let zd = e.htod(&[qz])?;
11547                    let _ = qmx;
11548                    let mut outd = e.zeros(1)?;
11549                    e.softmax_gather_filtered(
11550                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
11551                        sp_temp,
11552                    )?;
11553                    let qj = e.dtoh(&outd)?[0];
11554                    let u = host_u01(sp_seed, uctr);
11555                    uctr += 1;
11556                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
11557                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
11558                    // exactness signature (see `skey_probe`). Impossible when the draft was
11559                    // drawn from the same filtered distribution the verify reconstructs here;
11560                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
11561                    if skey_probe() && qj == 0.0 {
11562                        eprintln!(
11563                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
11564                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
11565                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
11566                        );
11567                    }
11568                    if accept {
11569                        n_acc += 1;
11570                    } else {
11571                        break;
11572                    }
11573                }
11574                let bonus = if n_acc == k_round {
11575                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
11576                    let col = base + k_round - 1;
11577                    let cb = col_buf.as_mut().unwrap();
11578                    e.copy_view_into(
11579                        cb,
11580                        0,
11581                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
11582                        n_vocab,
11583                    )?;
11584                    if pen_on {
11585                        let h = pen_hist_d.as_ref().unwrap();
11586                        let nh = h.len();
11587                        e.penalize_logits(
11588                            cb,
11589                            h,
11590                            nh,
11591                            sp.penalty_repeat,
11592                            sp.penalty_freq,
11593                            sp.penalty_present,
11594                            n_vocab,
11595                        )?;
11596                    }
11597                    if perturb_buf.is_none() {
11598                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
11599                    }
11600                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
11601                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
11602                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
11603                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
11604                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
11605                    // last gathered column, in both base arms. `th` is a threshold in e-units of
11606                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
11607                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
11608                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
11609                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
11610                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
11611                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
11612                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
11613                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
11614                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
11615                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
11616                    // and row_max is unused once nothing is masked), so this fix is a byte-level
11617                    // no-op for the untruncated serve default. One extra one-block filter_stats
11618                    // per full-accept round is the whole cost.
11619                    let (mx, th) = {
11620                        let rows0 = e.htod_i32(&[0])?;
11621                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11622                        let cb0 = col_buf.as_ref().unwrap();
11623                        e.filter_stats(
11624                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
11625                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
11626                        )?;
11627                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
11628                    };
11629                    let pb = perturb_buf.as_mut().unwrap();
11630                    let cb2 = col_buf.as_ref().unwrap();
11631                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
11632                    sctr += 1;
11633                    let td = e.argmax_token_device(pb, n_vocab)?;
11634                    e.dtoh_u32_one(&td)?
11635                } else {
11636                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
11637                    let cb = col_buf.as_mut().unwrap();
11638                    if n_acc > 0 || base == 1 {
11639                        let col = base + n_acc - 1;
11640                        e.copy_view_into(
11641                            cb,
11642                            0,
11643                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
11644                            n_vocab,
11645                        )?;
11646                    } else {
11647                        let lc = last_col_logits.as_ref().unwrap();
11648                        e.copy_into(cb, 0, lc, n_vocab)?;
11649                    }
11650                    if pen_on {
11651                        let h = pen_hist_d.as_ref().unwrap();
11652                        let nh = h.len();
11653                        e.penalize_logits(
11654                            cb,
11655                            h,
11656                            nh,
11657                            sp.penalty_repeat,
11658                            sp.penalty_freq,
11659                            sp.penalty_present,
11660                            n_vocab,
11661                        )?;
11662                    }
11663                    let cb2 = col_buf.as_ref().unwrap();
11664                    let sc = sctr;
11665                    sctr += 1;
11666                    // p-stats for the reject column: from col_stats when the col was gathered,
11667                    // else (j==0&&base==0) from last_col_stats.
11668                    let p_stats = if n_acc > 0 || base == 1 {
11669                        // col index within the gathered set == number of gathered cols before n_acc
11670                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
11671                        col_stats.get(gi).copied().unwrap_or_else(|| {
11672                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
11673                        })
11674                    } else {
11675                        last_col_stats.expect("sampled: last_col_stats unset at reject")
11676                    };
11677                    let q_stats = draft_stats[n_acc];
11678                    if let Some(map) = &d2t_dev {
11679                        if q_full_buf.is_none() {
11680                            q_full_buf = Some(e.zeros(n_vocab)?);
11681                        }
11682                        let qf = q_full_buf.as_mut().unwrap();
11683                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
11684                        let qf2 = q_full_buf.as_ref().unwrap();
11685                        e.residual_sample_filtered(
11686                            cb2,
11687                            Some(qf2),
11688                            n_vocab,
11689                            sp_temp,
11690                            sp_seed,
11691                            sc,
11692                            p_stats,
11693                            q_stats,
11694                            &mut sample_tok,
11695                        )?;
11696                    } else {
11697                        e.residual_sample_filtered(
11698                            cb2,
11699                            Some(&q_bufs[n_acc]),
11700                            n_vocab,
11701                            sp_temp,
11702                            sp_seed,
11703                            sc,
11704                            p_stats,
11705                            q_stats,
11706                            &mut sample_tok,
11707                        )?;
11708                    }
11709                    e.dtoh_u32(&sample_tok)?[0]
11710                };
11711                (n_acc, bonus)
11712            };
11713            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
11714            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
11715            // ordering). Walk the accepted drafts through the grammar in commit order; the
11716            // first illegal token truncates acceptance at its slot, and that slot's emission
11717            // is recomputed as the MASKED argmax of the target's own verify column — token-
11718            // identical to constrained plain greedy decode (an unmasked argmax that is
11719            // grammar-legal IS the masked argmax: masking only removes competitors). The
11720            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
11721            // measured in acceptance numbers, never hidden.
11722            let (n_acc, bonus) = match constraint.as_deref_mut() {
11723                None => (n_acc, bonus),
11724                Some(c) => {
11725                    fn ce(e2: String) -> Box<dyn std::error::Error> {
11726                        format!("constraint: {e2}").into()
11727                    }
11728                    let mut na = n_acc;
11729                    let mut cut = false;
11730                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
11731                        if c.is_allowed(d).map_err(ce)? {
11732                            c.consume(d).map_err(ce)?;
11733                        } else {
11734                            na = j;
11735                            cut = true;
11736                            dm_cut_tokens += n_acc - j;
11737                            break;
11738                        }
11739                    }
11740                    if cut {
11741                        dm_cuts += 1;
11742                    }
11743                    let mut bo = bonus;
11744                    if cut || !c.is_allowed(bo).map_err(ce)? {
11745                        let mut row = if na == 0 && base == 0 {
11746                            init_logits_host
11747                                .clone()
11748                                .ok_or("constraint: init logits missing (round-0 cut)")?
11749                        } else {
11750                            e.dtoh_view(
11751                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
11752                            )?
11753                        };
11754                        c.mask_logits(&mut row).map_err(ce)?;
11755                        bo = argmax(&row) as u32;
11756                    }
11757                    c.consume(bo).map_err(ce)?;
11758                    (na, bo)
11759                }
11760            };
11761            let mut successor_valid = false;
11762            if let Some((q_proxy, expected_d2)) = rejected_probe {
11763                let v_n = n_acc == 1 && bonus == expected_d2;
11764                eprintln!(
11765                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
11766                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
11767                );
11768            }
11769            if let Some(successor) = successor_attempt.as_ref() {
11770                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
11771                let generation = successor.generation;
11772                let q_proxy = successor.q_proxy;
11773                let expected_pending = successor.verify_tokens[0];
11774                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
11775                let fork = opti_fork
11776                    .as_mut()
11777                    .ok_or("optipipe successor resolution lost fork state")?;
11778                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
11779                if successor_valid {
11780                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11781                } else {
11782                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11783                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11784                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
11785                }
11786                let breaker_tripped = fork
11787                    .controller
11788                    .as_mut()
11789                    .expect("controller policy")
11790                    .resolve(successor_valid);
11791                if breaker_tripped {
11792                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11793                }
11794                eprintln!(
11795                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
11796                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
11797                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
11798                    generation.id, successor_valid, !successor_valid, breaker_tripped,
11799                );
11800                if !successor_valid {
11801                    let mut successor = successor_attempt
11802                        .take()
11803                        .expect("controller successor disappeared on miss");
11804                    successor.settle();
11805                    fork.retire(generation)?;
11806                }
11807            }
11808            total_drafted += k_round;
11809            total_accepted += n_acc;
11810            if let Some(t) = sess_telem {
11811                // Greedy, rejection-sampling, and grammar truncation all converge here after
11812                // the accept decision is already on host. Fixed-size relaxed atomics only.
11813                t.record_round(k_round, n_acc);
11814            }
11815            if spec_stats {
11816                st_len_hist[k_round] += 1;
11817                for j in 0..k_round {
11818                    st_drafted[j] += 1;
11819                }
11820                for j in 0..n_acc {
11821                    st_accepted[j] += 1;
11822                }
11823                if n_acc == k_round {
11824                    st_full += 1;
11825                }
11826            }
11827
11828            if debug_spec {
11829                eprintln!(
11830                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
11831                    out.len(),
11832                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
11833                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
11834                    // the GPU worker thread — a debug flag that killed the exact regime you would
11835                    // set it to investigate. See `debug_t_pred0`.
11836                    debug_t_pred0(sampled, base, last_pred, &preds)
11837                );
11838            }
11839
11840            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
11841            let commit_started = std::time::Instant::now();
11842            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
11843            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
11844            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
11845            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
11846            for j in 0..n_acc {
11847                if !session_mode && out.len() >= max_new {
11848                    break;
11849                }
11850                out.push(draft[j]);
11851            }
11852            if pen_on {
11853                pen_hist.extend_from_slice(&draft[0..n_acc]);
11854                pen_hist.push(bonus);
11855            }
11856            let bonus_emitted = session_mode || out.len() < max_new;
11857            if bonus_emitted {
11858                out.push(bonus);
11859            }
11860            last_token = bonus;
11861
11862            // --- 5. ROLLBACK + advance (§C) ---
11863            if n_acc == k_round && !spec_replay {
11864                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
11865                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
11866                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
11867                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
11868                // last_pred is dead in the pending path (t_pred reads verify col 0).
11869                //
11870                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
11871                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
11872                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
11873                // trunk hidden (the last verify column). set_len first: a p-min break may have
11874                // left one extra chain append at that slot. Partial accepts need NO fill (the
11875                // chain already covered every accepted position; round-start set_len truncates).
11876                let mut vh_seed = e.zeros(n_embd)?;
11877                e.copy_view_into(
11878                    &mut vh_seed,
11879                    0,
11880                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
11881                    n_embd,
11882                )?;
11883                if refresh {
11884                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
11885                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
11886                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
11887                    // the full stack (vx) is already resident from the verify. Replaces both the
11888                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
11889                    // (draft attention quality); exactness stays the verify's job.
11890                    scratch.set_len(e, pos)?;
11891                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
11892                    // (hidden of the last committed row before this verify batch).
11893                    let mut vxs = e.zeros(t_v * n_embd)?;
11894                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
11895                    if t_v > 1 {
11896                        e.copy_view_into(
11897                            &mut vxs,
11898                            n_embd,
11899                            &vx.slice(0..(t_v - 1) * n_embd),
11900                            (t_v - 1) * n_embd,
11901                        )?;
11902                    }
11903                    self.mtp_kv_fill(e, mtp, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
11904                } else {
11905                    scratch.set_len(e, pos + base + k_round - 1)?;
11906                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
11907                    let mut hp = e.zeros(n_embd)?;
11908                    if t_v >= 2 {
11909                        e.copy_view_into(
11910                            &mut hp,
11911                            0,
11912                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
11913                            n_embd,
11914                        )?;
11915                    } else {
11916                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
11917                    }
11918                    self.mtp_kv_fill(
11919                        e,
11920                        mtp,
11921                        &[draft[k_round - 1]],
11922                        &hp,
11923                        pos + base + k_round - 1,
11924                        &mut *scratch,
11925                        embd_dev,
11926                    )?;
11927                }
11928                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
11929                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
11930                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
11931                // col). Saves one MTP-block pass per round on top of the pairing fix.
11932                if !devacc_seeded {
11933                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
11934                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
11935                }
11936                pending = Some(bonus);
11937                if debug_spec {
11938                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
11939                }
11940            } else if !spec_replay && base + n_acc >= 1 {
11941                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
11942                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
11943                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
11944                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
11945                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
11946                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
11947                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
11948                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
11949                // accept (never compounds: the next verify recomputes true hiddens for all
11950                // committed columns).
11951                let j = base + n_acc;
11952                self.commit_verified_prefix(
11953                    e,
11954                    &mut *cache,
11955                    &snap,
11956                    ckpt.as_ref().unwrap(),
11957                    j,
11958                    devacc_seeded,
11959                    if devacc_seeded {
11960                        devacc_acc.as_ref().map(|a| (a, base, t_v))
11961                    } else {
11962                        None
11963                    },
11964                )?;
11965                let mut seed = e.zeros(n_embd)?;
11966                e.copy_view_into(
11967                    &mut seed,
11968                    0,
11969                    &vx.slice((j - 1) * n_embd..j * n_embd),
11970                    n_embd,
11971                )?;
11972                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
11973                // branch); without it the chain entries stand and only the tail truncates. Either
11974                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
11975                // (persistent mode), rope pos+j+1 (chain convention).
11976                if refresh {
11977                    scratch.set_len(e, pos)?;
11978                    let mut vxs = e.zeros(j * n_embd)?;
11979                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
11980                    if j > 1 {
11981                        e.copy_view_into(
11982                            &mut vxs,
11983                            n_embd,
11984                            &vx.slice(0..(j - 1) * n_embd),
11985                            (j - 1) * n_embd,
11986                        )?;
11987                    }
11988                    self.mtp_kv_fill(
11989                        e,
11990                        mtp,
11991                        &verify_tokens[0..j],
11992                        &vxs,
11993                        pos,
11994                        &mut *scratch,
11995                        embd_dev,
11996                    )?;
11997                } else {
11998                    scratch.set_len(e, pos + j)?;
11999                }
12000                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
12001                // bonus's predecessor (verify col j-1); no pseudo pass.
12002                if !devacc_seeded {
12003                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
12004                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
12005                }
12006                pending = Some(bonus);
12007                if debug_spec {
12008                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
12009                }
12010            } else if !spec_replay {
12011                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
12012                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
12013                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
12014                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
12015                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
12016                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
12017                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
12018                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
12019                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
12020                cache.rollback(e, &snap, 0)?;
12021                scratch.set_len(e, pos)?;
12022                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12023                pending = Some(bonus);
12024                if debug_spec {
12025                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
12026                }
12027            } else {
12028                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
12029                // this round survives, only possible before the first pending exists, ~round 0):
12030                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
12031                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
12032                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
12033                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
12034                // trunk hidden.
12035                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
12036                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
12037                if let Some(b) = pending.take() {
12038                    replay.push(b);
12039                }
12040                replay.extend_from_slice(&draft[0..n_acc]);
12041                replay.push(bonus);
12042                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
12043                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
12044                // last col exactly as before (byte-identical to the old _h_emb_dev call).
12045                let (rl_d, rx) = if self.qwen35_serving_class() {
12046                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
12047                    let mut hidden = e.uninit(replay.len() * n_embd)?;
12048                    for (row, &token) in replay.iter().enumerate() {
12049                        let (row_logits, row_hidden) =
12050                            self.spec_target_step_h(e, token, &mut *cache)?;
12051                        logits.extend_from_slice(&row_logits);
12052                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
12053                    }
12054                    (e.htod(&logits)?, hidden)
12055                } else {
12056                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
12057                };
12058                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
12059                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
12060                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
12061                last_pred = e.dtoh_u32(&preds_d)?[0];
12062                if sampled {
12063                    let lr0 = replay.len();
12064                    let lc = last_col_logits
12065                        .as_mut()
12066                        .expect("sampled: last_col_logits unset");
12067                    e.copy_view_into(
12068                        lc,
12069                        0,
12070                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
12071                        n_vocab,
12072                    )?;
12073                }
12074                let lr = replay.len();
12075                if lr >= 2 {
12076                    e.copy_view_into(
12077                        &mut h_seed_buf,
12078                        0,
12079                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
12080                        n_embd,
12081                    )?;
12082                } else {
12083                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
12084                    // last_token, whose own-row hidden fill_prev still holds.
12085                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12086                }
12087                // the bonus is COMMITTED here — it becomes the last committed row.
12088                let mut rh_last = e.zeros(n_embd)?;
12089                e.copy_view_into(
12090                    &mut rh_last,
12091                    0,
12092                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
12093                    n_embd,
12094                )?;
12095                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
12096                if debug_spec {
12097                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
12098                }
12099            }
12100            if devacc_seeded {
12101                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
12102                // consumed the old value (both slots carry the same value in every non-replay arm).
12103                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12104            }
12105            if successor_valid {
12106                let optimistic_scratch_len = successor_attempt
12107                    .as_ref()
12108                    .expect("valid controller successor disappeared")
12109                    .scratch_len;
12110                // The normal current-round commit refreshed/truncated the logical scratch tail.
12111                // Its optimistic successor row was already written physically, so restoring only
12112                // the retained logical length makes that row live for the carried round.
12113                scratch.set_len(e, optimistic_scratch_len)?;
12114            }
12115            if let Some(current) = current_opti.take() {
12116                opti_fork
12117                    .as_mut()
12118                    .ok_or("optipipe current retirement lost fork state")?
12119                    .retire(current.generation)?;
12120            }
12121            if successor_valid {
12122                let successor = successor_attempt
12123                    .take()
12124                    .expect("valid controller successor disappeared before promotion");
12125                let generation = successor.generation;
12126                opti_fork
12127                    .as_mut()
12128                    .ok_or("optipipe successor promotion lost fork state")?
12129                    .promote_successor_snapshot(&mut snap, generation);
12130                carried_opti = Some(successor);
12131            }
12132            if anatomy_on {
12133                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
12134                // only for this diagnostic so it does not disappear into the following draft's
12135                // first token readback.
12136                e.stream().synchronize()?;
12137                ph_commit += commit_started.elapsed().as_secs_f64();
12138            }
12139            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
12140            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
12141            // final position — the floor's position key reads the committed depth). Burst
12142            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
12143            // like gemma's burst arm.
12144            if adapt {
12145                let fl_now = floor_at(cache.pos);
12146                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
12147            }
12148            ph_mark(&mut ph_rest, phase_on);
12149            if let Some(p) = pipe {
12150                p.accept_end(round);
12151            }
12152            drop(pipe_accept);
12153            round += 1;
12154            // sse-cadence: this round's accepted drafts + bonus are committed (out is
12155            // append-only past step 4) — flush at round cadence.
12156            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12157        }
12158        if let Some(mut ticket) = carried_opti.take() {
12159            opti_fork
12160                .as_mut()
12161                .ok_or("optipipe tail drain lost fork state")?
12162                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
12163        }
12164        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
12165        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
12166        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
12167
12168        if spec_stats {
12169            let per_slot: Vec<String> = (0..k)
12170                .map(|j| {
12171                    if st_drafted[j] > 0 {
12172                        format!(
12173                            "{}/{}={:.3}",
12174                            st_accepted[j],
12175                            st_drafted[j],
12176                            st_accepted[j] as f64 / st_drafted[j] as f64
12177                        )
12178                    } else {
12179                        "0/0".into()
12180                    }
12181                })
12182                .collect();
12183            let acc = if total_drafted > 0 {
12184                total_accepted as f64 / total_drafted as f64
12185            } else {
12186                0.0
12187            };
12188            eprintln!(
12189                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
12190                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
12191                       tok_per_round={:.3}",
12192                per_slot.join(" "),
12193                (total_accepted + round) as f64 / round.max(1) as f64
12194            );
12195        }
12196        if constraint.is_some() {
12197            eprintln!(
12198                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
12199                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
12200                dm_clone_ns as f64 / 1e6,
12201                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
12202            );
12203        }
12204        if phase_on {
12205            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
12206            eprintln!(
12207                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
12208                ph_draft * 1e3,
12209                ph_draft / tot * 100.0,
12210                ph_verify * 1e3,
12211                ph_verify / tot * 100.0,
12212                ph_wait * 1e3,
12213                ph_wait / tot * 100.0,
12214                ph_rest * 1e3,
12215                ph_rest / tot * 100.0
12216            );
12217        }
12218        if anatomy_on {
12219            let rounds_f = round.max(1) as f64;
12220            let other = (ph_rest - ph_commit).max(0.0);
12221            eprintln!(
12222                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
12223                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
12224                ph_draft * 1e3 / rounds_f,
12225                ph_verify * 1e3 / rounds_f,
12226                ph_wait * 1e3 / rounds_f,
12227                ph_commit * 1e3 / rounds_f,
12228                other * 1e3 / rounds_f,
12229            );
12230        }
12231        let _pipe_tail = pipe.map(|p| p.primary());
12232        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
12233        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
12234        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
12235        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
12236        if let Some(slot) = sess_draft_slot.take() {
12237            *slot = Some(dctx);
12238        }
12239        let t_rounds = t_ent.elapsed();
12240        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
12241            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
12242            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
12243            // HERE, where the sampler, the session Philox counters and the penalty window are
12244            // all live and the boundary logits row still exists — that is the "make the state
12245            // available" half of the fix; the consuming burst then just emits it. `sctr` is
12246            // written to the session BELOW the draws so the advance is never lost.
12247            *next_pred_slot = Some(last_pred);
12248            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
12249            let mut stashed_pending = false;
12250            if let Some(b) = pending.take() {
12251                if !sampled {
12252                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
12253                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
12254                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
12255                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
12256                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
12257                    // OUT of `committed` (cache rows == committed); the consuming call
12258                    // prepends it once its verify commits the row. next_pred is unknowable
12259                    // without the commit pass — None; callers gate on pending_tok too.
12260                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
12261                    if let Some(slot) = sess_pending_slot.take() {
12262                        *slot = Some(b);
12263                    }
12264                    *next_pred_slot = None;
12265                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
12266                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
12267                    *last_h = Some(e.clone_dtod(&fill_prev)?);
12268                    stashed_pending = true;
12269                } else {
12270                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
12271                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
12272                    let pos_b = cache.pos;
12273                    scratch.set_len(e, pos_b)?;
12274                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
12275                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
12276                    // itself — the prediction AFTER the bonus never materialized; it would have
12277                    // been the next round's verify col 0). The commit's logits ARE that
12278                    // prediction — so they are also the row the next burst's boundary token
12279                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
12280                    *next_pred_slot = Some(if sample_boundary {
12281                        sample_boundary_token(
12282                            e,
12283                            &lg_b,
12284                            &sp,
12285                            &pen_hist,
12286                            &mut sctr,
12287                            "burst-tail-commit",
12288                        )?
12289                    } else {
12290                        argmax(&lg_b) as u32
12291                    });
12292                    self.mtp_kv_fill(e, mtp, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
12293                    *last_h = Some(hb);
12294                }
12295            } else {
12296                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
12297                *last_h = Some(e.clone_dtod(&fill_prev)?);
12298                if sample_boundary {
12299                    // No pending to commit, so the boundary row is the one `last_pred` was
12300                    // argmaxed from and the sampled path keeps it on device: the init feed's
12301                    // logits when the burst ran zero rounds, else the legacy-replay path's
12302                    // last verify column (both predict the token AFTER the last committed
12303                    // row). It is retained precisely because round 0's accept test needs it,
12304                    // so the draw costs no extra D2H of the [n_vocab] row.
12305                    match last_col_logits.as_ref() {
12306                        Some(lc) => {
12307                            *next_pred_slot = Some(sample_boundary_token_dev(
12308                                e,
12309                                lc,
12310                                n_vocab,
12311                                &sp,
12312                                &pen_hist,
12313                                &mut sctr,
12314                                "burst-tail-nopending",
12315                            )?);
12316                        }
12317                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
12318                        // burst always feeds or replays, so the row exists — but if it ever
12319                        // is, the stream takes a greedy token and SAYS so rather than
12320                        // silently regressing to the pre-lane behaviour.
12321                        None => eprintln!(
12322                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
12323                             (reason: no retained boundary logits row)"
12324                        ),
12325                    }
12326                }
12327            }
12328            *sctr_slot = sctr;
12329            *uctr_slot = uctr;
12330            committed.extend_from_slice(prompt);
12331            if let Some(cb) = carried_pending {
12332                // the consumed carry's cache row landed in round 0's verify (every pending
12333                // round commits col 0) — it joins `committed` here, in sequence order.
12334                committed.push(cb);
12335            }
12336            if stashed_pending {
12337                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
12338                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
12339                // 18446744073709551615 out of range for slice of length 0", killing the
12340                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
12341                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
12342                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
12343                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
12344                // did). So a burst that stashes a pending without emitting anything of its own —
12345                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
12346                // guard skipping every token under a tight budget — arrives here with
12347                // out.len() == 0 and stashed_pending == true.
12348                //
12349                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
12350                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
12351                // just above is already accounted. Saturating, not a min/assert: an empty `out`
12352                // here is a legitimate burst shape, not a corrupt state.
12353                let emitted = out.len().saturating_sub(1);
12354                committed.extend_from_slice(&out[..emitted]);
12355            } else {
12356                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
12357            }
12358            debug_assert_eq!(
12359                cache.pos,
12360                committed.len(),
12361                "session invariant: cache rows == committed tokens"
12362            );
12363            if setup_trace {
12364                e.stream().synchronize()?; // bound the async tail fill in the trace
12365                let t_tail = t_ent.elapsed();
12366                eprintln!(
12367                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
12368                    t_init.as_secs_f64() * 1e3,
12369                    (t_cap - t_init).as_secs_f64() * 1e3,
12370                    (t_fill - t_cap).as_secs_f64() * 1e3,
12371                    (t_rounds - t_fill).as_secs_f64() * 1e3,
12372                    (t_tail - t_rounds).as_secs_f64() * 1e3,
12373                    t_tail.as_secs_f64() * 1e3,
12374                    out.len(),
12375                    continuation
12376                );
12377            }
12378            return Ok((out, total_drafted, total_accepted));
12379        }
12380        out.truncate(max_new);
12381        Ok((out, total_drafted, total_accepted))
12382    }
12383
12384    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
12385    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
12386    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
12387    pub fn extract_dspark_anchors(
12388        &self,
12389        e: &Engine,
12390        tokens: &[u32],
12391        anchor_positions: &[usize],
12392        gamma: usize,
12393        top_k: usize,
12394        chunk: usize,
12395        temperature: f32,
12396    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
12397        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
12398            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
12399        }
12400        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
12401            return Err("DSpark anchor positions must be sorted and unique".into());
12402        }
12403        for &position in anchor_positions {
12404            if position == 0 || position + gamma >= tokens.len() {
12405                return Err(format!(
12406                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
12407                    tokens.len()
12408                )
12409                .into());
12410            }
12411        }
12412
12413        let n_vocab = self.output.out_features();
12414        let n_embd = self.cfg.n_embd as usize;
12415        let mut cache = crate::pp::new_cache(e, &self.cfg, tokens.len() + gamma + 8)?;
12416        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
12417        let embd_gpu = if spec_host_embd() {
12418            None
12419        } else {
12420            Some(
12421                self.embd_gpu
12422                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
12423            )
12424        };
12425        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
12426
12427        struct PendingRecord {
12428            position: usize,
12429            hidden: Option<Vec<f32>>,
12430            tokens: Vec<u32>,
12431            target_top_ids: Vec<Option<Vec<u32>>>,
12432            target_top_logits: Vec<Option<Vec<f32>>>,
12433            target_top_probs: Vec<Option<Vec<f32>>>,
12434            target_tail_probs: Vec<Option<f32>>,
12435        }
12436
12437        let mut pending: Vec<PendingRecord> = anchor_positions
12438            .iter()
12439            .map(|&position| PendingRecord {
12440                position,
12441                hidden: None,
12442                tokens: tokens[position..=position + gamma].to_vec(),
12443                target_top_ids: vec![None; gamma],
12444                target_top_logits: vec![None; gamma],
12445                target_top_probs: vec![None; gamma],
12446                target_tail_probs: vec![None; gamma],
12447            })
12448            .collect();
12449
12450        let mut start = 0usize;
12451        while start < tokens.len() {
12452            let end = (start + chunk).min(tokens.len());
12453            let chunk_tokens = &tokens[start..end];
12454            let (target_logits, hidden_rows) =
12455                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
12456            for record in &mut pending {
12457                let hidden_position = record.position - 1;
12458                if hidden_position >= start && hidden_position < end {
12459                    let local = hidden_position - start;
12460                    record.hidden = Some(
12461                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
12462                    );
12463                }
12464                for slot in 0..gamma {
12465                    let target_row = record.position + slot;
12466                    if target_row < start || target_row >= end {
12467                        continue;
12468                    }
12469                    let local = target_row - start;
12470                    let logits =
12471                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
12472                    let (ids, top_logits, probs, tail) =
12473                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
12474                    record.target_top_ids[slot] = Some(ids);
12475                    record.target_top_logits[slot] = Some(top_logits);
12476                    record.target_top_probs[slot] = Some(probs);
12477                    record.target_tail_probs[slot] = Some(tail);
12478                }
12479            }
12480            start = end;
12481        }
12482
12483        pending
12484            .into_iter()
12485            .map(|record| {
12486                let hidden = record
12487                    .hidden
12488                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
12489                let target_top_ids =
12490                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
12491                let target_top_logits = flatten_dspark_rows(
12492                    record.target_top_logits,
12493                    record.position,
12494                    "target logits",
12495                )?;
12496                let target_top_probs =
12497                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
12498                let target_tail_probs = record
12499                    .target_tail_probs
12500                    .into_iter()
12501                    .enumerate()
12502                    .map(|(slot, value)| {
12503                        value.ok_or_else(|| {
12504                            format!("missing DSpark tail at {} slot {slot}", record.position)
12505                        })
12506                    })
12507                    .collect::<Result<Vec<_>, _>>()?;
12508                Ok(DsparkAnchorRecord {
12509                    position: record.position,
12510                    hidden,
12511                    tokens: record.tokens,
12512                    target_top_ids,
12513                    target_top_logits,
12514                    target_top_probs,
12515                    target_tail_probs,
12516                })
12517            })
12518            .collect()
12519    }
12520
12521    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
12522    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
12523    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
12524    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
12525    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
12526    /// quant-induced head/hidden-state mismatch from text drift.
12527    ///
12528    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
12529    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
12530    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
12531    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
12532    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
12533    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
12534    ///              conditions on the corpus — deterministic and arm-comparable by design.
12535    ///
12536    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
12537    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
12538    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
12539    ///
12540    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
12541    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
12542    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
12543    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
12544    /// agreement vs this path — not usable as a training-data source).
12545    pub fn replay_acceptance(
12546        &self,
12547        e: &Engine,
12548        tokens: &[u32],
12549        k: usize,
12550        stride: usize,
12551        chunk: usize,
12552        mut hdump: Option<&mut std::fs::File>,
12553    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
12554        assert!(k >= 1 && stride >= 1 && chunk >= 2);
12555        let mtp = self
12556            .mtp
12557            .as_ref()
12558            .expect("replay_acceptance requires an MTP head");
12559        let n_vocab = self.output.out_features();
12560        let d_vocab = mtp
12561            .shared_head_head
12562            .as_ref()
12563            .unwrap_or(&self.output)
12564            .out_features();
12565        let n_embd = self.cfg.n_embd as usize;
12566        let t_total = tokens.len();
12567        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
12568        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
12569        let mut cache = crate::pp::new_cache(e, &self.cfg, t_total + k + 8)?;
12570        let mut scratch = MtpScratch::new(
12571            e,
12572            &self.cfg,
12573            t_total + k + 8,
12574            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
12575        )?;
12576        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
12577        let embd_gpu = if spec_host_embd() {
12578            None
12579        } else {
12580            Some(
12581                self.embd_gpu
12582                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
12583            )
12584        };
12585        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
12586
12587        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
12588        let mut bg: Vec<u32> = vec![0; t_total + 1];
12589        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
12590        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
12591        let mut seed_buf = e.zeros(n_embd)?;
12592        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
12593        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
12594        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
12595        let mut s = 0usize;
12596        while s < t_total {
12597            let cend = (s + chunk).min(t_total);
12598            let tc = cend - s;
12599            let ch = &tokens[s..cend];
12600            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
12601            //    the chunk's true hiddens.
12602            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
12603            for j in 0..tc {
12604                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12605            }
12606            let preds = e.dtoh_u32(&preds_d)?;
12607            for j in 0..tc {
12608                bg[s + j + 1] = preds[j];
12609            }
12610            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
12611            // checkpoint-quality metric (position j's logits score the GOLD next token).
12612            if nll_on {
12613                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
12614                if jmax > 0 {
12615                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
12616                    let rows: Vec<i32> = (0..jmax as i32).collect();
12617                    let idsd = e.htod_u32_v(&ids)?;
12618                    let rowsd = e.htod_i32(&rows)?;
12619                    let mut outd = e.zeros(jmax)?;
12620                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
12621                    for pr in e.dtoh(&outd)? {
12622                        nll_sum += -((pr.max(1e-30)) as f64).ln();
12623                        nll_cnt += 1;
12624                    }
12625                }
12626            }
12627            if let Some(f) = hdump.as_deref_mut() {
12628                use std::io::Write;
12629                let host: Vec<f32> = e.dtoh(&vx)?;
12630                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
12631                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
12632                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
12633                for v in &host[..tc * n_embd] {
12634                    let b = v.to_bits();
12635                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
12636                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
12637                }
12638                f.write_all(&bytes)?;
12639            }
12640            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
12641            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
12642            // per token saved; the forced trunk pass + hdump is all the mode needs).
12643            let chainless = stride > t_total;
12644            if chainless {
12645                e.copy_view_into(
12646                    &mut prev_last_h,
12647                    0,
12648                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
12649                    n_embd,
12650                )?;
12651                s = cend;
12652                continue;
12653            }
12654            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
12655            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
12656            let mut vxs = e.zeros(tc * n_embd)?;
12657            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
12658            if tc > 1 {
12659                e.copy_view_into(
12660                    &mut vxs,
12661                    n_embd,
12662                    &vx.slice(0..(tc - 1) * n_embd),
12663                    (tc - 1) * n_embd,
12664                )?;
12665            }
12666            scratch.set_len(e, s)?;
12667            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
12668            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
12669            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
12670            //    truncates those approximate appends before they can ever be read.
12671            let ps: Vec<usize> = (s..cend)
12672                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
12673                .collect();
12674            for &p in ps.iter().rev() {
12675                scratch.set_len(e, p)?;
12676                if p == s {
12677                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
12678                } else {
12679                    e.copy_view_into(
12680                        &mut seed_buf,
12681                        0,
12682                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
12683                        n_embd,
12684                    )?;
12685                }
12686                let mut e_tok = tokens[p];
12687                let mut d_seed = e.clone_dtod(&seed_buf)?;
12688                let mut drafts: Vec<u32> = Vec::with_capacity(k);
12689                for j in 0..k {
12690                    let (dl_d, h_nextn) = self.mtp_head_forward_dev(
12691                        e,
12692                        mtp,
12693                        e_tok,
12694                        &d_seed,
12695                        &mut scratch,
12696                        p + 1 + j,
12697                        embd_dev,
12698                        None, // acceptance-oracle walk: no grammar
12699                    )?;
12700                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
12701                    let idx = e.dtoh_u32_one(&tok_d)?;
12702                    let d = match &mtp.d2t {
12703                        Some(map) => map[idx as usize],
12704                        None => idx,
12705                    };
12706                    drafts.push(d);
12707                    e_tok = d;
12708                    d_seed = h_nextn;
12709                }
12710                // targets may live in a LATER chunk's bg — resolved after the walk.
12711                rows.push((p, drafts, Vec::new()));
12712            }
12713            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
12714            //    expect scratch.len == cend with exact rows).
12715            scratch.set_len(e, s)?;
12716            self.mtp_kv_fill(e, mtp, ch, &vxs, s, &mut scratch, embd_dev)?;
12717            e.copy_view_into(
12718                &mut prev_last_h,
12719                0,
12720                &vx.slice((tc - 1) * n_embd..tc * n_embd),
12721                n_embd,
12722            )?;
12723            s = cend;
12724        }
12725        for (p, drafts, targets) in rows.iter_mut() {
12726            for j in 0..drafts.len() {
12727                targets.push(bg[*p + 1 + j]);
12728            }
12729        }
12730        rows.sort_by_key(|r| r.0);
12731        if nll_cnt > 0 {
12732            let mean = nll_sum / nll_cnt as f64;
12733            println!(
12734                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
12735                mean.exp()
12736            );
12737        }
12738        Ok((rows, bg))
12739    }
12740}
12741
12742#[cfg(test)]
12743mod dspark_sparse_tests {
12744    use super::dspark_sparse_softmax_topk;
12745
12746    #[test]
12747    fn topk_keeps_full_softmax_mass_and_stable_ties() {
12748        let logits = [1.0f32, 3.0, 3.0, -2.0];
12749        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
12750        assert_eq!(ids, vec![1, 2]);
12751        assert_eq!(top_logits, vec![3.0, 3.0]);
12752        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
12753        let expected = 1.0 / denominator;
12754        assert!((probs[0] - expected).abs() < 1.0e-6);
12755        assert!((probs[1] - expected).abs() < 1.0e-6);
12756        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
12757        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
12758    }
12759}
12760
12761#[cfg(test)]
12762mod spec_replay_env_tests {
12763    use super::spec_replay_env_on;
12764
12765    #[test]
12766    fn replay_requires_literal_one() {
12767        assert!(!spec_replay_env_on(None));
12768        assert!(!spec_replay_env_on(Some("")));
12769        assert!(!spec_replay_env_on(Some("0")));
12770        assert!(!spec_replay_env_on(Some("true")));
12771        assert!(!spec_replay_env_on(Some("2")));
12772        assert!(spec_replay_env_on(Some("1")));
12773    }
12774}
12775
12776#[cfg(test)]
12777mod telem_tests {
12778    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
12779
12780    #[test]
12781    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
12782        let counters = SpecTelemetryCounters::default();
12783        for mask in [
12784            [true, true, true],
12785            [true, true, false],
12786            [true, false, false],
12787            [false, false, false],
12788        ] {
12789            let accepted = mask.iter().take_while(|&&value| value).count();
12790            counters.record_round(mask.len(), accepted);
12791        }
12792
12793        let snapshot = counters.snapshot();
12794        assert_eq!(
12795            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
12796            (4, 12, 6)
12797        );
12798        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
12799        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
12800        assert_eq!(snapshot.tau(), 1.5);
12801        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
12802        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
12803    }
12804
12805    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
12806    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
12807    #[test]
12808    fn delta_isolates_burst_contribution() {
12809        let mut t = SpecTelemetry::default();
12810        // "previous request": 2 rounds of k=3, accepts 3 then 1.
12811        for (kr, na) in [(3usize, 3usize), (3, 1)] {
12812            t.rounds += 1;
12813            t.drafted += kr as u64;
12814            t.accepted += na as u64;
12815            for j in 0..kr {
12816                t.pos_drafted[j] += 1;
12817            }
12818            for j in 0..na {
12819                t.pos_accepted[j] += 1;
12820            }
12821        }
12822        let before = t;
12823        // "this burst": 1 round k=3, accepts 2.
12824        t.rounds += 1;
12825        t.drafted += 3;
12826        t.accepted += 2;
12827        for j in 0..3 {
12828            t.pos_drafted[j] += 1;
12829        }
12830        for j in 0..2 {
12831            t.pos_accepted[j] += 1;
12832        }
12833        let d = t.delta_since(&before);
12834        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
12835        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
12836        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
12837        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
12838    }
12839
12840    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
12841    /// aggregation invariant.
12842    #[test]
12843    fn merge_accumulates_fieldwise() {
12844        let mut agg = SpecTelemetry::default();
12845        let mut d1 = SpecTelemetry {
12846            rounds: 2,
12847            drafted: 6,
12848            accepted: 4,
12849            ..Default::default()
12850        };
12851        d1.pos_drafted[0] = 2;
12852        d1.pos_accepted[0] = 2;
12853        let mut d2 = SpecTelemetry {
12854            rounds: 1,
12855            drafted: 3,
12856            accepted: 1,
12857            ..Default::default()
12858        };
12859        d2.pos_drafted[0] = 1;
12860        d2.pos_accepted[0] = 1;
12861        d2.pos_drafted[1] = 1;
12862        agg.merge(&d1);
12863        agg.merge(&d2);
12864        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
12865        assert_eq!(agg.pos_drafted[0], 3);
12866        assert_eq!(agg.pos_accepted[0], 3);
12867        assert_eq!(agg.pos_drafted[1], 1);
12868        assert_eq!(agg.pos_accepted[1], 0);
12869    }
12870
12871    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
12872    /// public metrics surface and must never publish a u64-wrapped garbage value.
12873    #[test]
12874    fn delta_saturates_never_wraps() {
12875        let small = SpecTelemetry {
12876            rounds: 1,
12877            drafted: 2,
12878            accepted: 1,
12879            ..Default::default()
12880        };
12881        let big = SpecTelemetry {
12882            rounds: 5,
12883            drafted: 15,
12884            accepted: 9,
12885            ..Default::default()
12886        };
12887        let d = small.delta_since(&big);
12888        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
12889    }
12890}
12891
12892#[cfg(test)]
12893mod opti_fork_tests {
12894    use super::{
12895        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
12896    };
12897
12898    #[test]
12899    fn controller_threshold_and_three_miss_breaker_are_exact() {
12900        let mut policy = OptiControllerPolicy {
12901            threshold: 0.7,
12902            consecutive_misses: 0,
12903            breaker_tripped: false,
12904        };
12905        assert!(!policy.admit(0.699_999));
12906        assert!(policy.admit(0.7));
12907        assert!(!policy.resolve(false));
12908        assert!(!policy.resolve(false));
12909        assert!(policy.resolve(false));
12910        assert!(policy.breaker_tripped);
12911        assert!(!policy.admit(1.0));
12912        assert!(
12913            !policy.resolve(true),
12914            "a resolved hit cannot re-arm a tripped request"
12915        );
12916        assert!(policy.breaker_tripped);
12917    }
12918
12919    #[test]
12920    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
12921        let mut policy = OptiControllerPolicy {
12922            threshold: 0.0,
12923            consecutive_misses: 0,
12924            breaker_tripped: false,
12925        };
12926        for _ in 0..16 {
12927            assert!(policy.admit(0.0));
12928            assert!(!policy.resolve(false));
12929        }
12930        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
12931            assert!(
12932                !policy.admit(invalid),
12933                "invalid q proxy must fail closed: {invalid}"
12934            );
12935        }
12936        assert!(!policy.breaker_tripped);
12937        assert_eq!(policy.consecutive_misses, 0);
12938    }
12939
12940    #[test]
12941    fn alternating_mode_flips_by_generation_not_round_parity() {
12942        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
12943        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
12944        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
12945        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
12946    }
12947
12948    #[test]
12949    fn live_generation_cannot_be_overwritten() {
12950        let mut tracker = OptiForkGenerationTracker::default();
12951        let g0 = tracker.reserve().unwrap();
12952        let g1 = tracker.reserve().unwrap();
12953        let err = tracker.reserve().unwrap_err().to_string();
12954        assert!(
12955            err.contains("still owns generation 0"),
12956            "unexpected error: {err}"
12957        );
12958        tracker.retire(g0).unwrap();
12959        let g2 = tracker.reserve().unwrap();
12960        assert_eq!((g2.id, g2.slot), (2, 0));
12961        tracker.retire(g1).unwrap();
12962        tracker.retire(g2).unwrap();
12963    }
12964
12965    #[test]
12966    fn teardown_rejects_a_stale_generation_tag() {
12967        let mut tracker = OptiForkGenerationTracker::default();
12968        let g0 = tracker.reserve().unwrap();
12969        tracker.retire(g0).unwrap();
12970        let err = tracker.retire(g0).unwrap_err().to_string();
12971        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
12972    }
12973}
12974
12975#[cfg(test)]
12976mod draft_graph_fallback_tests {
12977    use super::DraftGraphFallback;
12978
12979    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
12980    #[test]
12981    fn flip_is_loud_once_and_memoized_after() {
12982        let mut f = DraftGraphFallback::default();
12983        let line = f
12984            .mark_greedy("out of memory")
12985            .expect("first flip must return the warn line");
12986        assert!(
12987            line.contains("WARN"),
12988            "flip line must be warn-level: {line}"
12989        );
12990        assert!(
12991            line.contains("out of memory"),
12992            "flip line must carry the reason: {line}"
12993        );
12994        assert!(f.greedy_failed());
12995        // re-marking an already-failed graph is the memoization: quiet, still failed.
12996        assert!(f.mark_greedy("out of memory").is_none());
12997        assert!(f.greedy_failed());
12998        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
12999        assert!(!f.sampled_failed());
13000        let line_s = f
13001            .mark_sampled("capture unsupported")
13002            .expect("sampled flip is its own flip");
13003        assert!(
13004            line_s.contains("sampled"),
13005            "sampled flip names itself: {line_s}"
13006        );
13007        assert!(f.mark_sampled("capture unsupported").is_none());
13008    }
13009
13010    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
13011    /// and says so exactly when there was something to reset.
13012    #[test]
13013    fn reset_on_resume_clears_flags_and_logs_once() {
13014        let mut f = DraftGraphFallback::default();
13015        // clean session: resume is silent, nothing to reset.
13016        assert!(f.reset_on_resume().is_none());
13017        f.mark_greedy("oom").unwrap();
13018        f.mark_sampled("oom").unwrap();
13019        let note = f
13020            .reset_on_resume()
13021            .expect("a set flag must produce the reset note");
13022        assert!(
13023            note.contains("greedy+sampled"),
13024            "note names what was reset: {note}"
13025        );
13026        assert!(
13027            !f.greedy_failed() && !f.sampled_failed(),
13028            "both flags cleared"
13029        );
13030        // and the NEXT failure after a reset is a fresh flip — loud again.
13031        assert!(f.mark_greedy("oom again").is_some());
13032        let note2 = f.reset_on_resume().expect("greedy-only reset");
13033        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
13034    }
13035
13036    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
13037    /// they precede a fresh capture attempt whose own failure re-flips loudly.
13038    #[test]
13039    fn shape_change_clears_are_silent() {
13040        let mut f = DraftGraphFallback::default();
13041        f.mark_greedy("oom").unwrap();
13042        f.clear_greedy();
13043        assert!(!f.greedy_failed());
13044        f.mark_sampled("oom").unwrap();
13045        f.clear_sampled();
13046        assert!(!f.sampled_failed());
13047        // after a silent clear there is nothing left for resume to report.
13048        assert!(f.reset_on_resume().is_none());
13049    }
13050}
13051
13052/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
13053///
13054/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
13055/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
13056/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
13057/// than remembered.
13058#[cfg(test)]
13059mod sampled_graph_key_tests {
13060    use super::{SampledGraphKey, debug_t_pred0};
13061
13062    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
13063    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
13064        (k.seed, k.temp_bits, k.k)
13065    }
13066
13067    fn pure_temp_key() -> SampledGraphKey {
13068        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
13069        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
13070    }
13071
13072    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
13073    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
13074    #[test]
13075    fn vendor_filters_change_the_key() {
13076        let parked = pure_temp_key();
13077        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
13078        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
13079        assert_eq!(
13080            legacy_key(&parked),
13081            legacy_key(&vendor),
13082            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
13083        );
13084        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
13085        assert!(parked.pure_temp());
13086        assert!(!vendor.pure_temp());
13087    }
13088
13089    /// Each distribution-shaping field alone is enough to drop the parked graph.
13090    #[test]
13091    fn every_filter_field_is_keyed() {
13092        let base = pure_temp_key();
13093        for (what, other) in [
13094            (
13095                "top_k",
13096                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
13097            ),
13098            (
13099                "top_p",
13100                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
13101            ),
13102            (
13103                "min_p",
13104                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
13105            ),
13106            (
13107                "penalties",
13108                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
13109            ),
13110        ] {
13111            assert_ne!(base, other, "{what} must be part of the key");
13112            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
13113            assert_eq!(
13114                legacy_key(&base),
13115                legacy_key(&other),
13116                "{what} was invisible to the pre-fix key",
13117            );
13118        }
13119    }
13120
13121    /// The baked constants stay keyed (this half was always right — regression cover for it).
13122    #[test]
13123    fn baked_constants_stay_keyed() {
13124        let base = pure_temp_key();
13125        assert_ne!(
13126            base,
13127            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
13128            "seed"
13129        );
13130        assert_ne!(
13131            base,
13132            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
13133            "temp"
13134        );
13135        assert_ne!(
13136            base,
13137            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
13138            "k"
13139        );
13140        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
13141        assert_eq!(
13142            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
13143            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
13144        );
13145    }
13146
13147    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
13148    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
13149    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
13150    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
13151    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
13152    ///
13153    /// This test is the other end of that argument, asserted here rather than remembered in a
13154    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
13155    /// would silently become the unsound thing it is documented not to be.
13156    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
13157    #[test]
13158    fn seed_alone_still_rekeys_the_draft_graph() {
13159        let parked = pure_temp_key();
13160        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
13161        assert_ne!(
13162            parked, reseeded,
13163            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
13164             decision not to compare seed rests on exactly this",
13165        );
13166        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
13167        // because of a filter difference.
13168        assert!(parked.pure_temp() && reseeded.pure_temp());
13169    }
13170
13171    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
13172    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
13173    /// agree on the regime, so a graph that survives the drop is legal to launch.
13174    #[test]
13175    fn equal_keys_agree_on_the_regime() {
13176        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
13177        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
13178        assert_eq!(a, b);
13179        assert_eq!(a.pure_temp(), b.pure_temp());
13180        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
13181        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
13182        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
13183        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
13184    }
13185
13186    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
13187    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
13188    #[test]
13189    fn debug_print_survives_the_sampled_arm() {
13190        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
13191        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
13192        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
13193        // round 0 without a pending bonus still reports last_pred, in both arms.
13194        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
13195        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
13196        // greedy keeps the real prediction it always printed.
13197        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
13198        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
13199    }
13200}