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/// MTP-ROUTE verify graphs (`MEMRA_SPEC_VERIFY_GRAPH`), DEFAULT OFF pending its A/B.
240///
241/// The slice-4c capture already lives inside `qwen35_verify_tparallel` — it has simply
242/// never had a caller on this route ("stream rides the qwen35moe burst, graphs ride the
243/// dspark route"). The MTP spec round is the missing caller, and on this family the prize
244/// is an order larger than the dspark one the machinery was tuned against: measured with
245/// `MEMRA_SPEC_PHASE=1` on a 35B-A3B serve round, verify-ISSUE is 50-58% of the round and
246/// verify-WAIT is 0.0% — the host is never waiting for the device, it is spending its own
247/// time launching the trunk. Adding per-launch host cost (an nsys capture) inflates
248/// verify-issue to 74-78% and leaves wait at 0, which is what launch-bound looks like.
249/// Replay collapses that per-round issue into one graph launch; the body is the same
250/// kernels in the same order, so exactness is by construction and the gates arbitrate.
251pub(crate) fn spec_verify_graph_on() -> bool {
252    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
253    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() == Ok("1"))
254}
255/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
256/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
257/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
258/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
259/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
260/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
261/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
262/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
263/// 256-token run vs the serve session's thousands of rounds), and the two
264/// instruments must keep their own measured dispositions rather than share one flag.
265pub(crate) fn dspark_verify_graph_serve_on() -> bool {
266    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
267    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
268}
269/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
270/// pool's memory policy STATED instead of silently unbounded. The keyspace is
271/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
272/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
273/// on the q38 export — so the default (256) never engages there; the knob is the
274/// safety valve for a future export with a wider ladder. At the ceiling the pool
275/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
276/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
277/// cols-stashed layers inside one commit). No eviction by design: destroying a live
278/// exec graph re-opens the stale-address class the indirect tables exist to close,
279/// and the bounded keyspace makes reclaim worthless.
280pub(crate) fn dspark_vg_cap() -> usize {
281    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
282    *CAP.get_or_init(|| {
283        std::env::var("MEMRA_DSPARK_VG_MAX")
284            .ok()
285            .and_then(|v| v.parse().ok())
286            .unwrap_or(256)
287    })
288}
289/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
290/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
291/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
292/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
293/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
294/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
295/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
296/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
297/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
298/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
299/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
300/// empty partial the combine never reads, so the shared n_splits_max stride changes no
301/// bytes) and re-gated e2e by this lane's battery.
302pub(crate) fn dspark_fa_rows_on() -> bool {
303    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
304    *ON.get_or_init(|| {
305        std::env::var("MEMRA_DSPARK_FA_ROWS")
306            .map(|v| v != "0")
307            .unwrap_or(true)
308    })
309}
310
311/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
312///
313/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
314/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
315/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
316/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
317/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
318/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
319/// the flag crashed precisely the regime it exists to investigate.
320///
321/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
322/// indexing (an out-of-range pred there is a real bug and must still be loud).
323fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
324    if base == 0 {
325        return last_pred.to_string();
326    }
327    match preds.get(base - 1) {
328        Some(p) => p.to_string(),
329        // sampled: the greedy per-column argmax was never run for this round.
330        None => {
331            debug_assert!(
332                sampled,
333                "greedy spec: preds[{}] missing at base {base}",
334                base - 1
335            );
336            "n/a".to_string()
337        }
338    }
339}
340
341/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
342///
343/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
344/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
345/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
346/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
347/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
348/// not believe in — and `u * 0 < p` then accepts it unconditionally.
349///
350/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
351/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
352pub(crate) fn skey_probe() -> bool {
353    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
354    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
355}
356
357/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
358/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
359/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
360/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
361/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
362/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
363/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
364/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
365/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
366pub trait SpecConstraint {
367    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
368    /// masked argmax).
369    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
370    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
371    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
372    /// Is `tok` consumable in the CURRENT state?
373    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
374    /// Advance the state with an emitted token.
375    fn consume(&mut self, tok: u32) -> Result<(), String>;
376
377    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
378    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
379    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
380    // loose, research/constrained-full-20260803). These three methods let the engine mask the
381    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
382    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
383    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
384    // stays the correctness backstop and the emitted stream is unchanged by construction
385    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
386    // argmax; a cut slot is recomputed as the masked argmax either way).
387    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
388
389    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
390    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
391    fn draft_mask_enabled(&self) -> bool {
392        false
393    }
394    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
395    /// slot. Called once per spec round, before the first draft position.
396    fn draft_begin(&mut self) -> Result<(), String> {
397        Ok(())
398    }
399    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
400    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
401    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
402        Ok(None)
403    }
404    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
405    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
406    /// engine stops drafting; the token already pushed still goes through verify.
407    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
408        Ok(false)
409    }
410}
411
412/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
413/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
414/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
415/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
416/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
417/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
418/// verify emits the masked argmax as usual).
419fn upload_draft_mask(
420    e: &Engine,
421    c: &mut dyn SpecConstraint,
422    dst: &mut CudaSlice<u32>,
423    d2t: Option<&Vec<u32>>,
424    d_vocab: usize,
425    words: usize,
426) -> Result<bool, Box<dyn std::error::Error>> {
427    let Some(tw) = c
428        .draft_mask_words()
429        .map_err(|e2| format!("constraint: {e2}"))?
430    else {
431        return Ok(false);
432    };
433    let bit = |t: usize| -> bool {
434        let w = t >> 5;
435        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
436    };
437    let mut buf = vec![0u32; words];
438    match d2t {
439        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
440        Some(map) => {
441            for (i, &t) in map.iter().enumerate().take(d_vocab) {
442                if bit(t as usize) {
443                    buf[i >> 5] |= 1u32 << (i & 31);
444                }
445            }
446        }
447        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
448        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
449        None => {
450            let n = tw.len().min(words);
451            buf[..n].copy_from_slice(&tw[..n]);
452        }
453    }
454    if buf.iter().all(|w| *w == 0) {
455        return Ok(false);
456    }
457    e.htod_u32_into(dst, &buf)?;
458    Ok(true)
459}
460
461/// Keep the full token-embedding table in host memory and upload only the rows needed by each
462/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
463/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
464/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
465pub(crate) fn spec_host_embd() -> bool {
466    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
467    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
468}
469
470/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
471/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
472/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
473/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
474/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
475/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
476/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
477/// run-spec K=1..8 + acceptance identity arbitrate e2e).
478pub(crate) fn spec_fused_t() -> bool {
479    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
480    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
481    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
482    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
483    *F.get_or_init(|| {
484        std::env::var("MEMRA_SPEC_FUSED_T")
485            .map(|v| v != "0")
486            .unwrap_or(true)
487    })
488}
489
490/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
491/// Only call this on such buffers — the lean contract is "identical bytes by construction".
492fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
493    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
494}
495
496/// Scratch KV for the MTP block (one full-attn layer).
497///
498/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
499/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
500/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
501/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
502/// engine's "mtp_update" design). Entries come from two sources:
503///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
504///     hidden chain-approximate — the reference engine accepts the same);
505///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
506///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
507/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
508/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
509/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
510/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
511/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
512/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
513/// committed row across turns (the predecessor-pairing seed + fill anchor).
514/// Per-request sampling config for the sampled-spec serve path.
515#[derive(Clone, Copy, Debug)]
516pub struct SpecSampling {
517    pub temp: f32,
518    pub seed: u64,
519    pub top_k: i32,            // 0 = off
520    pub top_p: f32,            // 1.0 = off
521    pub min_p: f32,            // 0.0 = off
522    pub penalty_last_n: usize, // 0 = penalties off
523    pub penalty_repeat: f32,
524    pub penalty_freq: f32,
525    pub penalty_present: f32,
526}
527
528impl SpecSampling {
529    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
530    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
531    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
532    /// key their penalty arms off this.
533    pub fn pen_on(&self) -> bool {
534        self.penalty_last_n > 0
535            && (self.penalty_repeat != 1.0
536                || self.penalty_freq != 0.0
537                || self.penalty_present != 0.0)
538    }
539}
540
541/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
542/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
543/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
544/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
545/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
546/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
547/// is a distributional bug, not a style problem).
548pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
549    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
550    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
551    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
552    for _ in 0..10 {
553        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
554        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
555        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
556        c0 = n0;
557        c1 = n1;
558        c2 = n2;
559        c3 = n3;
560        k0 = k0.wrapping_add(0x9E3779B9);
561        k1 = k1.wrapping_add(0xBB67AE85);
562    }
563    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
564}
565
566/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
567/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
568pub const SPEC_TELEM_POS: usize = 8;
569
570/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
571/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
572/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
573/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
574/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
575/// in NEITHER drafted nor accepted.
576#[derive(Clone, Copy, Default, Debug)]
577pub struct SpecTelemetry {
578    /// verify rounds completed (a round-stream burst counts each of its M rounds).
579    pub rounds: u64,
580    /// tokens drafted / accepted across all rounds.
581    pub drafted: u64,
582    pub accepted: u64,
583    /// how often draft position j (0-based within a round's chain) was offered / accepted.
584    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
585    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
586    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
587    pub pos_drafted: [u64; SPEC_TELEM_POS],
588    pub pos_accepted: [u64; SPEC_TELEM_POS],
589}
590
591impl SpecTelemetry {
592    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
593    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
594    /// a wrapped counter.
595    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
596        let mut d = SpecTelemetry {
597            rounds: self.rounds.saturating_sub(prev.rounds),
598            drafted: self.drafted.saturating_sub(prev.drafted),
599            accepted: self.accepted.saturating_sub(prev.accepted),
600            ..Default::default()
601        };
602        for j in 0..SPEC_TELEM_POS {
603            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
604            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
605        }
606        d
607    }
608    /// Fieldwise `self += d` — the worker's per-model aggregation.
609    pub fn merge(&mut self, d: &SpecTelemetry) {
610        self.rounds += d.rounds;
611        self.drafted += d.drafted;
612        self.accepted += d.accepted;
613        for j in 0..SPEC_TELEM_POS {
614            self.pos_drafted[j] += d.pos_drafted[j];
615            self.pos_accepted[j] += d.pos_accepted[j];
616        }
617    }
618
619    /// Mean accepted draft-prefix length per verify round (tau).
620    pub fn tau(&self) -> f64 {
621        if self.rounds > 0 {
622            self.accepted as f64 / self.rounds as f64
623        } else {
624            0.0
625        }
626    }
627}
628
629/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
630/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
631/// launch, synchronization, allocation, or ordering dependency to the numeric path.
632struct SpecTelemetryCounters {
633    rounds: AtomicU64,
634    drafted: AtomicU64,
635    accepted: AtomicU64,
636    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
637    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
638}
639
640impl Default for SpecTelemetryCounters {
641    fn default() -> Self {
642        Self {
643            rounds: AtomicU64::new(0),
644            drafted: AtomicU64::new(0),
645            accepted: AtomicU64::new(0),
646            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
647            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
648        }
649    }
650}
651
652impl SpecTelemetryCounters {
653    fn record_round(&self, drafted: usize, accepted: usize) {
654        debug_assert!(accepted <= drafted);
655        self.rounds.fetch_add(1, Ordering::Relaxed);
656        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
657        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
658        for counter in self.pos_drafted.iter().take(drafted) {
659            counter.fetch_add(1, Ordering::Relaxed);
660        }
661        for counter in self.pos_accepted.iter().take(accepted) {
662            counter.fetch_add(1, Ordering::Relaxed);
663        }
664    }
665
666    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
667    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
668    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
669        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
670        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
671        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
672    }
673
674    fn snapshot(&self) -> SpecTelemetry {
675        SpecTelemetry {
676            rounds: self.rounds.load(Ordering::Relaxed),
677            drafted: self.drafted.load(Ordering::Relaxed),
678            accepted: self.accepted.load(Ordering::Relaxed),
679            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
680            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
681        }
682    }
683}
684
685pub struct SpecSession {
686    pub(crate) cache: Cache,
687    pub(crate) scratch: MtpScratch,
688    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
689    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
690    /// session must count them. Callers render output from this, not from their own echo.
691    pub committed: Vec<u32>,
692    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
693    pub(crate) last_h: Option<CudaSlice<f32>>,
694    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
695    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
696    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
697    pub next_pred: Option<u32>,
698    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
699    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
700    pub sctr: u32,
701    pub uctr: u32,
702    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
703    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
704    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
705    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
706    /// research/spec-serving-20260801). None before the first turn; error paths drop it
707    /// (next burst recaptures — serve retires errored sessions anyway).
708    pub(crate) draft_ctx: Option<DraftGraphCtx>,
709    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
710    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
711    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
712    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
713    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
714    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
715    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
716    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
717    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
718    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
719    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
720    pub pending_tok: Option<u32>,
721    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
722    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
723    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
724    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
725    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
726    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
727    /// accounting the loop already does — no syncs, no allocation. NOTE a
728    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
729    /// diff with [`SpecTelemetry::delta_since`] around each burst.
730    telem: SpecTelemetryCounters,
731    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
732    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
733    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
734    /// prime, result lands in `boundary_captures`.
735    pub capture_at: Option<usize>,
736    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
737    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
738    /// publication just isn't available for that request. Plural since
739    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
740    /// split (the shared-prefix class) and the stable pre-generation boundary (the
741    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
742    /// prefill tick publishes/checkpoints.
743    pub boundary_captures: Vec<SpecBoundaryCapture>,
744    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
745    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
746    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
747    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
748    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
749    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
750    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
751    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
752    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
753    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
754    /// prompt-end capture.
755    pub ckpt_at: Option<usize>,
756}
757impl SpecSession {
758    /// Context capacity of the session's caches (the server's ContextFull guard).
759    pub fn cache_max_ctx(&self) -> usize {
760        self.cache.max_ctx
761    }
762    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
763    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
764    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
765    /// the prime boundary), so no copy was taken at prime time.
766    pub fn cache_ref(&self) -> &Cache {
767        &self.cache
768    }
769    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
770    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
771    /// like the trunk KV — draft rows below the prompt end are append-only for the
772    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
773    /// committed length, never below the prime boundary, and the true-hidden refresh
774    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
775    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
776    /// prefix-addressable; the prefix cache already refuses that class end to end).
777    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
778        if self.scratch.kv.ring.is_some() {
779            return None;
780        }
781        Some((
782            &self.scratch.kv.k,
783            &self.scratch.kv.v,
784            self.scratch.kv.k_tok_bytes,
785            self.scratch.kv.v_tok_bytes,
786        ))
787    }
788    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
789    pub fn telemetry(&self) -> SpecTelemetry {
790        self.telem.snapshot()
791    }
792    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
793    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
794    /// `spec_rewind_to_checkpoint`.
795    pub fn rewind_pos(&self) -> Option<usize> {
796        self.turn_ckpt.as_ref().map(|c| c.pos)
797    }
798    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
799    pub fn rewind_is_resident(&self) -> bool {
800        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
801            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
802        })
803    }
804    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
805    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
806    /// session has never run a turn and has no prediction to hand over.
807    pub fn demote_ready(&self) -> bool {
808        self.pending_tok.is_none() && self.next_pred.is_some()
809    }
810    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
811    pub fn has_pending(&self) -> bool {
812        self.pending_tok.is_some()
813    }
814    /// Committed row count == cache rows (the session invariant), for the caller's own
815    /// `fed`-length cross-check at a handoff boundary.
816    pub fn committed_len(&self) -> usize {
817        self.committed.len()
818    }
819    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
820    /// cache + next-token prediction to the plain batched-decode path.
821    ///
822    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
823    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
824    /// tokenwise prime of the same `committed` sequence would have left it (that is the
825    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
826    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
827    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
828    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
829    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
830    /// a state indistinguishable from one the batched path produced itself: the batched tick
831    /// emits `next_pred`, feeds it into this same cache, and decodes on.
832    ///
833    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
834    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
835    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
836    /// path would silently skip a token.
837    ///
838    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
839    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
840    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
841    /// would mean an `mtp_kv_fill` over the whole committed history).
842    pub fn into_demoted(self) -> Option<(Cache, u32)> {
843        if self.pending_tok.is_some() {
844            return None;
845        }
846        let np = self.next_pred?;
847        debug_assert_eq!(
848            self.cache.pos,
849            self.committed.len(),
850            "demotion handoff: cache rows != committed tokens"
851        );
852        Some((self.cache, np))
853    }
854    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
855    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
856    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
857    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
858    pub fn reset_graph_fallback_on_resume(&mut self) {
859        if let Some(line) = self
860            .draft_ctx
861            .as_mut()
862            .and_then(|c| c.failed.reset_on_resume())
863        {
864            eprintln!("{line}");
865        }
866    }
867}
868
869/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
870///
871/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
872/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
873/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
874/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
875/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
876/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
877///
878/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
879/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
880/// position index, so it must be a real device COPY — that copy is the entire reason a spec
881/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
882/// below the boundary were written by this turn's fill and are never revisited (the per-round
883/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
884/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
885/// predecessor-pairing anchor the next prime's fill reads for its first row.
886///
887/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
888pub(crate) struct SpecCheckpoint {
889    snap: crate::cache::CacheSnapshot,
890    /// Committed length at the boundary (== cache.pos there, the session invariant).
891    pos: usize,
892    /// Pre-output_norm hidden of row `pos - 1`.
893    last_h: CudaSlice<f32>,
894}
895
896/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
897/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
898/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
899/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
900/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
901/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
902/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
903/// so the worker slices those from the live caches post-burst instead of copying at prime time.
904pub struct SpecBoundaryCapture {
905    pub snap: crate::cache::CacheSnapshot,
906    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
907    pub pos: usize,
908    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
909    pub logits: Vec<f32>,
910    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
911    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
912    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
913    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
914    pub last_h: Vec<f32>,
915}
916
917/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
918/// spec boundary capture carries for later restored-session fills. Failure is silent
919/// (`turn_ckpt` convention): the capture publishes without an anchor.
920fn capture_boundary_hidden(
921    e: &Engine,
922    h_rows: &CudaSlice<f32>,
923    pos: usize,
924    n_embd: usize,
925) -> Vec<f32> {
926    if pos == 0 || h_rows.len() < pos * n_embd {
927        return Vec::new();
928    }
929    let Ok(mut row) = e.uninit(n_embd) else {
930        return Vec::new();
931    };
932    if e.copy_view_into(
933        &mut row,
934        0,
935        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
936        n_embd,
937    )
938    .is_err()
939    {
940        return Vec::new();
941    }
942    e.dtoh(&row).unwrap_or_default()
943}
944
945/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
946/// Default ON: the token a burst emits at its own boundary is drawn from the request's
947/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
948/// every boundary) without touching greedy, which is byte-unaffected either way.
949pub fn spec_sampled_boundary_on() -> bool {
950    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
951    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
952}
953
954/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
955/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
956/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
957/// restores the pre-lane posture (each burst restarts the window from its own prompt
958/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
959/// must keep refusing penalized sampled prefix-cache restores, because the restored
960/// session's continuation burst is handed no prompt slice at all.
961pub fn spec_pen_session_on() -> bool {
962    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
963    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
964}
965
966/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
967/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
968/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
969/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
970/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
971/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
972pub fn spec_restore_republish_on() -> bool {
973    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
974    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
975}
976
977/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
978/// the argmax the pre-lane code would have emitted from the same row. This is how the
979/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
980fn spec_boundary_trace() -> bool {
981    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
982    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
983}
984
985/// llama-parity floor for the penalty window when the request does not ask for a bigger
986/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
987/// non-identity penalty, so this floor only matters to explicit small windows and to the
988/// CLI env path.
989const PEN_WINDOW_FLOOR: usize = 64;
990
991/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
992/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
993/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
994/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
995/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
996/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
997/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
998/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
999/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1000/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1001/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1002/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1003/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1004/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1005/// is a second thing to drift.
1006pub const PEN_WINDOW_MAX: usize = 8192;
1007
1008/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1009/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1010/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1011/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1012/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1013/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1014/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1015/// window through the SAME function (one definition of "the window" across both spec
1016/// routes and the gate binary's trunk-only reference arm).
1017pub fn pen_window_seed(
1018    session_committed: &[u32],
1019    burst_prompt: &[u32],
1020    penalty_last_n: usize,
1021) -> Vec<u32> {
1022    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1023    let take_prompt = burst_prompt.len().min(win);
1024    let take_sess = (win - take_prompt).min(session_committed.len());
1025    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1026    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1027    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1028    hist
1029}
1030
1031/// Draw a BOUNDARY token from the target distribution the request asked for
1032/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1033/// every burst boundary".
1034///
1035/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1036/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1037/// row after the last committed token on a continuation burst; the prefix-cache entry's
1038/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1039/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1040/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1041/// customer asked for a sampled token, so this draws one.
1042///
1043/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1044/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1045/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1046/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1047/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1048/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1049///
1050/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1051/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1052/// stream the accept walk uses — never a second, independently seeded stream (which would be
1053/// a new distributional bug: two streams from one seed correlate wherever their counters
1054/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1055/// to the cold session's own first draw from the same logits row, which is what preserves the
1056/// sampled-hit lane's per-seed hit==cold byte identity.
1057#[allow(clippy::too_many_arguments)]
1058pub fn sample_boundary_token_dev(
1059    e: &Engine,
1060    logits: &CudaSlice<f32>,
1061    n_vocab: usize,
1062    sp: &SpecSampling,
1063    pen_hist: &[u32],
1064    sctr: &mut u32,
1065    site: &str,
1066) -> Result<u32, Box<dyn std::error::Error>> {
1067    debug_assert!(
1068        sp.temp > 0.0,
1069        "boundary sampling is the sampled regime only"
1070    );
1071    // Own copy: penalize_logits mutates in place and the caller's row is live state
1072    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1073    let mut col = e.zeros(n_vocab)?;
1074    e.copy_into(&mut col, 0, logits, n_vocab)?;
1075    let pen_on = sp.penalty_last_n > 0
1076        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1077    if pen_on && !pen_hist.is_empty() {
1078        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1079        let w0 = pen_hist
1080            .len()
1081            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1082        let hist = &pen_hist[w0..];
1083        let hd = e.htod_u32_v(hist)?;
1084        e.penalize_logits(
1085            &mut col,
1086            &hd,
1087            hist.len(),
1088            sp.penalty_repeat,
1089            sp.penalty_freq,
1090            sp.penalty_present,
1091            n_vocab,
1092        )?;
1093    }
1094    let rows0 = e.htod_i32(&[0])?;
1095    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1096    e.filter_stats(
1097        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1098        sp.top_p, sp.min_p,
1099    )?;
1100    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1101    let mut perturb = e.zeros(n_vocab)?;
1102    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1103    *sctr = sctr.wrapping_add(1);
1104    let td = e.argmax_token_device(&perturb, n_vocab)?;
1105    let tok = e.dtoh_u32_one(&td)?;
1106    if spec_boundary_trace() {
1107        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1108        let raw = e.argmax_token_device(logits, n_vocab)?;
1109        let greedy = e.dtoh_u32_one(&raw)?;
1110        eprintln!(
1111            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1112             deviates={} temp={} sctr={}",
1113            (tok != greedy) as u8,
1114            sp.temp,
1115            sctr.wrapping_sub(1),
1116        );
1117    }
1118    Ok(tok)
1119}
1120
1121/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1122/// host `Vec<f32>`).
1123#[allow(clippy::too_many_arguments)]
1124pub fn sample_boundary_token(
1125    e: &Engine,
1126    logits: &[f32],
1127    sp: &SpecSampling,
1128    pen_hist: &[u32],
1129    sctr: &mut u32,
1130    site: &str,
1131) -> Result<u32, Box<dyn std::error::Error>> {
1132    let n_vocab = logits.len();
1133    let d = e.htod(logits)?;
1134    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1135}
1136
1137struct SpecPipeTraceClock {
1138    pair: usize,
1139    started: std::time::Instant,
1140}
1141
1142#[derive(Clone)]
1143struct SpecPipeTraceCtx {
1144    clock: std::sync::Arc<SpecPipeTraceClock>,
1145    round: usize,
1146    lane: usize,
1147}
1148
1149struct SpecPipeTraceMarker {
1150    trace: SpecPipeTraceCtx,
1151    phase: &'static str,
1152    edge: &'static str,
1153    slot: Option<usize>,
1154}
1155
1156unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1157    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1158    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1159    let slot = marker
1160        .slot
1161        .map(|v| v.to_string())
1162        .unwrap_or_else(|| "-".into());
1163    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1164    use std::io::Write as _;
1165    let stderr = std::io::stderr();
1166    let mut stderr = stderr.lock();
1167    let _ = writeln!(
1168        stderr,
1169        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1170         slot={slot} t_ms={t_ms:.3}",
1171        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1172    );
1173}
1174
1175fn enqueue_spec_pipe_trace_marker(
1176    stream: &cudarc::driver::CudaStream,
1177    trace: Option<&SpecPipeTraceCtx>,
1178    phase: &'static str,
1179    edge: &'static str,
1180    slot: Option<usize>,
1181) -> Result<(), Box<dyn std::error::Error>> {
1182    let Some(trace) = trace else {
1183        return Ok(());
1184    };
1185    let marker = Box::new(SpecPipeTraceMarker {
1186        trace: trace.clone(),
1187        phase,
1188        edge,
1189        slot,
1190    });
1191    let raw = Box::into_raw(marker);
1192    let result = unsafe {
1193        cudarc::driver::result::stream::launch_host_function(
1194            stream.cu_stream(),
1195            spec_pipe_trace_marker,
1196            raw.cast(),
1197        )
1198    };
1199    if let Err(err) = result {
1200        unsafe {
1201            drop(Box::from_raw(raw));
1202        }
1203        return Err(err.into());
1204    }
1205    Ok(())
1206}
1207
1208#[derive(Default)]
1209struct SpecPipeProgress {
1210    setup_done: [bool; 2],
1211    draft_done: [usize; 2],
1212    stage0_done: [usize; 2],
1213    verify_done: [usize; 2],
1214    accept_done: [usize; 2],
1215    finished: [bool; 2],
1216    aborted: bool,
1217}
1218
1219/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1220/// keeps its existing call stack and round locals; this object only orders phase entry. The
1221/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1222/// cannot be interleaved by the two host threads.
1223struct SpecPipeSync {
1224    progress: std::sync::Mutex<SpecPipeProgress>,
1225    changed: std::sync::Condvar,
1226    primary: std::sync::Mutex<()>,
1227    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1228}
1229
1230impl SpecPipeSync {
1231    fn new() -> Self {
1232        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1233        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1234            std::sync::Arc::new(SpecPipeTraceClock {
1235                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1236                started: std::time::Instant::now(),
1237            })
1238        });
1239        Self {
1240            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1241            changed: std::sync::Condvar::new(),
1242            primary: std::sync::Mutex::new(()),
1243            trace,
1244        }
1245    }
1246}
1247
1248#[derive(Clone)]
1249struct SpecPipeLane {
1250    sync: std::sync::Arc<SpecPipeSync>,
1251    lane: usize,
1252}
1253
1254impl SpecPipeLane {
1255    fn peer(&self) -> usize {
1256        1 - self.lane
1257    }
1258
1259    fn aborted() -> Box<dyn std::error::Error> {
1260        "paired speculative peer aborted".into()
1261    }
1262
1263    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1264        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1265            clock: clock.clone(),
1266            round,
1267            lane: self.lane,
1268        })
1269    }
1270
1271    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1272        let mut p = self.sync.progress.lock().unwrap();
1273        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1274            p = self.sync.changed.wait(p).unwrap();
1275        }
1276        if p.aborted {
1277            Err(Self::aborted())
1278        } else {
1279            Ok(())
1280        }
1281    }
1282
1283    fn setup_end(&self) {
1284        let mut p = self.sync.progress.lock().unwrap();
1285        p.setup_done[self.lane] = true;
1286        self.sync.changed.notify_all();
1287    }
1288
1289    fn draft_begin(
1290        &self,
1291        round: usize,
1292    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1293        let peer = self.peer();
1294        let mut p = self.sync.progress.lock().unwrap();
1295        loop {
1296            if p.aborted {
1297                return Err(Self::aborted());
1298            }
1299            let setup_ready =
1300                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1301            let prior_ready = p.accept_done[self.lane] >= round
1302                && (p.accept_done[peer] >= round || p.finished[peer]);
1303            let turn_ready = if self.lane == 0 {
1304                true
1305            } else {
1306                p.draft_done[0] > round || p.finished[0]
1307            };
1308            if setup_ready && prior_ready && turn_ready {
1309                break;
1310            }
1311            p = self.sync.changed.wait(p).unwrap();
1312        }
1313        drop(p);
1314        Ok(self.sync.primary.lock().unwrap())
1315    }
1316
1317    fn draft_end(&self, round: usize) {
1318        let mut p = self.sync.progress.lock().unwrap();
1319        p.draft_done[self.lane] = round + 1;
1320        self.sync.changed.notify_all();
1321    }
1322
1323    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1324    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1325    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1326        let peer = self.peer();
1327        let mut p = self.sync.progress.lock().unwrap();
1328        loop {
1329            if p.aborted {
1330                return Err(Self::aborted());
1331            }
1332            let ready = if self.lane == 0 {
1333                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1334            } else {
1335                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1336            };
1337            if ready {
1338                return Ok(self.lane == 0 || p.finished[peer]);
1339            }
1340            p = self.sync.changed.wait(p).unwrap();
1341        }
1342    }
1343
1344    fn stage0_end(&self, round: usize) {
1345        let mut p = self.sync.progress.lock().unwrap();
1346        p.stage0_done[self.lane] = round + 1;
1347        self.sync.changed.notify_all();
1348    }
1349
1350    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1351    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1352    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1353        let mut p = self.sync.progress.lock().unwrap();
1354        while !p.aborted
1355            && !(p.stage0_done[self.lane] > round
1356                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1357        {
1358            p = self.sync.changed.wait(p).unwrap();
1359        }
1360        if p.aborted {
1361            Err(Self::aborted())
1362        } else {
1363            Ok(())
1364        }
1365    }
1366
1367    fn verify_end(&self, round: usize) {
1368        let mut p = self.sync.progress.lock().unwrap();
1369        p.verify_done[self.lane] = round + 1;
1370        self.sync.changed.notify_all();
1371    }
1372
1373    fn accept_begin(
1374        &self,
1375        round: usize,
1376    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1377        let mut p = self.sync.progress.lock().unwrap();
1378        loop {
1379            if p.aborted {
1380                return Err(Self::aborted());
1381            }
1382            let ready = if self.lane == 0 {
1383                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1384            } else {
1385                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1386            };
1387            if ready {
1388                break;
1389            }
1390            p = self.sync.changed.wait(p).unwrap();
1391        }
1392        drop(p);
1393        Ok(self.sync.primary.lock().unwrap())
1394    }
1395
1396    fn accept_end(&self, round: usize) {
1397        let mut p = self.sync.progress.lock().unwrap();
1398        p.accept_done[self.lane] = round + 1;
1399        self.sync.changed.notify_all();
1400    }
1401
1402    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1403        self.sync.primary.lock().unwrap()
1404    }
1405
1406    fn finish(&self, failed: bool) {
1407        let mut p = self.sync.progress.lock().unwrap();
1408        p.finished[self.lane] = true;
1409        p.aborted |= failed;
1410        self.sync.changed.notify_all();
1411    }
1412}
1413
1414struct SpecPipeFinish<'a> {
1415    lane: &'a SpecPipeLane,
1416    closed: bool,
1417}
1418
1419impl<'a> SpecPipeFinish<'a> {
1420    fn new(lane: &'a SpecPipeLane) -> Self {
1421        Self {
1422            lane,
1423            closed: false,
1424        }
1425    }
1426
1427    fn close(&mut self, failed: bool) {
1428        self.lane.finish(failed);
1429        self.closed = true;
1430    }
1431}
1432
1433impl Drop for SpecPipeFinish<'_> {
1434    fn drop(&mut self) {
1435        if !self.closed {
1436            self.lane.finish(true);
1437        }
1438    }
1439}
1440
1441/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1442/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1443/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1444/// binds that context before touching the session, joins before returning, and never aliases the
1445/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1446/// session type Send.
1447struct SpecPipeSessionPtr(*mut SpecSession);
1448
1449unsafe impl Send for SpecPipeSessionPtr {}
1450
1451impl SpecPipeSessionPtr {
1452    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1453        unsafe { &mut *self.0 }
1454    }
1455}
1456
1457/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1458/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1459/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1460/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1461/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1462/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1463/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1464/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1465/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1466///
1467/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1468/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1469/// load-bearing:
1470///
1471/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1472///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1473///   This is all the key used to carry.
1474/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1475///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1476///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1477///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1478///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1479///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1480///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1481///
1482/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1483/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1484/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1485/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1486/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1487#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1488pub(crate) struct SampledGraphKey {
1489    seed: u64,
1490    temp_bits: u32,
1491    k: usize,
1492    top_k: i32,
1493    top_p_bits: u32,
1494    min_p_bits: u32,
1495    pen_on: bool,
1496}
1497
1498impl SampledGraphKey {
1499    pub(crate) fn new(
1500        seed: u64,
1501        temp: f32,
1502        k: usize,
1503        top_k: i32,
1504        top_p: f32,
1505        min_p: f32,
1506        pen_on: bool,
1507    ) -> Self {
1508        SampledGraphKey {
1509            seed,
1510            temp_bits: temp.to_bits(),
1511            k,
1512            top_k,
1513            top_p_bits: top_p.to_bits(),
1514            min_p_bits: min_p.to_bits(),
1515            pen_on,
1516        }
1517    }
1518
1519    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1520    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1521    /// the key can never drift apart (they were three separate expressions before this lane, and
1522    /// the launch site simply forgot to ask).
1523    pub(crate) fn pure_temp(&self) -> bool {
1524        self.top_k == 0
1525            && f32::from_bits(self.top_p_bits) >= 1.0
1526            && f32::from_bits(self.min_p_bits) <= 0.0
1527            && !self.pen_on
1528    }
1529}
1530
1531pub(crate) struct DraftGraphCtx {
1532    g_tok: CudaSlice<u32>,
1533    g_pos: CudaSlice<i32>,
1534    g_seed: CudaSlice<f32>,
1535    g_p: CudaSlice<f32>,
1536    g_ctr: CudaSlice<u32>,
1537    g_q: CudaSlice<f32>,
1538    g_perturb: CudaSlice<f32>,
1539    q_slots: Vec<CudaSlice<f32>>,
1540    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1541    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1542    /// per-position contents the host re-uploads before each replay (the graph-promote
1543    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1544    g_dmask: CudaSlice<u32>,
1545    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1546    graph_masked: bool,
1547    graph: Option<cudarc::driver::CudaGraph>,
1548    graph_s: Option<cudarc::driver::CudaGraph>,
1549    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1550    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1551    failed: DraftGraphFallback,
1552    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1553    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1554    s_key: Option<SampledGraphKey>,
1555    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1556    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1557    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1558    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1559    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1560    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1561    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1562    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1563    keeper: Vec<Box<dyn std::any::Any + Send>>,
1564    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1565}
1566
1567/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1568/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1569///
1570/// Three contracts:
1571/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1572///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1573///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1574///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1575///   fallback from paying a doomed capture attempt every burst).
1576/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1577///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1578///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1579///   actually set (quiet on the common clean-resume path).
1580/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1581///   capture attempt whose own failure would re-flip loudly.
1582#[derive(Default)]
1583pub(crate) struct DraftGraphFallback {
1584    greedy: bool,
1585    sampled: bool,
1586}
1587impl DraftGraphFallback {
1588    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1589        if self.greedy {
1590            return None;
1591        }
1592        self.greedy = true;
1593        Some(format!(
1594            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1595        ))
1596    }
1597    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1598        if self.sampled {
1599            return None;
1600        }
1601        self.sampled = true;
1602        Some(format!(
1603            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1604        ))
1605    }
1606    fn greedy_failed(&self) -> bool {
1607        self.greedy
1608    }
1609    fn sampled_failed(&self) -> bool {
1610        self.sampled
1611    }
1612    fn clear_greedy(&mut self) {
1613        self.greedy = false;
1614    }
1615    fn clear_sampled(&mut self) {
1616        self.sampled = false;
1617    }
1618    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1619    /// was set (so clean resumes stay quiet).
1620    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1621        if !self.greedy && !self.sampled {
1622            return None;
1623        }
1624        let which = match (self.greedy, self.sampled) {
1625            (true, true) => "greedy+sampled",
1626            (true, false) => "greedy",
1627            _ => "sampled",
1628        };
1629        self.greedy = false;
1630        self.sampled = false;
1631        Some(format!(
1632            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1633        ))
1634    }
1635}
1636
1637impl DraftGraphCtx {
1638    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1639        Ok(DraftGraphCtx {
1640            g_tok: e.alloc_u32_zeroed(1)?,
1641            g_pos: e.htod_i32(&[0])?,
1642            g_seed: e.zeros(n_embd)?,
1643            g_p: e.zeros(1)?,
1644            g_ctr: e.alloc_u32_zeroed(1)?,
1645            g_q: e.zeros(qlen)?,
1646            g_perturb: e.zeros(qlen)?,
1647            q_slots: Vec::new(),
1648            g_dmask: e.alloc_u32_zeroed(1)?,
1649            graph_masked: false,
1650            graph: None,
1651            graph_s: None,
1652            failed: DraftGraphFallback::default(),
1653            s_key: None,
1654            keeper: Vec::new(),
1655            keeper_s: Vec::new(),
1656        })
1657    }
1658}
1659
1660pub(crate) struct MtpScratch {
1661    kv: KvLayer,
1662    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1663    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1664    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1665    /// smaller host-indexed SWA ring instead.
1666    cap: usize,
1667    extra: Vec<MtpScratchPlane>,
1668}
1669
1670struct MtpScratchPlane {
1671    kv: KvLayer,
1672    cap: usize,
1673}
1674
1675fn mtp_scratch_layout(
1676    cfg: &memra_gguf::config::ModelConfig,
1677    geom: Option<&crate::hybrid::DraftGeom>,
1678) -> (usize, usize, usize, usize) {
1679    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1680    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1681    let head_dim_k = cfg.head_dim_k as usize;
1682    let head_dim_v = cfg.head_dim_v as usize;
1683    assert!(
1684        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1685        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1686    );
1687    let kv_dim_k = head_dim_k * n_head_kv;
1688    let kv_dim_v = head_dim_v * n_head_kv;
1689    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1690    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1691    let (kbb, vbb) = crate::kv_blk_bytes();
1692    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1693    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1694    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1695}
1696
1697fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
1698    assert!(head_count > 0, "MTP chain requires at least one head");
1699    step % head_count
1700}
1701
1702impl MtpScratch {
1703    fn alloc_plane(
1704        e: &Engine,
1705        cfg: &memra_gguf::config::ModelConfig,
1706        plan: &memra_gguf::model_plan::ModelPlan,
1707        cap: usize,
1708        geom: Option<&crate::hybrid::DraftGeom>,
1709    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
1710        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1711        let ring = if crate::cache::swa_ring_on()
1712            && crate::plan_backend::decode_batch_program(plan)
1713                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
1714        {
1715            let window = plan
1716                .layers
1717                .iter()
1718                .find_map(|layer| match layer.attention {
1719                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
1720                        Some(window as usize)
1721                    }
1722                    _ => None,
1723                })
1724                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
1725            Some(crate::cache::KvRing::new(
1726                crate::cache::swa_ring_rows(window, cap),
1727                window,
1728            ))
1729        } else {
1730            None
1731        };
1732        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1733        Ok(MtpScratchPlane {
1734            kv: KvLayer {
1735                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1736                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1737                kv_dim_k,
1738                kv_dim_v,
1739                k_tok_bytes,
1740                v_tok_bytes,
1741                len: 0,
1742                ring,
1743                len_d: e.htod_i32(&[0])?,
1744            },
1745            cap,
1746        })
1747    }
1748
1749    fn new(
1750        e: &Engine,
1751        cfg: &memra_gguf::config::ModelConfig,
1752        plan: &memra_gguf::model_plan::ModelPlan,
1753        cap: usize,
1754        geom: Option<&crate::hybrid::DraftGeom>,
1755    ) -> Result<Self, Box<dyn std::error::Error>> {
1756        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1757        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1758        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1759        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1760        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
1761        Ok(MtpScratch {
1762            kv: primary.kv,
1763            cap: primary.cap,
1764            extra: Vec::new(),
1765        })
1766    }
1767
1768    fn push_plane(
1769        &mut self,
1770        e: &Engine,
1771        cfg: &memra_gguf::config::ModelConfig,
1772        plan: &memra_gguf::model_plan::ModelPlan,
1773        geom: Option<&crate::hybrid::DraftGeom>,
1774    ) -> Result<(), Box<dyn std::error::Error>> {
1775        self.extra
1776            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
1777        Ok(())
1778    }
1779
1780    fn plane_count(&self) -> usize {
1781        1 + self.extra.len()
1782    }
1783
1784    fn plane(&self, index: usize) -> (&KvLayer, usize) {
1785        if index == 0 {
1786            (&self.kv, self.cap)
1787        } else {
1788            let plane = &self.extra[index - 1];
1789            (&plane.kv, plane.cap)
1790        }
1791    }
1792
1793    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
1794        if index == 0 {
1795            (&mut self.kv, self.cap)
1796        } else {
1797            let plane = &mut self.extra[index - 1];
1798            (&mut plane.kv, plane.cap)
1799        }
1800    }
1801
1802    fn set_plane_len(
1803        &mut self,
1804        e: &Engine,
1805        index: usize,
1806        n: usize,
1807    ) -> Result<(), Box<dyn std::error::Error>> {
1808        let (kv, _) = self.plane_mut(index);
1809        if kv.ring.as_ref().is_some_and(|ring| !ring.can_rewind_to(n)) {
1810            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1811        }
1812        kv.len = n;
1813        e.set_i32_one(&mut kv.len_d, n as i32)
1814    }
1815
1816    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1817    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1818    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1819    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1820        if !self.can_rewind_to(n) {
1821            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1822        }
1823        for index in 0..self.plane_count() {
1824            self.set_plane_len(e, index, n)?;
1825        }
1826        Ok(())
1827    }
1828
1829    fn can_rewind_to(&self, n: usize) -> bool {
1830        (0..self.plane_count()).all(|index| {
1831            self.plane(index)
1832                .0
1833                .ring
1834                .as_ref()
1835                .is_none_or(|ring| ring.can_rewind_to(n))
1836        })
1837    }
1838}
1839
1840/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1841/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1842/// full weight reads per round — recomputing columns the verify had already produced
1843/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1844/// to "after the first j verify columns" WITHOUT re-running the trunk:
1845/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1846///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1847///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1848///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1849///   pure-copy ring rebuild.
1850/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1851///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1852///   target: j <= t-1).
1853/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1854/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1855struct GdnStash {
1856    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1857    q_l2: CudaSlice<f32>,
1858    k_l2: CudaSlice<f32>,
1859    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1860    g_log: CudaSlice<f32>,
1861    beta: CudaSlice<f32>, // [t, num_v]
1862}
1863pub(crate) struct VerifyCkpt {
1864    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1865    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1866}
1867/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1868pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1869
1870/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1871/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1872/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1873/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1874/// layers between full-attention layers are shape-static given vt — no positions, no
1875/// t_kv, state addressed through pointer tables — so runs of them capture per
1876/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1877/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1878///
1879/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1880/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1881/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1882/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1883/// before and restored after — the graph's first real launch starts from the exact
1884/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1885/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1886/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1887pub(crate) struct DsparkVerifyGraphs {
1888    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1889    lin: Vec<usize>,
1890    lin_pos: std::collections::HashMap<usize, usize>,
1891    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1892    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1893    table_all: CudaSlice<u64>,
1894    host_table: Vec<u64>,
1895    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1896    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1897    stash_conv: Vec<CudaSlice<f32>>,
1898    stash_ssm: Vec<CudaSlice<f32>>,
1899    conv_words: usize,
1900    ssm_words: usize,
1901    /// Per-vt input/output staging (stable addresses the graphs bake).
1902    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1903    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1904    /// so the sink buffer must live (and persist) with the graphs, not with the round.
1905    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1906    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1907    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
1908    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
1909    save_conv: CudaSlice<f32>,
1910    save_ssm: CudaSlice<f32>,
1911    max_run: usize,
1912    n_embd: usize,
1913    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1914    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1915    pub(crate) round_slab: bool,
1916    // ---- slice 4c: full-verify single graph per (vt, rung) ----
1917    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
1918    fa: Vec<usize>,
1919    fa_pos: std::collections::HashMap<usize, usize>,
1920    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
1921    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
1922    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
1923    fa_table: CudaSlice<u64>,
1924    fa_host_table: Vec<u64>,
1925    t_cap: usize,
1926    /// Per-vt position staging for the captured bodies — contents refreshed per round
1927    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
1928    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
1929    /// Full-verify graphs keyed (vt, rung_end, hi).
1930    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
1931    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
1932    covered: usize,
1933    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
1934    /// full-verify capture walks all of them.
1935    walk_uniform: bool,
1936}
1937
1938struct DsparkSegGraph {
1939    graph: cudarc::driver::CudaGraph,
1940    _keeper: Vec<Box<dyn std::any::Any + Send>>,
1941}
1942
1943/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
1944/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
1945/// modes without a second copy of the math.
1946pub(crate) struct FaLayerArgs<'a> {
1947    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
1948    /// them per-z (append slot = pos, T_kv = pos + 1).
1949    pub pos_d: &'a CudaSlice<i32>,
1950    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
1951    /// arm builds/uses them (graph mode refuses that arm).
1952    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
1953    pub pos0: usize,
1954    pub seqs_append: bool,
1955    pub batch_fa_on: bool,
1956    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
1957    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
1958    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
1959    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
1960    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
1961    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
1962    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
1963    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
1964    /// for FA layers that never touch it.
1965    pub ckpt: Option<&'a mut VerifyCkpt>,
1966}
1967
1968// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
1969// no automatic trait; CUDA driver graph handles are context-scoped rather than
1970// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
1971// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
1972// single decode-stream thread.
1973unsafe impl Send for DsparkVerifyGraphs {}
1974
1975impl DsparkVerifyGraphs {
1976    /// Build for this cache's shape. None when there are no linear layers, sizes are
1977    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
1978    pub(crate) fn new(
1979        e: &Engine,
1980        cache: &Cache,
1981        t_max: usize,
1982        n_embd: usize,
1983    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
1984        let lin: Vec<usize> = (0..cache.recur.len())
1985            .filter(|&il| cache.recur[il].is_some())
1986            .collect();
1987        if lin.is_empty() || t_max < 2 {
1988            return Ok(None);
1989        }
1990        let first = cache.recur[lin[0]].as_ref().unwrap();
1991        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
1992        for &il in &lin {
1993            let rl = cache.recur[il].as_ref().unwrap();
1994            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
1995                return Ok(None);
1996            }
1997        }
1998        let n = lin.len();
1999        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2000        for (k, &il) in lin.iter().enumerate() {
2001            lin_pos.insert(il, k);
2002        }
2003        // longest run of consecutive linear layers (save-scratch sizing)
2004        let mut max_run = 1usize;
2005        let mut run = 1usize;
2006        for w in lin.windows(2) {
2007            if w[1] == w[0] + 1 {
2008                run += 1;
2009                max_run = max_run.max(run);
2010            } else {
2011                run = 1;
2012            }
2013        }
2014        let rows = t_max - 1;
2015        let mut stash_conv = Vec::with_capacity(n);
2016        let mut stash_ssm = Vec::with_capacity(n);
2017        for _ in 0..n {
2018            stash_conv.push(e.uninit(rows * conv_words)?);
2019            stash_ssm.push(e.uninit(rows * ssm_words)?);
2020        }
2021        let host_table = vec![0u64; n * 6];
2022        let table_all = e.htod_u64(&host_table)?;
2023        // slice 4c: full-attention census for the full-verify graphs.
2024        let fa: Vec<usize> = (0..cache.kv.len())
2025            .filter(|&il| cache.kv[il].is_some())
2026            .collect();
2027        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2028        for (k, &il) in fa.iter().enumerate() {
2029            fa_pos.insert(il, k);
2030        }
2031        let n_layers = cache.kv.len().max(cache.recur.len());
2032        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2033        let walk_uniform = (0..n_layers).all(|il| {
2034            cache.recur.get(il).is_some_and(|r| r.is_some())
2035                != cache.kv.get(il).is_some_and(|k| k.is_some())
2036        });
2037        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2038        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2039        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2040        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2041        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2042        let covered = (0..n_layers)
2043            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2044            .count();
2045        let t_cap = t_max;
2046        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2047        let fa_table = e.htod_u64(&fa_host_table)?;
2048        Ok(Some(Self {
2049            lin,
2050            lin_pos,
2051            table_all,
2052            host_table,
2053            stash_conv,
2054            stash_ssm,
2055            conv_words,
2056            ssm_words,
2057            stage: std::collections::HashMap::new(),
2058            tap_bufs: std::collections::HashMap::new(),
2059            graphs: std::collections::HashMap::new(),
2060            save_conv: e.uninit(n * conv_words)?,
2061            save_ssm: e.uninit(n * ssm_words)?,
2062            max_run,
2063            n_embd,
2064            round_slab: false,
2065            fa,
2066            fa_pos,
2067            fa_table,
2068            fa_host_table,
2069            t_cap,
2070            pos_stage: std::collections::HashMap::new(),
2071            full: std::collections::HashMap::new(),
2072            covered,
2073            walk_uniform,
2074        }))
2075    }
2076
2077    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2078    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2079    /// cache buffers land at new addresses; a stale table would read the wrong state).
2080    pub(crate) fn refresh_tables(
2081        &mut self,
2082        e: &Engine,
2083        cache: &Cache,
2084    ) -> Result<(), Box<dyn std::error::Error>> {
2085        use cudarc::driver::DevicePtr;
2086        {
2087            let s = &e.gpu.stream();
2088            for (k, &il) in self.lin.iter().enumerate() {
2089                let rl = cache.recur[il].as_ref().unwrap();
2090                let (pc, _g0) = rl.conv_state.device_ptr(s);
2091                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2092                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2093                let o = k * 6;
2094                self.host_table[o] = pc as u64;
2095                self.host_table[o + 1] = p0 as u64;
2096                self.host_table[o + 2] = p1 as u64;
2097                self.host_table[o + 3] = pc as u64;
2098                self.host_table[o + 4] = p1 as u64;
2099                self.host_table[o + 5] = p0 as u64;
2100            }
2101            for (k, &il) in self.fa.iter().enumerate() {
2102                let kvl = cache.kv[il].as_ref().unwrap();
2103                let (pk, _g0) = kvl.k.device_ptr(s);
2104                let (pv, _g1) = kvl.v.device_ptr(s);
2105                let o = k * 2 * self.t_cap;
2106                for z in 0..self.t_cap {
2107                    self.fa_host_table[o + 2 * z] = pk as u64;
2108                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2109                }
2110            }
2111        }
2112        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2113        if !self.fa_host_table.is_empty() {
2114            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2115        }
2116        Ok(())
2117    }
2118
2119    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2120    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2121    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2122    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2123    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2124    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2125    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2126    /// captured graph is bit-identical for every round the rung covers.
2127    #[allow(clippy::too_many_arguments)]
2128    pub(crate) fn full_rung(
2129        &self,
2130        model: &crate::hybrid::HybridModel,
2131        cache: &Cache,
2132        lo: usize,
2133        hi: usize,
2134        t: usize,
2135        seqs_arms_on: bool,
2136    ) -> Option<usize> {
2137        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2138            static ONCE: std::sync::Once = std::sync::Once::new();
2139            let len0 = self
2140                .fa
2141                .first()
2142                .and_then(|&il| cache.kv[il].as_ref())
2143                .map(|k| k.len);
2144            ONCE.call_once(|| {
2145                eprintln!(
2146                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2147                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2148                    self.lin.len(), self.fa.len(), self.t_cap, len0
2149                );
2150            });
2151        }
2152        if !self.walk_uniform
2153            || !seqs_arms_on
2154            || !dspark_fa_rows_on()
2155            || t < 2
2156            || lo != 0
2157            || hi > self.covered
2158            || t > self.t_cap
2159            || self.fa.is_empty()
2160        {
2161            return None;
2162        }
2163        let cfg = &model.cfg;
2164        let head_dim_global = cfg.head_dim_k as usize;
2165        let nkv = cfg.n_head_kv as usize;
2166        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2167        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2168        // projection stride (the body's guard, hoisted so ineligible models fall back
2169        // instead of refusing mid-capture).
2170        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2171        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2172        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2173            return None;
2174        }
2175        let len0 = kvl0.len;
2176        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2177        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2178            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2179            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2180        {
2181            return None;
2182        }
2183        let rung = t_kv_last.next_power_of_two().max(256);
2184        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2185            return None;
2186        }
2187        Some(rung)
2188    }
2189
2190    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2191    /// the residual + refresh the per-vt position staging, capture on first encounter
2192    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2193    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2194    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2195    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2196    #[allow(clippy::too_many_arguments)]
2197    pub(crate) fn run_full(
2198        &mut self,
2199        model: &crate::hybrid::HybridModel,
2200        e: &Engine,
2201        lo: usize,
2202        hi: usize,
2203        x: &CudaSlice<f32>,
2204        t: usize,
2205        pos0: usize,
2206        rung: usize,
2207        cache: &mut Cache,
2208    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2209        let n_embd = self.n_embd;
2210        if !self.stage.contains_key(&t) {
2211            let xin = e.uninit(t * n_embd)?;
2212            let xout = e.uninit(t * n_embd)?;
2213            self.stage.insert(t, (xin, xout));
2214        }
2215        if !self.pos_stage.contains_key(&t) {
2216            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2217        }
2218        // Per-round refresh: position contents + input staging (both addresses are baked
2219        // by the captured bodies; only their CONTENTS change round to round).
2220        {
2221            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2222            let pb = self.pos_stage.get_mut(&t).unwrap();
2223            e.htod_i32_into(pb, &pos_host)?;
2224            let (xin, _) = self.stage.get_mut(&t).unwrap();
2225            e.copy_into(xin, 0, x, t * n_embd)?;
2226        }
2227        let key = (t, rung, hi);
2228        if !self.full.contains_key(&key) {
2229            // The warmups EXECUTE the whole walk on live state — save every linear
2230            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2231            // graph mode never bumps host lens and the appends write this round's own
2232            // slots).
2233            for (k, &il) in self.lin.iter().enumerate() {
2234                let rl = cache.recur[il].as_ref().unwrap();
2235                e.copy_into(
2236                    &mut self.save_conv,
2237                    k * self.conv_words,
2238                    &rl.conv_state,
2239                    self.conv_words,
2240                )?;
2241                e.copy_into(
2242                    &mut self.save_ssm,
2243                    k * self.ssm_words,
2244                    &rl.ssm_state,
2245                    self.ssm_words,
2246                )?;
2247            }
2248            let (graph, keeper) = {
2249                let table_all = &self.table_all;
2250                let lin_pos = &self.lin_pos;
2251                let fa_pos = &self.fa_pos;
2252                let fa_table = &self.fa_table;
2253                let t_cap = self.t_cap;
2254                let stash_conv = &mut self.stash_conv;
2255                let stash_ssm = &mut self.stash_ssm;
2256                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2257                let (xin, xout) = self
2258                    .stage
2259                    .get_mut(&t)
2260                    .map(|(a, b)| (&*a, b))
2261                    .expect("stage bucket created above");
2262                let cache_ref: &mut Cache = cache;
2263                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2264                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2265                } else {
2266                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2267                };
2268                e.capture_graph_retained_flags(iflag, move |e| {
2269                    let mut xc: Option<CudaSlice<f32>> = None;
2270                    for il in lo..hi {
2271                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2272                        let nx = if let Some(&k) = lin_pos.get(&il) {
2273                            model.qwen35_tparallel_linear_layer(
2274                                e,
2275                                il,
2276                                xr,
2277                                t,
2278                                cache_ref,
2279                                None,
2280                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2281                                Some((table_all, k * 6)),
2282                            )?
2283                        } else if let Some(&kf) = fa_pos.get(&il) {
2284                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2285                            model.qwen35_tparallel_fa_layer(
2286                                e,
2287                                il,
2288                                xr,
2289                                t,
2290                                cache_ref,
2291                                FaLayerArgs {
2292                                    pos_d,
2293                                    pos_rows: &mut no_rows,
2294                                    pos0,
2295                                    seqs_append: true,
2296                                    batch_fa_on: true,
2297                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2298                                    stream: None,
2299                                    ckpt: None,
2300                                },
2301                            )?
2302                        } else {
2303                            return Err(format!(
2304                                "run_full: layer {il} is neither linear nor full-attention"
2305                            )
2306                            .into());
2307                        };
2308                        xc = Some(nx);
2309                    }
2310                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2311                    Ok(())
2312                })?
2313            };
2314            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2315            // is odd -> 3 runs = net one swap), then restore the device state the
2316            // warmups consumed (walk scope only — layers past hi never executed). The
2317            // launch below then behaves exactly like one run.
2318            if t % 2 == 1 {
2319                for &il in &self.lin {
2320                    if il < lo || il >= hi {
2321                        continue;
2322                    }
2323                    let rl = cache.recur[il].as_mut().unwrap();
2324                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2325                }
2326            }
2327            for (k, &il) in self.lin.iter().enumerate() {
2328                if il < lo || il >= hi {
2329                    continue;
2330                }
2331                let rl = cache.recur[il].as_mut().unwrap();
2332                let (cw, sw) = (self.conv_words, self.ssm_words);
2333                {
2334                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2335                    let win = sv.slice(k * cw..(k + 1) * cw);
2336                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2337                }
2338                {
2339                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2340                    let win = sv.slice(k * sw..(k + 1) * sw);
2341                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2342                }
2343            }
2344            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2345                if let Ok(c) = crate::graph_update::node_census(&graph) {
2346                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2347                }
2348            }
2349            self.full.insert(
2350                key,
2351                DsparkSegGraph {
2352                    graph,
2353                    _keeper: keeper,
2354                },
2355            );
2356        }
2357        self.full[&key].graph.launch()?;
2358        // Host bookkeeping for the replayed body (captured host code does not re-run):
2359        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2360        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2361        // head layer's kv) that the walk never touches.
2362        if t % 2 == 1 {
2363            for &il in &self.lin {
2364                if il < lo || il >= hi {
2365                    continue;
2366                }
2367                let rl = cache.recur[il].as_mut().unwrap();
2368                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2369            }
2370        }
2371        for &il in &self.fa {
2372            if il < lo || il >= hi {
2373                continue;
2374            }
2375            cache.kv[il].as_mut().unwrap().len += t;
2376        }
2377        let (_, xout) = self.stage.get(&t).unwrap();
2378        let mut out = e.uninit(t * n_embd)?;
2379        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2380        Ok(out)
2381    }
2382
2383    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2384    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2385    /// bracketed by a segment state save/restore), launch, then apply the host parity
2386    /// bookkeeping the captured body would have done. Returns the fresh residual.
2387    #[allow(clippy::too_many_arguments)]
2388    fn run_segment(
2389        &mut self,
2390        model: &crate::hybrid::HybridModel,
2391        e: &Engine,
2392        start: usize,
2393        end: usize,
2394        x: &CudaSlice<f32>,
2395        t: usize,
2396        cache: &mut Cache,
2397    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2398        let n_embd = self.n_embd;
2399        debug_assert!(end - start <= self.max_run);
2400        if !self.stage.contains_key(&t) {
2401            let xin = e.uninit(t * n_embd)?;
2402            let xout = e.uninit(t * n_embd)?;
2403            self.stage.insert(t, (xin, xout));
2404        }
2405        // Stage the residual at the bucket's baked input address.
2406        {
2407            let (xin, _) = self.stage.get_mut(&t).unwrap();
2408            e.copy_into(xin, 0, x, t * n_embd)?;
2409        }
2410        let key = (start, t);
2411        if !self.graphs.contains_key(&key) {
2412            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2413            // ssm of every segment layer first, restore after, so the graph's first real
2414            // launch starts from the exact pre-round state (bytes gated e2e).
2415            for (k, il) in (start..end).enumerate() {
2416                let rl = cache.recur[il].as_ref().unwrap();
2417                e.copy_into(
2418                    &mut self.save_conv,
2419                    k * self.conv_words,
2420                    &rl.conv_state,
2421                    self.conv_words,
2422                )?;
2423                e.copy_into(
2424                    &mut self.save_ssm,
2425                    k * self.ssm_words,
2426                    &rl.ssm_state,
2427                    self.ssm_words,
2428                )?;
2429            }
2430            let (graph, keeper) = {
2431                let table_all = &self.table_all;
2432                let lin_pos = &self.lin_pos;
2433                let stash_conv = &mut self.stash_conv;
2434                let stash_ssm = &mut self.stash_ssm;
2435                let (xin, xout) = self
2436                    .stage
2437                    .get_mut(&t)
2438                    .map(|(a, b)| (&*a, b))
2439                    .expect("stage bucket created above");
2440                let cache_ref: &mut Cache = cache;
2441                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2442                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2443                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2444                // = ~0.41 ms/round, most of the eager-launch savings. The captured
2445                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2446                // (every transient drops inside the capture region — the generic
2447                // capture path's census precedent, 1589/1589), so AUTO_FREE has
2448                // nothing to reclaim and the graph is legal to instantiate without
2449                // it; PRIORITY is the flag the gemma slotted door ships for exactly
2450                // this reason (both alternatives drop the scan; UPLOAD via
2451                // cuGraphInstantiateWithFlags is WithParams-only and refused).
2452                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2453                // the node census at capture (the ALLOC==FREE receipt).
2454                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2455                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2456                } else {
2457                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2458                };
2459                e.capture_graph_retained_flags(iflag, move |e| {
2460                    let mut xc: Option<CudaSlice<f32>> = None;
2461                    for il in start..end {
2462                        let k = lin_pos[&il];
2463                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2464                        let nx = model.qwen35_tparallel_linear_layer(
2465                            e,
2466                            il,
2467                            xr,
2468                            t,
2469                            cache_ref,
2470                            None,
2471                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
2472                            Some((table_all, k * 6)),
2473                        )?;
2474                        xc = Some(nx);
2475                    }
2476                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2477                    Ok(())
2478                })?
2479            };
2480            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2481            // is odd -> 3 runs = net one swap), then restore the device state the
2482            // warmups consumed. The launch below then behaves exactly like one run.
2483            if t % 2 == 1 {
2484                for il in start..end {
2485                    let rl = cache.recur[il].as_mut().unwrap();
2486                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2487                }
2488            }
2489            for (k, il) in (start..end).enumerate() {
2490                let rl = cache.recur[il].as_mut().unwrap();
2491                let (cw, sw) = (self.conv_words, self.ssm_words);
2492                {
2493                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2494                    let win = sv.slice(k * cw..(k + 1) * cw);
2495                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2496                }
2497                {
2498                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2499                    let win = sv.slice(k * sw..(k + 1) * sw);
2500                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2501                }
2502            }
2503            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2504                if let Ok(c) = crate::graph_update::node_census(&graph) {
2505                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
2506                }
2507            }
2508            self.graphs.insert(
2509                key,
2510                DsparkSegGraph {
2511                    graph,
2512                    _keeper: keeper,
2513                },
2514            );
2515        }
2516        self.graphs[&key].graph.launch()?;
2517        // Host parity bookkeeping for the replayed body (the captured host swaps do not
2518        // re-run at replay).
2519        if t % 2 == 1 {
2520            for il in start..end {
2521                let rl = cache.recur[il].as_mut().unwrap();
2522                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2523            }
2524        }
2525        let (_, xout) = self.stage.get(&t).unwrap();
2526        let mut out = e.uninit(t * n_embd)?;
2527        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2528        Ok(out)
2529    }
2530
2531    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
2532    fn can_capture(&self) -> bool {
2533        self.graphs.len() + self.full.len() < dspark_vg_cap()
2534    }
2535
2536    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
2537    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
2538    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
2539    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
2540    /// refusal would stash some layers in the ctx slabs and others in the round's cols
2541    /// while one commit reads only one of them.
2542    pub(crate) fn segments_ready(
2543        &self,
2544        model: &crate::hybrid::HybridModel,
2545        lo: usize,
2546        hi: usize,
2547        t: usize,
2548    ) -> bool {
2549        if self.can_capture() {
2550            return true;
2551        }
2552        let mut il = lo;
2553        while il < hi {
2554            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2555                let start = il;
2556                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2557                    il += 1;
2558                }
2559                if !self.graphs.contains_key(&(start, t)) {
2560                    return false;
2561                }
2562            } else {
2563                il += 1;
2564            }
2565        }
2566        true
2567    }
2568
2569    /// Widest verify window this pool was built for. A caller whose round exceeds it must
2570    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
2571    /// past them is a panic rather than a refusal.
2572    pub(crate) fn t_capacity(&self) -> usize {
2573        self.t_cap
2574    }
2575
2576    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
2577    /// `row` (0-based) of layer `il`. None for non-linear layers.
2578    pub(crate) fn slab_row(
2579        &self,
2580        e: &Engine,
2581        il: usize,
2582        row: usize,
2583    ) -> Option<(u64, u64, usize, usize)> {
2584        use cudarc::driver::DevicePtr;
2585        let k = *self.lin_pos.get(&il)?;
2586        let s = &e.gpu.stream();
2587        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
2588        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
2589        Some((
2590            pc as u64 + (row * self.conv_words * 4) as u64,
2591            ps as u64 + (row * self.ssm_words * 4) as u64,
2592            self.conv_words,
2593            self.ssm_words,
2594        ))
2595    }
2596}
2597
2598impl VerifyCkpt {
2599    fn new(n_layer: usize) -> Self {
2600        VerifyCkpt {
2601            gdn: (0..n_layer).map(|_| None).collect(),
2602            cols: (0..n_layer).map(|_| None).collect(),
2603        }
2604    }
2605}
2606
2607/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
2608/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
2609/// a logical round number.
2610struct VerifyBoundaryTicket {
2611    rt: &'static crate::pp::PpNRt,
2612    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2613    slot: usize,
2614    pos0: usize,
2615    t: usize,
2616    payload: usize,
2617    n_st: usize,
2618    pipelined: bool,
2619    pp_anatomy: bool,
2620    pp_started: std::time::Instant,
2621    reverse_ms: f64,
2622    stage0_ms: f64,
2623    tx_ms: f64,
2624    trace: Option<SpecPipeTraceCtx>,
2625}
2626
2627/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
2628/// increment-2 controller can also be armed by the server's fresh-process research door.
2629#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2630pub enum OptiForkGateMode {
2631    Disabled,
2632    Hit,
2633    Miss,
2634    Alternate,
2635    Abort,
2636    Controller,
2637}
2638
2639static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
2640static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
2641    std::sync::atomic::AtomicU32::new(0);
2642static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2643static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2644static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2645static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2646static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2647static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2648static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2649static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2650static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2651static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2652    std::sync::atomic::AtomicU64::new(0);
2653static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2654    std::sync::atomic::AtomicU64::new(0);
2655static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2656
2657impl OptiForkGateMode {
2658    fn code(self) -> u8 {
2659        match self {
2660            Self::Disabled => 0,
2661            Self::Hit => 1,
2662            Self::Miss => 2,
2663            Self::Alternate => 3,
2664            Self::Abort => 4,
2665            Self::Controller => 5,
2666        }
2667    }
2668
2669    fn configured() -> Self {
2670        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2671            1 => Self::Hit,
2672            2 => Self::Miss,
2673            3 => Self::Alternate,
2674            4 => Self::Abort,
2675            5 => Self::Controller,
2676            _ => Self::Disabled,
2677        }
2678    }
2679
2680    fn action(self, generation: u64) -> OptiForkAction {
2681        match self {
2682            Self::Hit => OptiForkAction::Hit,
2683            Self::Miss => OptiForkAction::Miss,
2684            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2685            Self::Alternate => OptiForkAction::Miss,
2686            Self::Abort => OptiForkAction::Abort,
2687            Self::Disabled | Self::Controller => {
2688                unreachable!("non-forced mode cannot choose a forced fork action")
2689            }
2690        }
2691    }
2692
2693    fn is_forced(self) -> bool {
2694        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2695    }
2696}
2697
2698/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2699pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2700    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2701}
2702
2703/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2704/// two-token draft-probability product. Serving can call this only through its explicit
2705/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2706pub fn set_optipipe_controller_threshold(threshold: f32) {
2707    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2708    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2709    set_optipipe_gate_mode(OptiForkGateMode::Controller);
2710}
2711
2712#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2713pub struct OptiForkGateStats {
2714    pub attempts: u64,
2715    pub hits: u64,
2716    pub misses: u64,
2717    pub abort_drains: u64,
2718    pub refusals: u64,
2719    pub gate_checks: u64,
2720    pub gate_admits: u64,
2721    pub gate_rejects: u64,
2722    pub reconciles: u64,
2723    pub wasted_draft_tokens: u64,
2724    pub shadow_draft_tokens: u64,
2725    pub breaker_trips: u64,
2726}
2727
2728#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2729pub struct OptiForkStateIdentity {
2730    pub trunk_kv_bytes: usize,
2731    pub recurrent_bytes: usize,
2732    pub scratch_kv_bytes: usize,
2733    pub hidden_bytes: usize,
2734}
2735
2736pub fn reset_optipipe_gate_stats() {
2737    for counter in [
2738        &OPTI_FORK_ATTEMPTS,
2739        &OPTI_FORK_HITS,
2740        &OPTI_FORK_MISSES,
2741        &OPTI_FORK_ABORT_DRAINS,
2742        &OPTI_FORK_REFUSALS,
2743        &OPTI_GATE_CHECKS,
2744        &OPTI_GATE_ADMITS,
2745        &OPTI_GATE_REJECTS,
2746        &OPTI_RECONCILES,
2747        &OPTI_WASTED_DRAFT_TOKENS,
2748        &OPTI_SHADOW_DRAFT_TOKENS,
2749        &OPTI_BREAKER_TRIPS,
2750    ] {
2751        counter.store(0, std::sync::atomic::Ordering::Relaxed);
2752    }
2753}
2754
2755pub fn optipipe_gate_stats() -> OptiForkGateStats {
2756    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2757    OptiForkGateStats {
2758        attempts: load(&OPTI_FORK_ATTEMPTS),
2759        hits: load(&OPTI_FORK_HITS),
2760        misses: load(&OPTI_FORK_MISSES),
2761        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2762        refusals: load(&OPTI_FORK_REFUSALS),
2763        gate_checks: load(&OPTI_GATE_CHECKS),
2764        gate_admits: load(&OPTI_GATE_ADMITS),
2765        gate_rejects: load(&OPTI_GATE_REJECTS),
2766        reconciles: load(&OPTI_RECONCILES),
2767        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2768        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2769        breaker_trips: load(&OPTI_BREAKER_TRIPS),
2770    }
2771}
2772
2773#[derive(Clone, Copy, Debug)]
2774struct OptiControllerPolicy {
2775    threshold: f32,
2776    consecutive_misses: u8,
2777    breaker_tripped: bool,
2778}
2779
2780impl OptiControllerPolicy {
2781    fn configured() -> Self {
2782        Self {
2783            threshold: f32::from_bits(
2784                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2785            ),
2786            consecutive_misses: 0,
2787            breaker_tripped: false,
2788        }
2789    }
2790
2791    fn admit(&self, q_proxy: f32) -> bool {
2792        q_proxy.is_finite()
2793            && (0.0..=1.0).contains(&q_proxy)
2794            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2795    }
2796
2797    /// Returns true exactly when this resolution newly trips the three-miss breaker.
2798    fn resolve(&mut self, hit: bool) -> bool {
2799        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2800        // every optimistic opportunity, so the safety breaker is measured separately and must
2801        // not silently turn this arm into "three attempts then serial".
2802        if self.threshold == 0.0 {
2803            self.consecutive_misses = 0;
2804            return false;
2805        }
2806        if hit {
2807            self.consecutive_misses = 0;
2808            return false;
2809        }
2810        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2811        if !self.breaker_tripped && self.consecutive_misses >= 3 {
2812            self.breaker_tripped = true;
2813            return true;
2814        }
2815        false
2816    }
2817}
2818
2819#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2820enum OptiForkAction {
2821    Hit,
2822    Miss,
2823    Abort,
2824}
2825
2826#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2827struct OptiForkGeneration {
2828    id: u64,
2829    slot: usize,
2830}
2831
2832#[derive(Default)]
2833struct OptiForkGenerationTracker {
2834    next: u64,
2835    live: [Option<u64>; 2],
2836}
2837
2838impl OptiForkGenerationTracker {
2839    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2840        let generation = OptiForkGeneration {
2841            id: self.next,
2842            slot: (self.next & 1) as usize,
2843        };
2844        if let Some(live) = self.live[generation.slot] {
2845            return Err(format!(
2846                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2847                generation.slot,
2848            )
2849            .into());
2850        }
2851        self.next += 1;
2852        self.live[generation.slot] = Some(generation.id);
2853        Ok(generation)
2854    }
2855
2856    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2857        match self.live[generation.slot] {
2858            Some(id) if id == generation.id => {
2859                self.live[generation.slot] = None;
2860                Ok(())
2861            }
2862            other => Err(format!(
2863                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2864                generation.id, generation.slot,
2865            )
2866            .into()),
2867        }
2868    }
2869}
2870
2871struct OptiForkSeedGeneration {
2872    h_seed: CudaSlice<f32>,
2873    fill_prev: CudaSlice<f32>,
2874    scratch_len: usize,
2875}
2876
2877/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2878/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2879/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2880/// device ownership.
2881fn opti_snapshot_stage_owned(
2882    e: &Engine,
2883    cache: &Cache,
2884    rt: &'static crate::pp::PpNRt,
2885    fence: &[usize],
2886) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2887    let n = cache.kv.len();
2888    let mut snapshot = crate::cache::CacheSnapshot {
2889        kv_len: vec![None; n],
2890        tp_kv_len: vec![None; n],
2891        conv: (0..n).map(|_| None).collect(),
2892        ssm: (0..n).map(|_| None).collect(),
2893        pos: cache.pos,
2894    };
2895    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
2896    Ok(snapshot)
2897}
2898
2899fn opti_snapshot_stage_owned_into(
2900    e: &Engine,
2901    cache: &Cache,
2902    rt: &'static crate::pp::PpNRt,
2903    fence: &[usize],
2904    snapshot: &mut crate::cache::CacheSnapshot,
2905) -> Result<(), Box<dyn std::error::Error>> {
2906    if fence.len() != rt.n_stages() + 1
2907        || snapshot.kv_len.len() != cache.kv.len()
2908        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
2909    {
2910        return Err("optipipe stage-owned snapshot shape mismatch".into());
2911    }
2912    for stage in 0..rt.n_stages() {
2913        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
2914    }
2915    snapshot.pos = cache.pos;
2916    Ok(())
2917}
2918
2919/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
2920/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
2921/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
2922/// either point would capture one side of the fork at the wrong generation.
2923fn opti_snapshot_one_stage_owned_into(
2924    e: &Engine,
2925    cache: &Cache,
2926    rt: &'static crate::pp::PpNRt,
2927    fence: &[usize],
2928    stage: usize,
2929    snapshot: &mut crate::cache::CacheSnapshot,
2930) -> Result<(), Box<dyn std::error::Error>> {
2931    if fence.len() != rt.n_stages() + 1
2932        || snapshot.kv_len.len() != cache.kv.len()
2933        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
2934        || stage >= rt.n_stages()
2935    {
2936        return Err("optipipe single-stage snapshot shape mismatch".into());
2937    }
2938    let _scope = rt.enter(stage);
2939    let owner = rt.engine(stage, e);
2940    for il in fence[stage]..fence[stage + 1] {
2941        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
2942        snapshot.tp_kv_len[il] = cache.tp_kv[il]
2943            .as_ref()
2944            .map(crate::tp::ResidentTpKvCache::committed_len);
2945        match &cache.recur[il] {
2946            Some(recur) => {
2947                match snapshot.conv[il].as_mut() {
2948                    Some(dst) => {
2949                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
2950                    }
2951                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
2952                }
2953                match snapshot.ssm[il].as_mut() {
2954                    Some(dst) => {
2955                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
2956                    }
2957                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
2958                }
2959            }
2960            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
2961                return Err(
2962                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
2963                );
2964            }
2965            None => {}
2966        }
2967    }
2968    snapshot.pos = cache.pos;
2969    Ok(())
2970}
2971
2972/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
2973/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
2974/// resolve, so the reconcile tables and conditional restores are stage-local.
2975struct OptiForkState {
2976    mode: OptiForkGateMode,
2977    controller: Option<OptiControllerPolicy>,
2978    generations: OptiForkGenerationTracker,
2979    active_snapshot_slot: usize,
2980    alternate_snapshot: crate::cache::CacheSnapshot,
2981    seeds: [OptiForkSeedGeneration; 2],
2982    rt: &'static crate::pp::PpNRt,
2983    fence: [usize; 3],
2984    split: usize,
2985    len_ptrs: CudaSlice<u64>,
2986    saved_lens: CudaSlice<i32>,
2987    forced_acc: CudaSlice<u32>,
2988    valid: CudaSlice<u32>,
2989    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2990    logical_payload_bytes: [usize; 2],
2991}
2992
2993struct OptiForkTicket {
2994    generation: OptiForkGeneration,
2995    boundary: Option<VerifyBoundaryTicket>,
2996    drain: std::sync::Arc<cudarc::driver::CudaStream>,
2997    settled: bool,
2998}
2999
3000struct OptiControllerTicket {
3001    generation: OptiForkGeneration,
3002    boundary: Option<VerifyBoundaryTicket>,
3003    ckpt: Option<VerifyCkpt>,
3004    verify_tokens: [u32; 2],
3005    draft_prob: f32,
3006    eager_seed: Option<CudaSlice<f32>>,
3007    q_proxy: f32,
3008    scratch_len: usize,
3009    issued_at: std::time::Instant,
3010    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3011    settled: bool,
3012}
3013
3014struct OptiControllerPrepared {
3015    verify_tokens: [u32; 2],
3016    draft_prob: f32,
3017    eager_seed: Option<CudaSlice<f32>>,
3018    q_proxy: f32,
3019    scratch_len: usize,
3020}
3021
3022impl OptiControllerTicket {
3023    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3024        self.boundary
3025            .take()
3026            .expect("controller boundary ticket already consumed")
3027    }
3028
3029    fn take_ckpt(&mut self) -> VerifyCkpt {
3030        self.ckpt
3031            .take()
3032            .expect("controller verify checkpoint already consumed")
3033    }
3034
3035    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3036        self.eager_seed.take()
3037    }
3038
3039    fn settle(&mut self) {
3040        self.settled = true;
3041    }
3042}
3043
3044impl Drop for OptiControllerTicket {
3045    fn drop(&mut self) {
3046        if !self.settled {
3047            let _ = self.drain.synchronize();
3048            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3049        }
3050    }
3051}
3052
3053impl OptiForkTicket {
3054    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3055        self.boundary
3056            .take()
3057            .expect("fork ticket boundary already consumed")
3058    }
3059
3060    fn settle(&mut self) {
3061        self.settled = true;
3062    }
3063}
3064
3065impl Drop for OptiForkTicket {
3066    fn drop(&mut self) {
3067        if !self.settled {
3068            let _ = self.drain.synchronize();
3069            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3070        }
3071    }
3072}
3073
3074impl OptiForkState {
3075    #[allow(clippy::too_many_arguments)]
3076    fn new(
3077        e: &Engine,
3078        cache: &Cache,
3079        mode: OptiForkGateMode,
3080        alternate_snapshot: crate::cache::CacheSnapshot,
3081        h_seed: &CudaSlice<f32>,
3082        fill_prev: &CudaSlice<f32>,
3083        rt: &'static crate::pp::PpNRt,
3084        split: usize,
3085        n_layer: usize,
3086    ) -> Result<Self, Box<dyn std::error::Error>> {
3087        let fence = [0, split, n_layer];
3088        let mut logical_payload_bytes = [0usize; 2];
3089        for stage in 0..2 {
3090            for il in fence[stage]..fence[stage + 1] {
3091                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3092                    .as_ref()
3093                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3094                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3095                    .as_ref()
3096                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3097            }
3098        }
3099        let seeds = [
3100            OptiForkSeedGeneration {
3101                h_seed: e.clone_dtod(h_seed)?,
3102                fill_prev: e.clone_dtod(fill_prev)?,
3103                scratch_len: 0,
3104            },
3105            OptiForkSeedGeneration {
3106                h_seed: e.clone_dtod(h_seed)?,
3107                fill_prev: e.clone_dtod(fill_prev)?,
3108                scratch_len: 0,
3109            },
3110        ];
3111        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3112            let _stage = rt.enter(0);
3113            let e0 = rt.engine(0, e);
3114            (
3115                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3116                e0.htod_i32(&vec![0; split])?,
3117                e0.alloc_u32_zeroed(2)?,
3118                e0.alloc_u32_zeroed(1)?,
3119                e0.stream(),
3120            )
3121        };
3122        logical_payload_bytes[0] += seeds
3123            .iter()
3124            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3125            .sum::<usize>();
3126        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3127            + saved_lens.len() * std::mem::size_of::<i32>()
3128            + forced_acc.len() * std::mem::size_of::<u32>()
3129            + valid.len() * std::mem::size_of::<u32>();
3130        Ok(Self {
3131            mode,
3132            controller: (mode == OptiForkGateMode::Controller)
3133                .then(OptiControllerPolicy::configured),
3134            generations: OptiForkGenerationTracker::default(),
3135            active_snapshot_slot: 0,
3136            alternate_snapshot,
3137            seeds,
3138            rt,
3139            fence,
3140            split,
3141            len_ptrs,
3142            saved_lens,
3143            forced_acc,
3144            valid,
3145            stage0_stream,
3146            logical_payload_bytes,
3147        })
3148    }
3149
3150    fn reserve(
3151        &mut self,
3152        current_snapshot: &mut crate::cache::CacheSnapshot,
3153    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3154        let generation = self.generations.reserve()?;
3155        if generation.slot != self.active_snapshot_slot {
3156            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3157            self.active_snapshot_slot = generation.slot;
3158        }
3159        Ok(generation)
3160    }
3161
3162    fn capture_seed(
3163        &mut self,
3164        e: &Engine,
3165        generation: OptiForkGeneration,
3166        h_seed: &CudaSlice<f32>,
3167        fill_prev: &CudaSlice<f32>,
3168        scratch_len: usize,
3169    ) -> Result<(), Box<dyn std::error::Error>> {
3170        let seed = &mut self.seeds[generation.slot];
3171        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3172        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3173        seed.scratch_len = scratch_len;
3174        Ok(())
3175    }
3176
3177    fn ticket(
3178        &self,
3179        generation: OptiForkGeneration,
3180        boundary: VerifyBoundaryTicket,
3181    ) -> OptiForkTicket {
3182        OptiForkTicket {
3183            generation,
3184            boundary: Some(boundary),
3185            drain: self.stage0_stream.clone(),
3186            settled: false,
3187        }
3188    }
3189
3190    #[allow(clippy::too_many_arguments)]
3191    fn controller_ticket(
3192        &self,
3193        generation: OptiForkGeneration,
3194        boundary: VerifyBoundaryTicket,
3195        ckpt: VerifyCkpt,
3196        verify_tokens: [u32; 2],
3197        draft_prob: f32,
3198        eager_seed: Option<CudaSlice<f32>>,
3199        q_proxy: f32,
3200        scratch_len: usize,
3201    ) -> OptiControllerTicket {
3202        OptiControllerTicket {
3203            generation,
3204            boundary: Some(boundary),
3205            ckpt: Some(ckpt),
3206            verify_tokens,
3207            draft_prob,
3208            eager_seed,
3209            q_proxy,
3210            scratch_len,
3211            issued_at: std::time::Instant::now(),
3212            drain: self.stage0_stream.clone(),
3213            settled: false,
3214        }
3215    }
3216
3217    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3218        self.generations.reserve()
3219    }
3220
3221    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3222        &mut self.alternate_snapshot
3223    }
3224
3225    fn promote_successor_snapshot(
3226        &mut self,
3227        current_snapshot: &mut crate::cache::CacheSnapshot,
3228        generation: OptiForkGeneration,
3229    ) {
3230        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3231        self.active_snapshot_slot = generation.slot;
3232    }
3233
3234    fn queue_actual_reconcile(
3235        &mut self,
3236        e: &Engine,
3237        snapshot: &crate::cache::CacheSnapshot,
3238        acc: &CudaSlice<u32>,
3239        optimistic_pending: u32,
3240        base: usize,
3241    ) -> Result<(), Box<dyn std::error::Error>> {
3242        let saved: Vec<i32> = (0..self.split)
3243            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3244            .collect();
3245        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3246        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3247        // the validity/reconcile kernels must never peer-read acc before it is written. The
3248        // increment-1 harness uses primary stage 0, where stream order already provides this.
3249        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3250            self.rt.fence_stages_behind(&e.stream())?;
3251        }
3252        let _stage = self.rt.enter(0);
3253        let e0 = self.rt.engine(0, e);
3254        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3255        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3256        e0.spec_fork_reconcile_kv(
3257            &self.len_ptrs,
3258            &self.saved_lens,
3259            acc,
3260            &self.valid,
3261            base,
3262            self.split,
3263        )
3264    }
3265
3266    fn finish_actual_reconcile(
3267        &mut self,
3268        e: &Engine,
3269        cache: &mut Cache,
3270        snapshot: &crate::cache::CacheSnapshot,
3271        n_acc: usize,
3272        base: usize,
3273        hit: bool,
3274    ) -> Result<(), Box<dyn std::error::Error>> {
3275        if hit {
3276            return Ok(());
3277        }
3278        let len_delta = base + n_acc;
3279        for il in 0..self.split {
3280            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3281                kv.len = saved + len_delta;
3282            }
3283        }
3284        {
3285            let _stage = self.rt.enter(1);
3286            let e1 = self.rt.engine(1, e);
3287            for il in self.split..self.fence[2] {
3288                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3289                    kv.len = saved + len_delta;
3290                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3291                }
3292            }
3293        }
3294        self.rt.publish_to(0, &e.stream())?;
3295        Ok(())
3296    }
3297
3298    fn cancel_controller_ticket(
3299        &mut self,
3300        e: &Engine,
3301        cache: &mut Cache,
3302        scratch: &mut MtpScratch,
3303        snapshot: &crate::cache::CacheSnapshot,
3304        ticket: &mut OptiControllerTicket,
3305    ) -> Result<(), Box<dyn std::error::Error>> {
3306        {
3307            let _stage = self.rt.enter(0);
3308            let e0 = self.rt.engine(0, e);
3309            for il in 0..self.split {
3310                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3311                    kv.len = saved;
3312                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3313                }
3314            }
3315        }
3316        scratch.set_len(e, snapshot.pos)?;
3317        ticket.settle();
3318        self.generations.retire(ticket.generation)?;
3319        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3320        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3321        eprintln!(
3322            "[opti-controller] tail-drain generation={} slot={}",
3323            ticket.generation.id, ticket.generation.slot,
3324        );
3325        Ok(())
3326    }
3327
3328    #[allow(clippy::too_many_arguments)]
3329    fn reconcile(
3330        &mut self,
3331        e: &Engine,
3332        cache: &mut Cache,
3333        scratch: &mut MtpScratch,
3334        snapshot: &crate::cache::CacheSnapshot,
3335        h_seed: &mut CudaSlice<f32>,
3336        fill_prev: &mut CudaSlice<f32>,
3337        generation: OptiForkGeneration,
3338        action: OptiForkAction,
3339        optimistic_pending: u32,
3340    ) -> Result<(), Box<dyn std::error::Error>> {
3341        debug_assert!(action != OptiForkAction::Abort);
3342        let miss_started = std::time::Instant::now();
3343        let keep = action == OptiForkAction::Hit;
3344        let saved: Vec<i32> = (0..self.split)
3345            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3346            .collect();
3347        let seed = &self.seeds[generation.slot];
3348        {
3349            let _stage = self.rt.enter(0);
3350            let e0 = self.rt.engine(0, e);
3351            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3352            let forced = if keep {
3353                [1u32, optimistic_pending]
3354            } else {
3355                [0u32, optimistic_pending]
3356            };
3357            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3358            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3359            e0.spec_fork_reconcile_kv(
3360                &self.len_ptrs,
3361                &self.saved_lens,
3362                &self.forced_acc,
3363                &self.valid,
3364                0,
3365                self.split,
3366            )?;
3367            for il in 0..self.split {
3368                if let Some(recur) = cache.recur[il].as_mut() {
3369                    let conv = snapshot.conv[il]
3370                        .as_ref()
3371                        .ok_or("optipipe stage0 snapshot missing conv state")?;
3372                    let ssm = snapshot.ssm[il]
3373                        .as_ref()
3374                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
3375                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3376                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3377                }
3378            }
3379            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3380            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3381        }
3382
3383        if keep {
3384            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3385            return Ok(());
3386        }
3387
3388        for il in 0..self.split {
3389            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3390                kv.len = saved;
3391            }
3392        }
3393        scratch.set_len(e, seed.scratch_len)?;
3394        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3395        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3396        let caller = e.stream();
3397        self.rt.publish_to(0, &caller)?;
3398        caller.synchronize()?;
3399        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3400        eprintln!(
3401            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3402            generation.id, generation.slot,
3403        );
3404        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3405        Ok(())
3406    }
3407
3408    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3409        self.generations.retire(generation)
3410    }
3411}
3412
3413fn rewind_tp_kv_verified_prefix(
3414    tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3415    saved_lens: &[Option<usize>],
3416    accepted: usize,
3417) -> Result<(), Box<dyn std::error::Error>> {
3418    if tp_kv.len() != saved_lens.len() {
3419        return Err("spec TP KV snapshot shape mismatch".into());
3420    }
3421    for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3422        match (cache.as_mut(), *saved) {
3423            (Some(cache), Some(saved)) => {
3424                let target = saved
3425                    .checked_add(accepted)
3426                    .ok_or("spec TP KV committed length overflow")?;
3427                cache.rewind_to(target)?;
3428            }
3429            (None, None) => {}
3430            _ => {
3431                return Err(
3432                    format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3433                );
3434            }
3435        }
3436    }
3437    Ok(())
3438}
3439
3440impl HybridModel {
3441    fn mtp_head_count(&self) -> usize {
3442        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3443    }
3444
3445    fn mtp_head_at(&self, index: usize) -> &MtpHead {
3446        if index == 0 {
3447            self.mtp.as_ref().expect("MTP head 0 is unavailable")
3448        } else {
3449            &self.mtp_extra[index - 1]
3450        }
3451    }
3452
3453    fn new_mtp_scratch(
3454        &self,
3455        e: &Engine,
3456        cap: usize,
3457    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3458        let mut scratch = MtpScratch::new(
3459            e,
3460            &self.cfg,
3461            &self.plan,
3462            cap,
3463            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3464        )?;
3465        for head in &self.mtp_extra {
3466            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3467        }
3468        Ok(scratch)
3469    }
3470
3471    fn opti_graph_draft_step(
3472        &self,
3473        e: &Engine,
3474        mtp: &MtpHead,
3475        dctx: &mut DraftGraphCtx,
3476        scratch: &mut MtpScratch,
3477        d_vocab: usize,
3478    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3479        dctx.graph
3480            .as_ref()
3481            .ok_or("optipipe controller requires the greedy draft graph")?
3482            .launch()?;
3483        scratch.kv.len += 1;
3484        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
3485        if (idx as usize) >= d_vocab {
3486            return Err(
3487                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
3488            );
3489        }
3490        let probability = e.dtoh(&dctx.g_p)?[0];
3491        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3492            return Err(format!("optipipe draft probability is invalid: {probability}").into());
3493        }
3494        let token = match &mtp.d2t {
3495            Some(map) => map[idx as usize],
3496            None => idx,
3497        };
3498        if token != idx {
3499            e.set_u32_one(&mut dctx.g_tok, token)?;
3500        }
3501        Ok((token, probability))
3502    }
3503
3504    #[allow(clippy::too_many_arguments)]
3505    fn opti_controller_draft_step(
3506        &self,
3507        e: &Engine,
3508        mtp: &MtpHead,
3509        dctx: &mut DraftGraphCtx,
3510        scratch: &mut MtpScratch,
3511        d_vocab: usize,
3512        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
3513        eager_pos: usize,
3514        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3515    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3516        if dctx.graph.is_some() {
3517            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
3518        }
3519        let (input_token, input_seed) = eager_state
3520            .take()
3521            .ok_or("optipipe eager continuation seed is unavailable")?;
3522        let (logits, next_seed) = self.mtp_head_forward_dev(
3523            e,
3524            mtp,
3525            input_token,
3526            &input_seed,
3527            scratch,
3528            eager_pos,
3529            embd_dev,
3530            None,
3531        )?;
3532        let token_d = e.argmax_token_device(&logits, d_vocab)?;
3533        let idx = e.dtoh_u32_one(&token_d)?;
3534        if (idx as usize) >= d_vocab {
3535            return Err(format!(
3536                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
3537            )
3538            .into());
3539        }
3540        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
3541        let probability = e.dtoh(&probability_d)?[0];
3542        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3543            return Err(
3544                format!("optipipe eager draft probability is invalid: {probability}").into(),
3545            );
3546        }
3547        let token = match &mtp.d2t {
3548            Some(map) => map[idx as usize],
3549            None => idx,
3550        };
3551        *eager_state = Some((token, next_seed));
3552        Ok((token, probability))
3553    }
3554
3555    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3556    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3557    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3558    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3559    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3560    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3561    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3562    /// transfer + host argmax per draft token from the K-token draft chain.
3563    #[allow(clippy::too_many_arguments)]
3564    fn mtp_head_forward_dev(
3565        &self,
3566        e: &Engine,
3567        mtp: &MtpHead,
3568        e_tok: u32,
3569        h_seed: &CudaSlice<f32>,
3570        scratch: &mut MtpScratch,
3571        mtp_pos: usize,
3572        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3573        mask: Option<(&CudaSlice<u32>, usize)>,
3574    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3575        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3576    }
3577
3578    #[allow(clippy::too_many_arguments)]
3579    fn mtp_head_forward_dev_at(
3580        &self,
3581        e: &Engine,
3582        mtp: &MtpHead,
3583        e_tok: u32,
3584        h_seed: &CudaSlice<f32>,
3585        scratch: &mut MtpScratch,
3586        scratch_index: usize,
3587        mtp_pos: usize,
3588        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3589        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3590        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3591        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3592        mask: Option<(&CudaSlice<u32>, usize)>,
3593    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3594        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3595        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3596        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3597        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3598        static ANAT_NS: [AtomicU64; 5] = [
3599            AtomicU64::new(0),
3600            AtomicU64::new(0),
3601            AtomicU64::new(0),
3602            AtomicU64::new(0),
3603            AtomicU64::new(0),
3604        ];
3605        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3606        let anat = {
3607            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3608            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3609        };
3610        if anat {
3611            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3612        }
3613        let t_all = std::time::Instant::now();
3614        let mut t_ph = std::time::Instant::now();
3615        let mut anat_mark = |i: usize,
3616                             e: &Engine,
3617                             t: &mut std::time::Instant|
3618         -> Result<(), Box<dyn std::error::Error>> {
3619            if anat {
3620                e.stream().synchronize()?;
3621                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3622                *t = std::time::Instant::now();
3623            }
3624            Ok(())
3625        };
3626        let cfg = &self.cfg;
3627        let n_embd = cfg.n_embd as usize;
3628        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3629        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3630        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3631        let eps = cfg.rms_eps;
3632        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3633
3634        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3635        // expands this one row on CPU and transfers n_embd f32 values instead.
3636        let e_emb = match embd_dev {
3637            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3638            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
3639        };
3640
3641        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3642        let mut e_norm = e.zeros(n_embd)?;
3643        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3644        let mut h_norm = e.zeros(n_embd)?;
3645        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3646
3647        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3648        let mut concat = e.zeros(2 * n_embd)?;
3649        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3650        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3651
3652        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3653        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3654
3655        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3656        let mut a_norm = e.zeros(di)?;
3657        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3658        anat_mark(0, e, &mut t_ph)?;
3659
3660        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3661        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3662        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3663        // advances only the device counter).
3664        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3665            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3666            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
3667            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
3668            // whose host-side mirror the caller does).
3669            (Mixer::Full(fa), Some(g)) => {
3670                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3671            }
3672            (Mixer::Full(fa), None) => {
3673                let out = self.mtp_full_attn_dc(
3674                    e,
3675                    fa,
3676                    &a_norm,
3677                    &pos_d,
3678                    scratch,
3679                    scratch_index,
3680                    mtp.geom.as_ref(),
3681                )?;
3682                scratch.plane_mut(scratch_index).0.len += 1;
3683                out
3684            }
3685            (Mixer::Linear(_), _) => {
3686                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3687            }
3688            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
3689        };
3690        anat_mark(1, e, &mut t_ph)?;
3691
3692        // op 7: x1 = inpSA + attn_out
3693        let mut x1 = e.zeros(di)?;
3694        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3695
3696        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
3697        let mut z = e.zeros(di)?;
3698        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3699
3700        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3701        let ffn_out = match &mtp.ffn {
3702            crate::hybrid::Ffn::Dense {
3703                ffn_gate,
3704                ffn_up,
3705                ffn_down,
3706            } => {
3707                let n_ff = ffn_gate.out_features();
3708                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3709                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3710                    (
3711                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3712                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3713                    )
3714                } else {
3715                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3716                };
3717                let mut act = e.zeros(n_ff)?;
3718                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3719                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3720                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3721                // passes None, which is `ffn_act`'s dispatch verbatim.
3722                Self::ffn_act_lim(
3723                    e,
3724                    &self.cfg,
3725                    &gate,
3726                    &up,
3727                    1.0,
3728                    1.0,
3729                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
3730                    &mut act,
3731                    n_ff,
3732                )?;
3733                e.matmul(ffn_down, &act, 1)?
3734            }
3735            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3736            // so they never alias trunk layer 0's cache keys.
3737            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3738        };
3739        anat_mark(2, e, &mut t_ph)?;
3740
3741        // op 10: h_nextn = x1 + ffn_out (at di)
3742        let mut h_inner = e.zeros(di)?;
3743        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3744
3745        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3746        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3747        let h_nextn = match mtp.geom.as_ref() {
3748            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3749            None => h_inner,
3750        };
3751
3752        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3753        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3754        let mut final_h = e.zeros(n_embd)?;
3755        e.rms_norm(
3756            &h_nextn,
3757            final_norm.float_data(),
3758            &mut final_h,
3759            n_embd,
3760            1,
3761            eps,
3762        )?;
3763
3764        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3765        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3766        let mut logits = e.matmul(head, &final_h, 1)?;
3767        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3768        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3769        if let Some((mask_d, mw)) = mask {
3770            let d_vocab = head.out_features();
3771            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3772        }
3773        anat_mark(3, e, &mut t_ph)?;
3774        if anat {
3775            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3776            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3777            if n % 128 == 0 {
3778                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3779                eprintln!(
3780                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3781                    us(0),
3782                    us(1),
3783                    us(2),
3784                    us(3),
3785                    us(4)
3786                );
3787            }
3788        }
3789        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3790        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3791        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3792    }
3793
3794    #[allow(clippy::too_many_arguments)]
3795    fn mtp_chain_forward_dev(
3796        &self,
3797        e: &Engine,
3798        tokens: &[u32],
3799        seeds: &[CudaSlice<f32>],
3800        scratch: &mut MtpScratch,
3801        committed_scratch_len: usize,
3802        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3803        mask: Option<(&CudaSlice<u32>, usize)>,
3804    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3805        if tokens.is_empty() || tokens.len() != seeds.len() {
3806            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3807        }
3808        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3809        let head = self.mtp_head_at(index);
3810        scratch.set_plane_len(e, index, committed_scratch_len)?;
3811
3812        let mut last = None;
3813        for row in 0..tokens.len() {
3814            let is_last = row + 1 == tokens.len();
3815            last = Some(self.mtp_head_forward_dev_at(
3816                e,
3817                head,
3818                tokens[row],
3819                &seeds[row],
3820                scratch,
3821                index,
3822                committed_scratch_len + row + 1,
3823                embd_dev,
3824                if is_last { mask } else { None },
3825            )?);
3826        }
3827        Ok(last.expect("non-empty MTP prefix produced no row"))
3828    }
3829
3830    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3831    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3832    /// the dc path, and all three are properties of this arch's MTP block:
3833    ///
3834    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3835    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3836    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3837    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3838    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3839    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3840    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
3841    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3842    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3843    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3844    ///    resolved `Step35MtpGeom`, never from `cfg`.
3845    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3846    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3847    ///    fused-into-wq `q_gate_split` form the dc arm handles.
3848    ///
3849    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3850    /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3851    /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3852    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3853    ///
3854    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3855    /// caller must not mirror.
3856    fn mtp_step35_attn(
3857        &self,
3858        e: &Engine,
3859        fa: &FullAttnLayer,
3860        g: &crate::hybrid::Step35MtpGeom,
3861        h: &CudaSlice<f32>,
3862        pos_d: &CudaSlice<i32>,
3863        scratch: &mut MtpScratch,
3864        scratch_index: usize,
3865    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3866        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3867        let eps = self.cfg.rms_eps;
3868        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3869        let n_embd = self.cfg.n_embd as usize;
3870        let gw = fa
3871            .attn_gate
3872            .as_ref()
3873            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3874
3875        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3876            && e.uses_q8_1_fast(&fa.wk)
3877            && e.uses_q8_1_fast(&fa.wv)
3878            && e.uses_q8_1_fast(gw)
3879        {
3880            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3881            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3882                Some(t3) => t3,
3883                None => (
3884                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3885                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3886                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3887                ),
3888            };
3889            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
3890        } else {
3891            (
3892                e.matmul(&fa.wq, h, 1)?,
3893                e.matmul(&fa.wk, h, 1)?,
3894                e.matmul(&fa.wv, h, 1)?,
3895                e.matmul(gw, h, 1)?,
3896            )
3897        };
3898
3899        let mut q = e.uninit(nh * hd)?;
3900        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
3901        let mut k = e.uninit(nkv * hd)?;
3902        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
3903        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
3904        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
3905        // the resolved flag, not the constant, so an all-full sibling stays correct.
3906        let ff = if g.swa {
3907            None
3908        } else {
3909            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
3910        };
3911        #[cfg(debug_assertions)]
3912        if let Some(ff) = ff {
3913            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
3914        }
3915        e.rope_neox2(
3916            &mut q,
3917            &mut k,
3918            pos_d,
3919            hd,
3920            g.n_rot,
3921            nh,
3922            nkv,
3923            1,
3924            g.rope_base,
3925            1.0,
3926            ff,
3927        )?;
3928
3929        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
3930        // length on the host anyway, and the windowed view below needs it there to compute the
3931        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
3932        // dc-family consumer of this scratch still agree.
3933        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
3934        assert!(
3935            kv.len < scratch_cap,
3936            "step35 MTP scratch overflow ({} >= {})",
3937            kv.len,
3938            scratch_cap
3939        );
3940        let next_len = kv.len + 1;
3941        let (off, t_kv) = if g.swa && next_len > g.window {
3942            (next_len - g.window, g.window)
3943        } else {
3944            (0, next_len)
3945        };
3946        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
3947        e.append_kv_quantized(
3948            &k,
3949            &v0,
3950            &mut kv.k,
3951            &mut kv.v,
3952            write_row,
3953            kv.kv_dim_k,
3954            kv.kv_dim_v,
3955            kv.k_tok_bytes,
3956            kv.v_tok_bytes,
3957            false,
3958        )?;
3959        kv.len = next_len;
3960        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3961        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
3962        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
3963        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
3964        // therefore live, not theoretical.
3965        let physical = kv.physical_rows(off, off + t_kv)?;
3966        let k_view = e.view_u8_range(
3967            &kv.k,
3968            physical.start * kv.k_tok_bytes,
3969            physical.end * kv.k_tok_bytes,
3970        );
3971        let v_view = e.view_u8_range(
3972            &kv.v,
3973            physical.start * kv.v_tok_bytes,
3974            physical.end * kv.v_tok_bytes,
3975        );
3976        let mut attn = e.uninit(nh * hd)?;
3977        e.fa_decode_kvmod(
3978            &q,
3979            &k_view,
3980            &v_view,
3981            &mut attn,
3982            hd,
3983            nh,
3984            nkv,
3985            t_kv,
3986            scale,
3987            kv.k_tok_bytes,
3988            kv.v_tok_bytes,
3989            false,
3990        )?;
3991
3992        let mut ag = e.uninit(nh * hd)?;
3993        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
3994        Ok(e.matmul(&fa.wo, &ag, 1)?)
3995    }
3996
3997    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
3998    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
3999    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4000    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4001    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4002    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4003    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4004    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4005    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4006    fn mtp_full_attn_dc(
4007        &self,
4008        e: &Engine,
4009        fa: &FullAttnLayer,
4010        h: &CudaSlice<f32>,
4011        pos_d: &CudaSlice<i32>,
4012        scratch: &mut MtpScratch,
4013        scratch_index: usize,
4014        geom: Option<&crate::hybrid::DraftGeom>,
4015    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4016        let cfg = &self.cfg;
4017        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4018        let geometry = cfg.full_attention_geometry_at(mtp_il);
4019        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4020        let n_head_kv = geom
4021            .map(|g| g.n_head_kv)
4022            .unwrap_or(geometry.n_head_kv as usize);
4023        let head_dim = geometry.head_dim_k as usize;
4024        let eps = cfg.rms_eps;
4025        let scale = geometry.attention_scale();
4026        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4027        let bucket_max = scratch.plane(scratch_index).1;
4028
4029        let (qf, mut k, v) =
4030            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4031                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4032                (
4033                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4034                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4035                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4036                )
4037            } else {
4038                (
4039                    e.matmul(&fa.wq, h, 1)?,
4040                    e.matmul(&fa.wk, h, 1)?,
4041                    e.matmul(&fa.wv, h, 1)?,
4042                )
4043            };
4044        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4045        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4046        let (mut q, gate) = if gated {
4047            let mut q = e.zeros(n_head * head_dim)?;
4048            let mut gate = e.zeros(n_head * head_dim)?;
4049            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4050            (q, Some(gate))
4051        } else {
4052            (qf, None)
4053        };
4054
4055        let mut qn = e.zeros(n_head * head_dim)?;
4056        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4057        q = qn;
4058        let mut kn = e.zeros(n_head_kv * head_dim)?;
4059        e.rms_norm(
4060            &k,
4061            fa.k_norm.float_data(),
4062            &mut kn,
4063            head_dim,
4064            n_head_kv,
4065            eps,
4066        )?;
4067        k = kn;
4068        let rope_dims = geometry.n_rot as usize;
4069        e.rope_neox(
4070            &mut q,
4071            pos_d,
4072            head_dim,
4073            rope_dims,
4074            n_head,
4075            1,
4076            geometry.rope_base,
4077            1.0,
4078        )?;
4079        e.rope_neox(
4080            &mut k,
4081            pos_d,
4082            head_dim,
4083            rope_dims,
4084            n_head_kv,
4085            1,
4086            geometry.rope_base,
4087            1.0,
4088        )?;
4089
4090        let kv = scratch.plane_mut(scratch_index).0;
4091        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4092        e.append_kv_quantized_dc(
4093            &k,
4094            &v,
4095            &mut kv.k,
4096            &mut kv.v,
4097            &kv.len_d,
4098            kv.kv_dim_k,
4099            kv.kv_dim_v,
4100            kv.k_tok_bytes,
4101            kv.v_tok_bytes,
4102            false,
4103        )?;
4104        e.inc_seqlen(&mut kv.len_d)?;
4105        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4106        // key range from the device counter.
4107        let k_view = e.view_u8(&kv.k, kv.k.len());
4108        let v_view = e.view_u8(&kv.v, kv.v.len());
4109        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4110        let mut attn = e.zeros(n_head * head_dim)?;
4111        e.fa_decode_dc(
4112            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4113            scale, ktb, vtb, false,
4114        )?;
4115
4116        let attn_g = match &gate {
4117            Some(gate) => {
4118                let mut gsig = e.zeros(n_head * head_dim)?;
4119                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4120                let mut ag = e.zeros(n_head * head_dim)?;
4121                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4122                ag
4123            }
4124            None => attn,
4125        };
4126        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4127    }
4128
4129    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4130    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4131    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4132    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4133    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4134    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4135    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4136    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4137    #[allow(clippy::too_many_arguments)]
4138    fn mtp_kv_fill_at(
4139        &self,
4140        e: &Engine,
4141        mtp: &MtpHead,
4142        tokens: &[u32],
4143        h: &CudaSlice<f32>,
4144        pos0: usize,
4145        scratch: &mut MtpScratch,
4146        scratch_index: usize,
4147        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4148    ) -> Result<(), Box<dyn std::error::Error>> {
4149        let cfg = &self.cfg;
4150        let n_embd = cfg.n_embd as usize;
4151        let eps = cfg.rms_eps;
4152        let t = tokens.len();
4153        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4154        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4155        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4156        let Mixer::Full(fa) = &mtp.mixer else {
4157            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4158        };
4159        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4160        let pos_d = e.htod_i32(&pos_vec)?;
4161
4162        // ops A/1/2: embed + the two input norms, T-wide.
4163        let e_emb = match embd_dev {
4164            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4165            None => e.htod(&self.embd.gather(n_embd, tokens))?,
4166        };
4167        let mut e_norm = e.zeros(t * n_embd)?;
4168        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4169        let mut h_norm = e.zeros(t * n_embd)?;
4170        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4171
4172        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4173        let mut concat = e.zeros(t * 2 * n_embd)?;
4174        for i in 0..t {
4175            e.copy_view_into(
4176                &mut concat,
4177                i * 2 * n_embd,
4178                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4179                n_embd,
4180            )?;
4181            e.copy_view_into(
4182                &mut concat,
4183                i * 2 * n_embd + n_embd,
4184                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4185                n_embd,
4186            )?;
4187        }
4188
4189        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4190        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4191        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4192        let mut a_norm = e.zeros(t * di)?;
4193        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4194
4195        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4196        // the fill only has to leave correct K/V rows behind for later chains to attend over.
4197        let n_head_kv = mtp
4198            .geom
4199            .as_ref()
4200            .map(|g| g.n_head_kv)
4201            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4202            .unwrap_or_else(|| {
4203                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4204                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4205            });
4206        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4207        let geometry = cfg.full_attention_geometry_at(mtp_il);
4208        let head_dim = geometry.head_dim_k as usize;
4209        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4210        let v = e.matmul(&fa.wv, &a_norm, t)?;
4211        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4212        e.rms_norm(
4213            &k,
4214            fa.k_norm.float_data(),
4215            &mut kn,
4216            head_dim,
4217            n_head_kv * t,
4218            eps,
4219        )?;
4220        k = kn;
4221        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4222        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4223        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4224        // writes K rows the attention arm then re-derives at a different theta: correct-looking
4225        // output with dead acceptance, invisible to the exactness gates.
4226        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4227            Some(s) => (
4228                s.n_rot,
4229                s.rope_base,
4230                if s.swa {
4231                    None
4232                } else {
4233                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4234                },
4235            ),
4236            None => (geometry.n_rot as usize, geometry.rope_base, None),
4237        };
4238        #[cfg(debug_assertions)]
4239        if let Some(ff) = ff {
4240            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4241        }
4242        match ff {
4243            Some(f) => e.rope_neox_ff(
4244                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4245            )?,
4246            None => e.rope_neox(
4247                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4248            )?,
4249        }
4250
4251        let kv = scratch.plane_mut(scratch_index).0;
4252        // Match the trunk prime contract: a chunk may need the aligned window immediately before
4253        // its first row, so preserve that prefix when the physical tail rebases at wrap.
4254        let retain_from = kv
4255            .ring
4256            .as_ref()
4257            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
4258            .unwrap_or(0);
4259        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4260        for i in 0..t {
4261            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4262            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4263            e.append_kv_quantized_view(
4264                &k_row,
4265                &v_row,
4266                &mut kv.k,
4267                &mut kv.v,
4268                write_row + i,
4269                kv.kv_dim_k,
4270                kv.kv_dim_v,
4271                kv.k_tok_bytes,
4272                kv.v_tok_bytes,
4273                false,
4274            )?;
4275        }
4276        kv.len = pos0 + t;
4277        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4278        Ok(())
4279    }
4280
4281    #[allow(clippy::too_many_arguments)]
4282    fn mtp_kv_fill_all(
4283        &self,
4284        e: &Engine,
4285        tokens: &[u32],
4286        h: &CudaSlice<f32>,
4287        pos0: usize,
4288        scratch: &mut MtpScratch,
4289        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4290    ) -> Result<(), Box<dyn std::error::Error>> {
4291        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4292        for index in 0..self.mtp_head_count() {
4293            self.mtp_kv_fill_at(
4294                e,
4295                self.mtp_head_at(index),
4296                tokens,
4297                h,
4298                pos0,
4299                scratch,
4300                index,
4301                embd_dev,
4302            )?;
4303        }
4304        Ok(())
4305    }
4306
4307    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4308    /// every varying input device-resident —
4309    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4310    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4311    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4312    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4313    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4314    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4315    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4316    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4317    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4318    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4319    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4320    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4321    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4322    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4323    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4324    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4325    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4326    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4327    #[allow(clippy::too_many_arguments)]
4328    fn mtp_head_forward_cap(
4329        &self,
4330        e: &Engine,
4331        mtp: &MtpHead,
4332        tok_d: &mut CudaSlice<u32>,
4333        pos_d: &mut CudaSlice<i32>,
4334        h_seed_d: &mut CudaSlice<f32>,
4335        p_d: &mut CudaSlice<f32>,
4336        scratch: &mut MtpScratch,
4337        with_prob: bool,
4338        with_head: bool,
4339        embd_gpu: &CudaSlice<u8>,
4340        embd_qt: i32,
4341        embd_rb: usize,
4342        d_vocab: usize,
4343        sampled_cap: Option<(
4344            &mut CudaSlice<u32>,
4345            &mut CudaSlice<f32>,
4346            &mut CudaSlice<f32>,
4347            u64,
4348            f32,
4349        )>,
4350        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4351        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4352        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4353        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4354        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4355        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4356        mask_cap: Option<(&CudaSlice<u32>, usize)>,
4357    ) -> Result<(), Box<dyn std::error::Error>> {
4358        let cfg = &self.cfg;
4359        let n_embd = cfg.n_embd as usize;
4360        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
4361        // whose device-counter key bound always starts at row 0 — it cannot express this block's
4362        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
4363        // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
4364        // refuses step35 heads explicitly (SWA refusal), so the eager chain
4365        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
4366        // panic) is what the two capture sites and the round-stream capture already handle by
4367        // degrading to eager / stream-off.
4368        if mtp.step35.is_some() {
4369            return Err(
4370                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4371                        block's SWA view offset; same root cause as the dc decode refusal) — the \
4372                        eager draft chain serves this arch"
4373                    .into(),
4374            );
4375        }
4376        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4377        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4378        let eps = cfg.rms_eps;
4379        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4380        let mut e_norm = e.zeros(n_embd)?;
4381        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4382        let mut h_norm = e.zeros(n_embd)?;
4383        e.rms_norm(
4384            &*h_seed_d,
4385            mtp.hnorm.float_data(),
4386            &mut h_norm,
4387            n_embd,
4388            1,
4389            eps,
4390        )?;
4391        let mut concat = e.zeros(2 * n_embd)?;
4392        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4393        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4394        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4395        let mut a_norm = e.zeros(di)?;
4396        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4397        let attn_out = match &mtp.mixer {
4398            Mixer::Full(fa) => {
4399                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
4400            }
4401            Mixer::Linear(_) => {
4402                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4403            }
4404            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4405        };
4406        let mut x1 = e.zeros(di)?;
4407        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4408        let mut z = e.zeros(di)?;
4409        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4410        let ffn_out = match &mtp.ffn {
4411            crate::hybrid::Ffn::Dense {
4412                ffn_gate,
4413                ffn_up,
4414                ffn_down,
4415            } => {
4416                let n_ff = ffn_gate.out_features();
4417                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4418                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4419                    (
4420                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4421                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4422                    )
4423                } else {
4424                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4425                };
4426                let mut act = e.zeros(n_ff)?;
4427                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
4428                e.matmul(ffn_down, &act, 1)?
4429            }
4430            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
4431            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
4432            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
4433            // error arm degrades the caller to eager/stream-off.
4434            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
4435                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4436            }
4437            crate::hybrid::Ffn::Moe(_) => {
4438                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
4439            }
4440        };
4441        let mut h_inner = e.zeros(di)?;
4442        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4443        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4444        let h_nextn = match mtp.geom.as_ref() {
4445            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4446            None => h_inner,
4447        };
4448        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4449        let final_h = if with_head || spec_hpost() {
4450            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4451            let mut fh = e.zeros(n_embd)?;
4452            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4453            Some(fh)
4454        } else {
4455            None
4456        };
4457        if with_head {
4458            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4459            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4460            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4461            // before the argmax — proposals become legal by construction. Contents-only
4462            // per-replay upload keeps the capture valid.
4463            if let Some((mask_d, mw)) = mask_cap {
4464                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4465            }
4466            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
4467                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4468                // own buffer is pool-recycled after the capture body returns, so it can't be the
4469                // retention target), bump the device event counter, gumbel-perturb reading it,
4470                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4471                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4472                e.sctr_inc(ctr_d)?;
4473                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
4474                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4475                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4476                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4477                if with_prob {
4478                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4479                }
4480            } else {
4481                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4482                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4483                // p-min under a draft mask reads the MASKED row: confidence relative to the
4484                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4485                // is the right semantics for "does the drafter know what comes next here" and
4486                // the same row the pick came from. Draft-quality only — verify arbitrates.
4487                if with_prob {
4488                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4489                }
4490            }
4491        }
4492        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4493        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4494        if let Some((out, slot, d2t)) = stream_pack {
4495            e.pack_tok_p(tok_d, p_d, out, slot)?;
4496            if let Some(map) = d2t {
4497                e.tok_map_u32(tok_d, map)?;
4498            }
4499        }
4500        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4501        if spec_hpost() {
4502            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4503        } else {
4504            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4505        }
4506        // advance the draft rope position in-graph.
4507        e.inc_seqlen(pos_d)?;
4508        Ok(())
4509    }
4510
4511    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4512    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4513    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4514    /// Advances `cache.pos` by T.
4515    pub fn decode_step_t(
4516        &self,
4517        e: &Engine,
4518        tokens: &[u32],
4519        pos0: usize,
4520        cache: &mut Cache,
4521    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4522        if self.is_gemma4_e4b() {
4523            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4524        }
4525        if self.gemma_batch_program() {
4526            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4527        }
4528        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4529    }
4530
4531    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4532    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4533    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4534    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4535    pub fn decode_step_t_h(
4536        &self,
4537        e: &Engine,
4538        tokens: &[u32],
4539        pos0: usize,
4540        cache: &mut Cache,
4541    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4542        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4543    }
4544
4545    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4546    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4547    pub fn decode_step_t_h_emb(
4548        &self,
4549        e: &Engine,
4550        tokens: &[u32],
4551        pos0: usize,
4552        cache: &mut Cache,
4553        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4554    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4555        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4556        Ok((e.dtoh(&logits_d)?, h_seed))
4557    }
4558
4559    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4560    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4561    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4562    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4563    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4564    pub fn decode_step_t_h_emb_dev(
4565        &self,
4566        e: &Engine,
4567        tokens: &[u32],
4568        pos0: usize,
4569        cache: &mut Cache,
4570        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4571    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4572        let n_embd = self.cfg.n_embd as usize;
4573        let t = tokens.len();
4574        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4575        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4576        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4577        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4578        Ok((logits, hs))
4579    }
4580
4581    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4582    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4583    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4584    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4585    /// retains/copies — they never change what any kernel computes).
4586    fn decode_step_t_core(
4587        &self,
4588        e: &Engine,
4589        tokens: &[u32],
4590        pos0: usize,
4591        cache: &mut Cache,
4592        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4593        mut ckpt: Option<&mut VerifyCkpt>,
4594    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4595        self.decode_step_t_core_stream(
4596            e,
4597            tokens,
4598            pos0,
4599            cache,
4600            embd_dev,
4601            ckpt.take(),
4602            None,
4603            None,
4604            None,
4605            None,
4606        )
4607    }
4608
4609    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4610    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4611    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4612    fn decode_step_t_core_vg(
4613        &self,
4614        e: &Engine,
4615        tokens: &[u32],
4616        pos0: usize,
4617        cache: &mut Cache,
4618        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4619        mut ckpt: Option<&mut VerifyCkpt>,
4620        graphs: Option<&mut DsparkVerifyGraphs>,
4621    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4622        self.decode_step_t_core_stream(
4623            e,
4624            tokens,
4625            pos0,
4626            cache,
4627            embd_dev,
4628            ckpt.take(),
4629            None,
4630            None,
4631            None,
4632            graphs,
4633        )
4634    }
4635
4636    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
4637    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
4638    fn decode_step_t_core_pipelined(
4639        &self,
4640        e: &Engine,
4641        tokens: &[u32],
4642        pos0: usize,
4643        cache: &mut Cache,
4644        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4645        mut ckpt: Option<&mut VerifyCkpt>,
4646        pipe: &SpecPipeLane,
4647        round: usize,
4648    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4649        let fence = crate::pp::pp_cuts(self.layers.len())
4650            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
4651        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
4652            return Err("two-session speculative pipeline requires the PP verify split".into());
4653        }
4654        let interval_fence = pipe.stage0_begin(round)?;
4655        let ticket = self.verify_stage0_issue(
4656            e,
4657            tokens,
4658            pos0,
4659            cache,
4660            embd_dev,
4661            ckpt.as_deref_mut(),
4662            None,
4663            &fence,
4664            Some(interval_fence),
4665            pipe.trace(round),
4666        )?;
4667        pipe.stage0_end(round);
4668        pipe.stage1_begin(round)?;
4669        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
4670        pipe.verify_end(round);
4671        Ok(result)
4672    }
4673
4674    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4675    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4676    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4677    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4678    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4679    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4680    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4681    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4682    #[allow(clippy::too_many_arguments)]
4683    fn decode_step_t_core_stream(
4684        &self,
4685        e: &Engine,
4686        tokens: &[u32],
4687        pos0: usize,
4688        cache: &mut Cache,
4689        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4690        mut ckpt: Option<&mut VerifyCkpt>,
4691        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4692        pp_pipe: Option<bool>,
4693        vtok_dev: Option<&CudaSlice<u32>>,
4694        graphs: Option<&mut DsparkVerifyGraphs>,
4695    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4696        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4697        // exactly as the eager and batched steps do. This is the single funnel every verify
4698        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4699        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4700        // is untouched.
4701        //
4702        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4703        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4704        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4705        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4706        // or a placement whose PpNRt fails to build — so a config that would still walk the
4707        // whole trunk on one stream refuses instead of regressing 28x.
4708        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4709            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
4710                if vtok_dev.is_some() {
4711                    return Err(
4712                        "device-token dspark verify (slice-2 deferred readback) has no PP \
4713                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
4714                         route on one device"
4715                            .into(),
4716                    );
4717                }
4718                return self.decode_step_t_core_ppn(
4719                    e,
4720                    tokens,
4721                    pos0,
4722                    cache,
4723                    embd_dev,
4724                    ckpt.take(),
4725                    stream,
4726                    &fence,
4727                    pp_pipe,
4728                );
4729            }
4730        }
4731        crate::pp::refuse_unsplit_if_remote(
4732            "decode_step_t (spec verify)",
4733            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4734             split (decode_step_t_core_ppn); or run spec on one device",
4735        )?;
4736        let cfg = &self.cfg;
4737        let n_embd = cfg.n_embd as usize;
4738        let eps = cfg.rms_eps;
4739        let t = tokens.len();
4740        let pos_d = match stream {
4741            Some((_, ctr)) => {
4742                let mut p = e.alloc_uninit::<i32>(t)?;
4743                e.pos_iota(ctr, &mut p, t)?;
4744                p
4745            }
4746            None => {
4747                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4748                e.htod_i32(&pos_vec)?
4749            }
4750        };
4751
4752        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
4753        let x = match (stream, embd_dev) {
4754            (Some((vtok, _)), Some((g, qt, rb))) => {
4755                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4756            }
4757            (None, Some((g, qt, rb))) => match vtok_dev {
4758                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
4759                // bit-identical rows to the host-token arm (same per-dtype deq).
4760                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
4761                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4762            },
4763            _ => {
4764                assert!(
4765                    vtok_dev.is_none(),
4766                    "device-token verify requires the resident embed table (embd_dev)"
4767                );
4768                e.htod(&self.embd.gather(n_embd, tokens))?
4769            }
4770        };
4771
4772        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
4773        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
4774        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
4775        let x = self.verify_layers(
4776            e,
4777            x,
4778            0,
4779            self.layers.len(),
4780            &pos_d,
4781            pos0,
4782            t,
4783            cache,
4784            ckpt.take(),
4785            stream,
4786            graphs,
4787        )?;
4788
4789        let mut hn = vbuf(e, t * n_embd)?;
4790        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
4791        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
4792        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
4793        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
4794        let eager_tail = self.sliding_gated_moe_batch_program()
4795            && std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1");
4796        if eager_tail {
4797            let n_vocab = self.cfg.n_vocab as usize;
4798            let mut logits = vbuf(e, t * n_vocab)?;
4799            for r in 0..t {
4800                let mut row = e.uninit(n_embd)?;
4801                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4802                let mut hr = e.uninit(n_embd)?;
4803                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
4804                let lr = e.matmul(&self.output, &hr, 1)?;
4805                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
4806                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
4807            }
4808            if stream.is_none() {
4809                cache.pos += t;
4810            }
4811            return Ok((logits, if spec_hpost() { hn } else { x }));
4812        }
4813        let serving_head =
4814            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
4815        let logits = if serving_head {
4816            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
4817            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
4818            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
4819            // serve one batched numeric class at every live width, including B=1. Keep the
4820            // verify head in that same class; other generic families retain the decode-exact
4821            // head that their run-spec contract pins.
4822            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4823            e.matmul(&self.output, &hn, t)?
4824        } else {
4825            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4826            e.matmul_decode_exact(&self.output, &hn, t)?
4827        };
4828        // stream: the device pos counter owns position; host mirror reconciles at drain.
4829        if stream.is_none() {
4830            cache.pos += t;
4831        }
4832        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
4833        Ok((logits, if spec_hpost() { hn } else { x }))
4834    }
4835
4836    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
4837    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
4838    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
4839    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
4840    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
4841    /// the payload).
4842    ///
4843    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
4844    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
4845    /// receipts):
4846    ///
4847    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4848    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4849    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4850    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4851    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
4852    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4853    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4854    ///
4855    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4856    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4857    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4858    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4859    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
4860    ///
4861    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4862    ///    sharded loader leaves the table with stage 0 by construction).
4863    ///
4864    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4865    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4866    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4867    ///    model, every round.
4868    ///
4869    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4870    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4871    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4872    /// through the primary context by UVA — the same read the batched serving epilogue's
4873    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4874    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4875    ///
4876    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4877    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4878    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4879    ///
4880    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4881    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4882    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4883    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4884    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4885    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4886    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4887    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4888    #[allow(clippy::too_many_arguments)]
4889    fn decode_step_t_core_ppn(
4890        &self,
4891        e: &Engine,
4892        tokens: &[u32],
4893        pos0: usize,
4894        cache: &mut Cache,
4895        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4896        mut ckpt: Option<&mut VerifyCkpt>,
4897        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4898        fence: &[usize],
4899        pp_pipe: Option<bool>,
4900    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4901        let ticket = self.verify_stage0_issue(
4902            e,
4903            tokens,
4904            pos0,
4905            cache,
4906            embd_dev,
4907            ckpt.as_deref_mut(),
4908            stream,
4909            fence,
4910            pp_pipe,
4911            None,
4912        )?;
4913        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
4914    }
4915
4916    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
4917    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
4918    #[allow(clippy::too_many_arguments)]
4919    fn verify_stage0_issue(
4920        &self,
4921        e: &Engine,
4922        tokens: &[u32],
4923        pos0: usize,
4924        cache: &mut Cache,
4925        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4926        mut ckpt: Option<&mut VerifyCkpt>,
4927        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4928        fence: &[usize],
4929        pp_pipe: Option<bool>,
4930        trace: Option<SpecPipeTraceCtx>,
4931    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
4932        assert!(
4933            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
4934            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
4935             (the gemma4 arms have their own decode_step_t twins)"
4936        );
4937        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
4938            return Err(
4939                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
4940                 boundary itself is host-staged, but device-resident verify still peer-reads \
4941                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
4942                 serving on this host class; spec requires local per-stage inputs first."
4943                    .into(),
4944            );
4945        }
4946        let rt = crate::pp::PpNRt::get(e)?;
4947        let n_st = fence.len() - 1;
4948        assert_eq!(
4949            rt.n_stages(),
4950            n_st,
4951            "PpNRt stage count {} != fence stages {n_st}",
4952            rt.n_stages()
4953        );
4954        let n_embd = self.cfg.n_embd as usize;
4955        let t = tokens.len();
4956        let payload = t * n_embd;
4957        if pp_pipe.is_some() {
4958            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
4959        }
4960        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
4961        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
4962        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
4963        // the report below names exactly two stages and must never imply it measured middle ones.
4964        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
4965        let pp_started = std::time::Instant::now();
4966        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
4967        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
4968        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
4969        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
4970        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
4971        // stage stream and the wait would self-order into a no-op.
4972        let caller_stream = e.stream();
4973        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
4974        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
4975        // the primary stream still holds queued reads of them — with event tracking elided,
4976        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
4977        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
4978        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
4979        // stage stream behind the caller before enqueueing new stage work.
4980        let reverse_started = std::time::Instant::now();
4981        if pp_pipe != Some(false) {
4982            rt.fence_stages_behind(&caller_stream)?;
4983        }
4984        if pp_pipe == Some(true) {
4985            // Both session verifies must alternate boundary slots even when the ordinary
4986            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
4987            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
4988            rt.prepare_overlap_slots(0, payload)?;
4989        }
4990        if pp_anatomy {
4991            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
4992            // prices any primary-stream rollback/refresh tail inherited from the prior round.
4993            for s in 0..n_st {
4994                let _st = rt.enter(s);
4995                rt.engine(s, e).stream().synchronize()?;
4996            }
4997            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
4998        }
4999
5000        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5001        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5002        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5003            match stream {
5004                Some((_, ctr)) => {
5005                    let mut p = es.alloc_uninit::<i32>(t)?;
5006                    es.pos_iota(ctr, &mut p, t)?;
5007                    Ok(p)
5008                }
5009                None => {
5010                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5011                    es.htod_i32(&pos_vec)
5012                }
5013            }
5014        };
5015
5016        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5017        let slot = {
5018            let _st0 = rt.enter(0);
5019            let e0 = rt.engine(0, e);
5020            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5021            let stage0_started = std::time::Instant::now();
5022            let pos_d = stage_pos(e0)?;
5023            let x = match (stream, embd_dev) {
5024                (Some((vtok, _)), Some((g, qt, rb))) => {
5025                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5026                }
5027                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5028                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5029            };
5030            let x = self.verify_layers(
5031                e0,
5032                x,
5033                fence[0],
5034                fence[1],
5035                &pos_d,
5036                pos0,
5037                t,
5038                cache,
5039                ckpt.as_deref_mut(),
5040                stream,
5041                None,
5042            )?;
5043            if pp_anatomy {
5044                e0.stream().synchronize()?;
5045                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5046            }
5047            let tx_started = std::time::Instant::now();
5048            let slot = if pp_pipe.is_some() {
5049                rt.tx_pipelined(0, &x, payload)?
5050            } else {
5051                rt.tx(0, &x, payload)?
5052            };
5053            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5054            if pp_anatomy {
5055                e0.stream().synchronize()?;
5056                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5057            }
5058            slot
5059            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5060        };
5061
5062        Ok(VerifyBoundaryTicket {
5063            rt,
5064            caller_stream,
5065            slot,
5066            pos0,
5067            t,
5068            payload,
5069            n_st,
5070            pipelined: pp_pipe.is_some(),
5071            pp_anatomy,
5072            pp_started,
5073            reverse_ms,
5074            stage0_ms,
5075            tx_ms,
5076            trace,
5077        })
5078    }
5079
5080    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5081    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5082    #[allow(clippy::too_many_arguments)]
5083    fn verify_stage1_finish(
5084        &self,
5085        e: &Engine,
5086        ticket: VerifyBoundaryTicket,
5087        cache: &mut Cache,
5088        mut ckpt: Option<&mut VerifyCkpt>,
5089        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5090        fence: &[usize],
5091        publish_to_caller: bool,
5092    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5093        let VerifyBoundaryTicket {
5094            rt,
5095            caller_stream,
5096            slot,
5097            pos0,
5098            t,
5099            payload,
5100            n_st,
5101            pipelined,
5102            pp_anatomy,
5103            pp_started,
5104            reverse_ms,
5105            stage0_ms,
5106            tx_ms,
5107            trace,
5108        } = ticket;
5109        let n_embd = self.cfg.n_embd as usize;
5110        let eps = self.cfg.rms_eps;
5111        let mut slot = slot;
5112        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5113        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5114            match stream {
5115                Some((_, ctr)) => {
5116                    let mut p = es.alloc_uninit::<i32>(t)?;
5117                    es.pos_iota(ctr, &mut p, t)?;
5118                    Ok(p)
5119                }
5120                None => {
5121                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5122                    es.htod_i32(&pos_vec)
5123                }
5124            }
5125        };
5126
5127        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5128        for s in 1..n_st - 1 {
5129            let _st = rt.enter(s);
5130            let es = rt.engine(s, e);
5131            let pos_d = stage_pos(es)?;
5132            let x = rt.rx(s - 1, slot, payload)?;
5133            let x = self.verify_layers(
5134                es,
5135                x,
5136                fence[s],
5137                fence[s + 1],
5138                &pos_d,
5139                pos0,
5140                t,
5141                cache,
5142                ckpt.as_deref_mut(),
5143                stream,
5144                None,
5145            )?;
5146            slot = if pipelined {
5147                rt.tx_pipelined(s, &x, payload)?
5148            } else {
5149                rt.tx(s, &x, payload)?
5150            };
5151        }
5152
5153        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5154        let _stl = rt.enter(n_st - 1);
5155        let el = rt.engine(n_st - 1, e);
5156        let pos_d = stage_pos(el)?;
5157        let rx_started = std::time::Instant::now();
5158        let x = rt.rx(n_st - 2, slot, payload)?;
5159        if pp_anatomy {
5160            el.stream().synchronize()?;
5161            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5162        }
5163        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5164        let stage1_started = std::time::Instant::now();
5165        let x = self.verify_layers(
5166            el,
5167            x,
5168            fence[n_st - 1],
5169            fence[n_st],
5170            &pos_d,
5171            pos0,
5172            t,
5173            cache,
5174            ckpt.as_deref_mut(),
5175            stream,
5176            None,
5177        )?;
5178
5179        let mut hn = vbuf(el, payload)?;
5180        let logits = if self.sliding_gated_moe_batch_program() {
5181            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5182            // Verify must not switch numeric class merely because the same session speculates.
5183            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5184            el.matmul(&self.output, &hn, t)?
5185        } else {
5186            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5187            el.matmul_decode_exact(&self.output, &hn, t)?
5188        };
5189        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
5190        if pp_anatomy {
5191            el.stream().synchronize()?;
5192            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5193        }
5194        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5195        // stream. Order the caller's stream behind that work before the buffers escape this
5196        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5197        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5198        // the following arm's KV in the same process).
5199        if publish_to_caller {
5200            rt.publish_to(n_st - 1, &caller_stream)?;
5201        }
5202        if pp_anatomy {
5203            if publish_to_caller {
5204                caller_stream.synchronize()?;
5205            }
5206            eprintln!(
5207                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5208                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5209                pp_started.elapsed().as_secs_f64() * 1e3,
5210            );
5211        }
5212        // stream: the device pos counter owns position; host mirror reconciles at drain.
5213        if stream.is_none() {
5214            cache.pos += t;
5215        }
5216        Ok((logits, if spec_hpost() { hn } else { x }))
5217    }
5218
5219    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5220    ///
5221    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5222    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5223    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5224    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5225    /// bytes when a request moves from batched plain serving into speculative verify. Run the
5226    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5227    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5228    /// every norm/projection/FFN uses exactly the live serving dispatch.
5229    #[allow(clippy::too_many_arguments)]
5230    fn step35_verify_batch_layers(
5231        &self,
5232        e: &Engine,
5233        mut x: CudaSlice<f32>,
5234        lo: usize,
5235        hi: usize,
5236        pos0: usize,
5237        t: usize,
5238        cache: &mut Cache,
5239    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5240        let n_embd = self.cfg.n_embd as usize;
5241        if !self.uses_sliding_gated_moe_program() {
5242            return Err(
5243                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5244            );
5245        }
5246        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5247        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5248        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5249        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5250        // and the tap path keep the batch-layer class.
5251        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5252        let eager_verify = *VE
5253            .get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1"))
5254            && lo == 0
5255            && hi == self.layers.len();
5256        if eager_verify {
5257            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5258            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5259            // column runs the UNMODIFIED t=1 attention program via the col-select door and
5260            // the ordinary residual/FFN body. Values per column are bit-equal to the
5261            // row-outer walk: rms over the materialized residual == the fused add+norm
5262            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5263            // kernel, and every downstream op IS the t=1 program.
5264            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5265            let tcol =
5266                *TCOL.get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() == Ok("1"));
5267            if tcol && t >= 2 && t <= 8 {
5268                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5269                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5270                // syncs serialize the stream, so the split is for TARGETING amortization
5271                // work only — never a perf claim.
5272                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5273                let prof =
5274                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5275                let mut prof_ms = [0f64; 3];
5276                let eps = self.cfg.rms_eps;
5277                let mut x_t = x;
5278                let mut h_t = e.uninit(t * n_embd)?;
5279                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5280                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5281                // pageable htod was an in-stream engine turnaround x t x 45).
5282                let mut pos_rows = Vec::with_capacity(t);
5283                for r in 0..t {
5284                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5285                }
5286                let mut ok = true;
5287                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5288                // stashes `gated` instead of joining per column; one b4_tcol per rank +
5289                // one slab join produce every column's `mixed` after the attention pass.
5290                // Bit-exact per column (t=1 b4 program per column; elementwise join).
5291                // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
5292                // MoE layer deferred, the residual norm runs as one t-grid launch
5293                // (per-row program == t=1) and the FFN as ONE two-column device-routed
5294                // sweep + per-column shexp — the two columns' expert weights dedup
5295                // through L2 instead of reading HBM twice.
5296                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5297                let ffn_batch =
5298                    *FFN2.get_or_init(|| std::env::var("MEMRA_TCOL_FFN").as_deref() == Ok("1"));
5299                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5300                let mut mixed_row = e.uninit(n_embd)?;
5301                for il in lo..hi {
5302                    let layer = &self.layers[il];
5303                    let mut seg = std::time::Instant::now();
5304                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5305                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5306                        ok = false;
5307                        break;
5308                    }
5309                    if prof {
5310                        e.stream().synchronize()?;
5311                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5312                        seg = std::time::Instant::now();
5313                    }
5314                    let mut next = e.uninit(t * n_embd)?;
5315                    // Columns whose o_proj was deferred (their FFN runs after the join).
5316                    // A NON-deferred column's FFN must run INSIDE the column loop: the
5317                    // oproj-tail handoff is a single cell that the same column's
5318                    // residual_norm_ffn consumes before the next column's finish.
5319                    let mut deferred: Vec<usize> = Vec::new();
5320                    let mut ffn_col =
5321                        |r: usize,
5322                         mixed: &CudaSlice<f32>,
5323                         next: &mut CudaSlice<f32>|
5324                         -> Result<(), Box<dyn std::error::Error>> {
5325                            let mut x_row = e.uninit(n_embd)?;
5326                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5327                            let (x1, ffn_out) =
5328                                self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5329                            let mut x2 = e.uninit(n_embd)?;
5330                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5331                            e.dtod_copy_into(&x2, next, r * n_embd)?;
5332                            Ok(())
5333                        };
5334                    for r in 0..t {
5335                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5336                        let row_pos = &pos_rows[r];
5337                        crate::tp::set_verify_tcol(Some(r));
5338                        if oproj_batch {
5339                            crate::tp::set_tcol_oproj_defer(Some(r));
5340                        }
5341                        let mixed = match &layer.mixer {
5342                            crate::hybrid::Mixer::Full(fa) => {
5343                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5344                            }
5345                            _ => Err("step35 verify expects full attention".into()),
5346                        };
5347                        crate::tp::set_verify_tcol(None);
5348                        crate::tp::set_tcol_oproj_defer(None);
5349                        let mixed = mixed?;
5350                        if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5351                            deferred.push(r);
5352                        } else {
5353                            ffn_col(r, &mixed, &mut next)?;
5354                        }
5355                    }
5356                    if prof {
5357                        e.stream().synchronize()?;
5358                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5359                        seg = std::time::Instant::now();
5360                    }
5361                    if !deferred.is_empty() {
5362                        let mixed_t = self.step35_verify_oproj_tcol(e, il, t)?;
5363                        let o_out = mixed_t.len() / t;
5364                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5365                        // program == t=1; bit-identical to the oproj-tail join per the
5366                        // M2 verbatim-program contract) feeding the two-column routed
5367                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5368                        // to the per-column body.
5369                        let mut batched = false;
5370                        if ffn_batch && t == 2 && deferred.len() == t && o_out == n_embd {
5371                            let mut x1_t = e.uninit(t * n_embd)?;
5372                            let mut z_t = e.uninit(t * n_embd)?;
5373                            e.add_rms_norm(
5374                                &x_t,
5375                                &mixed_t,
5376                                layer.post_attn_norm.float_data(),
5377                                &mut x1_t,
5378                                &mut z_t,
5379                                n_embd,
5380                                t,
5381                                eps,
5382                            )?;
5383                            if let Some(ffn_t) = self.step35_verify_moe_t2(e, il, &z_t)? {
5384                                let mut x2_t = e.uninit(t * n_embd)?;
5385                                e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5386                                next = x2_t;
5387                                batched = true;
5388                            }
5389                        }
5390                        if !batched {
5391                            for &r in &deferred {
5392                                e.dtod_copy_view(
5393                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
5394                                    &mut mixed_row,
5395                                )?;
5396                                ffn_col(r, &mixed_row, &mut next)?;
5397                            }
5398                        }
5399                    }
5400                    if prof {
5401                        e.stream().synchronize()?;
5402                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5403                    }
5404                    drop(ffn_col);
5405                    x_t = next;
5406                }
5407                if prof {
5408                    eprintln!(
5409                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
5410                        prof_ms[0], prof_ms[1], prof_ms[2]
5411                    );
5412                }
5413                if ok {
5414                    return Ok(x_t);
5415                }
5416                // fall through to the row-outer walk on ineligible layers
5417                x = x_t;
5418            }
5419            let mut next = e.uninit(t * n_embd)?;
5420            for r in 0..t {
5421                let mut row = e.uninit(n_embd)?;
5422                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5423                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5424                let out = self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?;
5425                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5426            }
5427            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
5428            // row-outer walk does not materialize); the door is a step37 MTP bring-up
5429            // surface where taps are unused.
5430            return Ok(next);
5431        }
5432        let mut ph_last = std::time::Instant::now();
5433        for il in lo..hi {
5434            let mut next = e.uninit(t * n_embd)?;
5435            for r in 0..t {
5436                let mut row = e.uninit(n_embd)?;
5437                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5438                // The caller owns this verify's position. During controller overlap, cache.pos
5439                // still describes generation N while this stage-0 walk belongs to N+1.
5440                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5441                let mut one = [&mut *cache];
5442                let out = self.step35_decode_batch_layers(
5443                    e,
5444                    row,
5445                    &mut one,
5446                    &[(pos0 + r) as i32],
5447                    &row_pos,
5448                    il,
5449                    il + 1,
5450                    &mut ph_last,
5451                )?;
5452                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5453            }
5454            self.dflash_tap(e, cache, il, &next, t)?;
5455            x = next;
5456        }
5457        Ok(x)
5458    }
5459
5460    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
5461    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
5462    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
5463    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
5464    /// prefix-keep, not all-or-nothing).
5465    pub(crate) fn dspark_verify_t_am(
5466        &self,
5467        e: &Engine,
5468        tokens: &[u32],
5469        pos0: usize,
5470        cache: &mut Cache,
5471    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
5472        let (logits, _hn) = self.decode_step_t_core_stream(
5473            e, tokens, pos0, cache, None, None, None, None, None, None,
5474        )?;
5475        let t = tokens.len();
5476        let v = self.output.out_features();
5477        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5478        for r in 0..t {
5479            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5480        }
5481        Ok(e.dtoh_u32(&am_d)?)
5482    }
5483
5484    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
5485    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
5486    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
5487    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
5488    pub(crate) fn dspark_verify_t_logits(
5489        &self,
5490        e: &Engine,
5491        tokens: &[u32],
5492        pos0: usize,
5493        cache: &mut Cache,
5494    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5495        let (logits, _hn) = self.decode_step_t_core_stream(
5496            e, tokens, pos0, cache, None, None, None, None, None, None,
5497        )?;
5498        Ok(logits)
5499    }
5500
5501    /// DSpark verify with the MTP column-stash armed: identical forward to
5502    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
5503    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
5504    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
5505    pub(crate) fn dspark_verify_t_am_ckpt(
5506        &self,
5507        e: &Engine,
5508        tokens: &[u32],
5509        pos0: usize,
5510        cache: &mut Cache,
5511    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5512        let mut ck = VerifyCkpt::new(self.layers.len());
5513        let (logits, _hn) = self.decode_step_t_core_stream(
5514            e,
5515            tokens,
5516            pos0,
5517            cache,
5518            None,
5519            Some(&mut ck),
5520            None,
5521            None,
5522            None,
5523            None,
5524        )?;
5525        let t = tokens.len();
5526        let v = self.output.out_features();
5527        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5528        for r in 0..t {
5529            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5530        }
5531        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
5532    }
5533
5534    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
5535    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
5536    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
5537    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
5538    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
5539    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
5540    pub(crate) fn dspark_verify_t_am_ckpt_dev(
5541        &self,
5542        e: &Engine,
5543        vtok: &CudaSlice<u32>,
5544        t: usize,
5545        pos0: usize,
5546        cache: &mut Cache,
5547        embd_dev: (&CudaSlice<u8>, i32, usize),
5548        graphs: Option<&mut DsparkVerifyGraphs>,
5549    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5550        debug_assert!(
5551            vtok.len() >= t,
5552            "verify window exceeds the device token buffer"
5553        );
5554        // The slab flag is a per-round statement: clear it here so a verify that never
5555        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
5556        // stale `true` steering the commit at slabs the round never wrote.
5557        let mut graphs = graphs;
5558        if let Some(g) = graphs.as_deref_mut() {
5559            g.round_slab = false;
5560        }
5561        let mut ck = VerifyCkpt::new(self.layers.len());
5562        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
5563        // arm's established pattern — spec.rs stream-mode verify does the same).
5564        let dummy = vec![0u32; t];
5565        let (logits, _hn) = self.decode_step_t_core_stream(
5566            e,
5567            &dummy,
5568            pos0,
5569            cache,
5570            Some(embd_dev),
5571            Some(&mut ck),
5572            None,
5573            None,
5574            Some(vtok),
5575            graphs,
5576        )?;
5577        let v = self.output.out_features();
5578        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5579        for r in 0..t {
5580            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5581        }
5582        Ok((am_d, DsparkVerifyCkpt(ck)))
5583    }
5584
5585    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
5586    pub(crate) fn dspark_verify_t_logits_ckpt(
5587        &self,
5588        e: &Engine,
5589        tokens: &[u32],
5590        pos0: usize,
5591        cache: &mut Cache,
5592    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5593        let mut ck = VerifyCkpt::new(self.layers.len());
5594        let (logits, _hn) = self.decode_step_t_core_stream(
5595            e,
5596            tokens,
5597            pos0,
5598            cache,
5599            None,
5600            Some(&mut ck),
5601            None,
5602            None,
5603            None,
5604            None,
5605        )?;
5606        Ok((logits, DsparkVerifyCkpt(ck)))
5607    }
5608
5609    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
5610    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
5611    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
5612    pub(crate) fn dspark_commit_prefix(
5613        &self,
5614        e: &Engine,
5615        cache: &mut Cache,
5616        snap: &crate::cache::CacheSnapshot,
5617        ckpt: &DsparkVerifyCkpt,
5618        keep: usize,
5619    ) -> Result<(), Box<dyn std::error::Error>> {
5620        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
5621    }
5622
5623    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
5624    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
5625    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
5626    /// from the stash of column keep-1), slab-addressed and batched into two copy
5627    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
5628    pub(crate) fn dspark_commit_prefix_slab(
5629        &self,
5630        e: &Engine,
5631        cache: &mut Cache,
5632        snap: &crate::cache::CacheSnapshot,
5633        ctx: &DsparkVerifyGraphs,
5634        keep: usize,
5635    ) -> Result<(), Box<dyn std::error::Error>> {
5636        use cudarc::driver::DevicePtr;
5637        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
5638        let mut conv_src: Vec<u64> = Vec::new();
5639        let mut ssm_src: Vec<u64> = Vec::new();
5640        let mut conv_dst: Vec<u64> = Vec::new();
5641        let mut ssm_dst: Vec<u64> = Vec::new();
5642        for il in 0..self.layers.len() {
5643            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
5644                kvl.len = saved + keep;
5645                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
5646            }
5647            if let Some(rl) = cache.recur[il].as_ref() {
5648                let (pc, ps, _cw, _sw) = ctx
5649                    .slab_row(e, il, keep - 1)
5650                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
5651                conv_src.push(pc);
5652                ssm_src.push(ps);
5653                let st = &e.gpu.stream();
5654                let (dc, _g0) = rl.conv_state.device_ptr(st);
5655                let (ds, _g1) = rl.ssm_state.device_ptr(st);
5656                conv_dst.push(dc as u64);
5657                ssm_dst.push(ds as u64);
5658            }
5659        }
5660        let n = conv_src.len();
5661        if n > 0 {
5662            if state_copy_batch_on() {
5663                let mut tt = vec![0u64; 2 * n];
5664                tt[..n].copy_from_slice(&conv_src);
5665                tt[n..].copy_from_slice(&conv_dst);
5666                let ct = e.htod_u64(&tt)?;
5667                tt[..n].copy_from_slice(&ssm_src);
5668                tt[n..].copy_from_slice(&ssm_dst);
5669                let st = e.htod_u64(&tt)?;
5670                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
5671                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
5672            } else {
5673                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
5674                let row = keep - 1;
5675                for il in 0..self.layers.len() {
5676                    let Some(rl) = cache.recur[il].as_mut() else {
5677                        continue;
5678                    };
5679                    let k = ctx.lin_pos[&il];
5680                    {
5681                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
5682                        let win = sv.slice(row * cw..(row + 1) * cw);
5683                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
5684                    }
5685                    {
5686                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
5687                        let win = sv.slice(row * sw..(row + 1) * sw);
5688                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
5689                    }
5690                }
5691            }
5692        }
5693        cache.pos = snap.pos + keep;
5694        Ok(())
5695    }
5696
5697    /// Qwen35-family verify trunk in the live serving numeric class.
5698    ///
5699    /// Serving intentionally keeps this architecture in the generic batched program even at
5700    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
5701    ///
5702    /// Two arms, one numeric class:
5703    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
5704    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
5705    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
5706    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
5707    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
5708    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
5709    ///   program its isolated serving step would). One weight read per layer per round
5710    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
5711    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
5712    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
5713    ///   serving layer body, preserving single-session autoregressive cache order (the
5714    ///   correctness reference; also the rollback seam for the t-parallel arm).
5715    ///
5716    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
5717    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
5718    #[allow(clippy::too_many_arguments)]
5719    fn qwen35_verify_batch_layers(
5720        &self,
5721        e: &Engine,
5722        x: CudaSlice<f32>,
5723        lo: usize,
5724        hi: usize,
5725        pos0: usize,
5726        t: usize,
5727        cache: &mut Cache,
5728        ckpt: Option<&mut VerifyCkpt>,
5729        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5730        graphs: Option<&mut DsparkVerifyGraphs>,
5731    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5732        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
5733        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
5734        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
5735        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
5736        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
5737        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
5738        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
5739            || !self.batched_serving_numeric_class()
5740            || t > 16;
5741        if rowwise {
5742            if stream.is_some() {
5743                // rowwise replays per row with host cache.pos — irreconcilable with a
5744                // device position counter. Burst callers must keep t <= 16 and the
5745                // ROWWISE env unset; refusing beats silently mispositioned rows.
5746                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
5747                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
5748                    .into());
5749            }
5750            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
5751        } else {
5752            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
5753        }
5754    }
5755
5756    /// The per-row correctness reference: replay each verify row through the authoritative
5757    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
5758    #[allow(clippy::too_many_arguments)]
5759    fn qwen35_verify_rowwise(
5760        &self,
5761        e: &Engine,
5762        mut x: CudaSlice<f32>,
5763        lo: usize,
5764        hi: usize,
5765        pos0: usize,
5766        t: usize,
5767        cache: &mut Cache,
5768        mut ckpt: Option<&mut VerifyCkpt>,
5769    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5770        let n_embd = self.cfg.n_embd as usize;
5771        let saved_pos = cache.pos;
5772        let mut ph_last = std::time::Instant::now();
5773        for il in lo..hi {
5774            let mut next = e.uninit(t * n_embd)?;
5775            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
5776                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5777                    Some(Vec::with_capacity(t - 1))
5778                } else {
5779                    None
5780                };
5781            for r in 0..t {
5782                cache.pos = pos0 + r;
5783                let mut row = e.uninit(n_embd)?;
5784                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5785                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5786                let mut one = [&mut *cache];
5787                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
5788                let out = match self.decode_batch_layers(
5789                    e,
5790                    row,
5791                    &mut one,
5792                    &ctx,
5793                    &row_pos,
5794                    &mut ph_last,
5795                ) {
5796                    Ok(out) => out,
5797                    Err(error) => {
5798                        cache.pos = saved_pos;
5799                        return Err(error);
5800                    }
5801                };
5802                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5803                if r + 1 < t {
5804                    if let Some(states) = col_states.as_mut() {
5805                        let recur = cache.recur[il]
5806                            .as_ref()
5807                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
5808                        states.push((
5809                            e.clone_dtod(&recur.conv_state)?,
5810                            e.clone_dtod(&recur.ssm_state)?,
5811                        ));
5812                    }
5813                }
5814            }
5815            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
5816                checkpoint.cols[il] = Some(states);
5817            }
5818            x = next;
5819        }
5820        cache.pos = saved_pos;
5821        Ok(x)
5822    }
5823
5824    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
5825    ///
5826    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
5827    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
5828    /// pins the serving batch tier already carries:
5829    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
5830    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
5831    ///     alone;
5832    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
5833    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
5834    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
5835    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
5836    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
5837    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
5838    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
5839    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
5840    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
5841    /// program its isolated B=1 serving step would.
5842    ///
5843    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
5844    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
5845    #[allow(clippy::too_many_arguments)]
5846    fn qwen35_verify_tparallel(
5847        &self,
5848        e: &Engine,
5849        mut x: CudaSlice<f32>,
5850        lo: usize,
5851        hi: usize,
5852        pos0: usize,
5853        t: usize,
5854        cache: &mut Cache,
5855        mut ckpt: Option<&mut VerifyCkpt>,
5856        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5857        mut graphs: Option<&mut DsparkVerifyGraphs>,
5858    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5859        let seqs_append =
5860            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
5861        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
5862
5863        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
5864        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
5865        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
5866        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
5867        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
5868        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
5869        // full-verify bodies).
5870        if stream.is_some() && graphs.is_some() {
5871            return Err(
5872                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
5873                        cannot arm together"
5874                    .into(),
5875            );
5876        }
5877        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
5878        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
5879        // moves the kv caches). Then:
5880        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
5881        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
5882        //    full-verify graph per (vt, rung) — linear layers through the shared
5883        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
5884        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
5885        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
5886        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
5887        //    the full-attention layers run eager (batched rows when eligible).
5888        if let Some(g) = graphs.as_deref_mut() {
5889            g.refresh_tables(e, cache)?;
5890            g.round_slab = false;
5891            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
5892                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
5893                // full capture past the ceiling falls through to the segment/eager arms.
5894                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
5895                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
5896                    g.round_slab = true;
5897                    return Ok(out);
5898                }
5899            }
5900            // Round-atomic ceiling check for the segment door: if any linear run in this
5901            // walk would need a NEW capture past the ceiling, the whole round runs the
5902            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
5903            // would corrupt the commit).
5904            if !g.segments_ready(self, lo, hi, t) {
5905                graphs = None;
5906            }
5907        }
5908        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
5909        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
5910        let pos_d = match stream {
5911            Some((_, ctr)) => {
5912                let mut p = e.alloc_uninit::<i32>(t)?;
5913                e.pos_iota(ctr, &mut p, t)?;
5914                p
5915            }
5916            None => {
5917                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
5918                e.htod_i32(&pos_host)?
5919            }
5920        };
5921        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
5922        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
5923        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
5924        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
5925        // rides the dc rows kernels and never reaches the fallback).
5926        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
5927        let mut il = lo;
5928        while il < hi {
5929            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
5930                let mut end = il;
5931                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
5932                    end += 1;
5933                }
5934                let g = graphs.as_deref_mut().expect("checked above");
5935                x = g.run_segment(self, e, il, end, &x, t, cache)?;
5936                g.round_slab = true;
5937                il = end;
5938                continue;
5939            }
5940            let layer = &self.layers[il];
5941            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
5942                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
5943                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
5944                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
5945                x = self.qwen35_tparallel_linear_layer(
5946                    e,
5947                    il,
5948                    &x,
5949                    t,
5950                    cache,
5951                    ckpt.as_deref_mut(),
5952                    None,
5953                    None,
5954                )?;
5955                il += 1;
5956                continue;
5957            }
5958            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
5959            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
5960            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
5961            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
5962            // run (lane/draftcost-moe).
5963            x = self.qwen35_tparallel_fa_layer(
5964                e,
5965                il,
5966                &x,
5967                t,
5968                cache,
5969                FaLayerArgs {
5970                    pos_d: &pos_d,
5971                    pos_rows: &mut pos_rows,
5972                    pos0,
5973                    seqs_append,
5974                    batch_fa_on,
5975                    graph_cap: None,
5976                    stream,
5977                    ckpt: ckpt.as_deref_mut(),
5978                },
5979            )?;
5980            il += 1;
5981        }
5982        Ok(x)
5983    }
5984
5985    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
5986    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
5987    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
5988    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
5989    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
5990    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
5991    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
5992    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
5993    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
5994    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
5995    /// original singles chain, byte-for-byte.
5996    #[allow(clippy::too_many_arguments)]
5997    fn qwen35_tparallel_dense_ffn(
5998        &self,
5999        e: &Engine,
6000        ffn_gate: &crate::model::GpuTensor,
6001        ffn_up: &crate::model::GpuTensor,
6002        ffn_down: &crate::model::GpuTensor,
6003        zn: &CudaSlice<f32>,
6004        t: usize,
6005        n_embd: usize,
6006    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6007        let n_ff = ffn_gate.out_features();
6008        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6009        if Engine::tk_ffn_dual_on() {
6010            if let Some(((g, gs), (u, us))) =
6011                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6012            {
6013                if e.uses_q8_1_fast(ffn_down) {
6014                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6015                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6016                }
6017                let mut act = e.uninit(t * n_ff)?;
6018                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6019                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6020                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6021            }
6022        }
6023        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6024        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6025        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6026        let mut act = e.uninit(t * n_ff)?;
6027        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6028        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6029        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6030    }
6031
6032    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6033    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6034    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6035    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6036    ///
6037    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6038    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6039    ///   generation's cache lands at new addresses that only the per-verify table refresh
6040    ///   knows — the slice-3 baked-address lesson);
6041    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6042    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6043    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6044    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
6045    ///   round whose rows all sit inside the rung;
6046    /// - the host len bump moves to the replay caller (captured host code does not
6047    ///   re-run at replay).
6048    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6049    /// host-branches on t_kv and must never be captured.
6050    #[allow(clippy::too_many_arguments)]
6051    fn qwen35_tparallel_fa_layer(
6052        &self,
6053        e: &Engine,
6054        il: usize,
6055        x: &CudaSlice<f32>,
6056        t: usize,
6057        cache: &mut Cache,
6058        args: FaLayerArgs<'_>,
6059    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6060        use cudarc::driver::DevicePtr;
6061        let cfg = &self.cfg;
6062        let n_embd = cfg.n_embd as usize;
6063        let eps = cfg.rms_eps;
6064        let head_dim_global = cfg.head_dim_k as usize;
6065        let layer = &self.layers[il];
6066        let FaLayerArgs {
6067            pos_d,
6068            pos_rows,
6069            pos0,
6070            seqs_append,
6071            batch_fa_on,
6072            graph_cap,
6073            stream,
6074            mut ckpt,
6075        } = args;
6076
6077        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6078        let anorm = layer.attn_norm.float_data();
6079        let mut xn = e.uninit(t * n_embd)?;
6080        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6081        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6082
6083        let mixed: CudaSlice<f32> = match &layer.mixer {
6084            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6085            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6086            // per-row serving-kernel chain cannot run (host state swaps keyed on host
6087            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6088            // rebuild — the per-row chain only produces per-column clones). GDN rides
6089            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6090            // and its one-scan recurrence is pinned bit-identical to T chained T=1
6091            // steps (its header + kernel-check). Position-independent, so no counter
6092            // plumbing is needed. Guards mirror the generic call site exactly.
6093            Mixer::Linear(la) if stream.is_some() => {
6094                if !(t >= 3 || (t == 2 && spec_m2()))
6095                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6096                    || !e.uses_q8_1_fast(&la.ssm_out)
6097                {
6098                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6099                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6100                        .into());
6101                }
6102                let want = ckpt.is_some();
6103                let (out, stash) =
6104                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6105                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6106                    ck.gdn[il] = Some(st);
6107                }
6108                out
6109            }
6110            Mixer::Linear(_) => {
6111                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6112            }
6113            Mixer::Full(fa) => {
6114                let geometry = cfg.full_attention_geometry_at(il as u32);
6115                let n_head = geometry.n_head as usize;
6116                let n_head_kv = geometry.n_head_kv as usize;
6117                let head_dim = geometry.head_dim_k as usize;
6118                let rope_dims = geometry.n_rot as usize;
6119                let rope_base = geometry.rope_base;
6120                let scale = geometry.attention_scale();
6121                // Batched projections: one weight read serves all T rows.
6122                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6123                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6124                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6125                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6126                    [&fa.wq, &fa.wk, &fa.wv],
6127                    &hq,
6128                    &hd,
6129                    t,
6130                )? {
6131                    Some(mut g3) => {
6132                        let v = g3.pop().unwrap();
6133                        let k = g3.pop().unwrap();
6134                        let qf = g3.pop().unwrap();
6135                        (qf, k, v)
6136                    }
6137                    None => (
6138                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6139                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6140                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6141                    ),
6142                };
6143                let gated =
6144                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6145                let (mut q, gate) = if gated {
6146                    let mut qs = e.uninit(t * n_head * head_dim)?;
6147                    let mut gs = e.uninit(t * n_head * head_dim)?;
6148                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6149                    (qs, Some(gs))
6150                } else {
6151                    (qf, None)
6152                };
6153                let mut qn = e.uninit(t * n_head * head_dim)?;
6154                e.rms_norm(
6155                    &q,
6156                    fa.q_norm.float_data(),
6157                    &mut qn,
6158                    head_dim,
6159                    t * n_head,
6160                    eps,
6161                )?;
6162                q = qn;
6163                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6164                e.rms_norm(
6165                    &k,
6166                    fa.k_norm.float_data(),
6167                    &mut kn,
6168                    head_dim,
6169                    t * n_head_kv,
6170                    eps,
6171                )?;
6172                k = kn;
6173                e.rope_neox(
6174                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6175                )?;
6176                e.rope_neox(
6177                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6178                )?;
6179
6180                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6181                // draft), each through the b_n=1 serving kernels at its own t_kv.
6182                let q_dim = n_head * head_dim;
6183                let kv_dim = n_head_kv * head_dim;
6184                let mut attn = e.uninit(t * q_dim)?;
6185                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6186                    let kvl = cache.kv[il].as_ref().unwrap();
6187                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
6188                    // the batched twins; the per-row fallback reads pair 0 (same cache
6189                    // for every row of one layer). Graph mode reads the ctx table.
6190                    let local: Option<CudaSlice<u64>> = match graph_cap {
6191                        Some(_) => None,
6192                        None => {
6193                            let s = &e.gpu.stream();
6194                            let (pk, _g) = kvl.k.device_ptr(s);
6195                            let (pv, _g2) = kvl.v.device_ptr(s);
6196                            let mut tbl = Vec::with_capacity(2 * t);
6197                            for _ in 0..t {
6198                                tbl.push(pk as u64);
6199                                tbl.push(pv as u64);
6200                            }
6201                            Some(e.htod_u64(&tbl)?)
6202                        }
6203                    };
6204                    (
6205                        kvl.kv_dim_k,
6206                        kvl.kv_dim_v,
6207                        kvl.k_tok_bytes,
6208                        kvl.v_tok_bytes,
6209                        kvl.len,
6210                        local,
6211                    )
6212                };
6213                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6214                    Some((tb, off, _)) => (tb, off),
6215                    None => (kv_local.as_ref().expect("built above"), 0),
6216                };
6217                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6218                // section batches into the z-batched serving twins when every row of
6219                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6220                // guards are evaluated at the round's FIRST and LAST t_kv — the
6221                // eligibility window (vec floor .. v4 max) and each split-ladder rung
6222                // are intervals in t_kv, so ends-inside means all-inside (the straddle
6223                // law). Appending all T rows before any attend is read-equivalent to
6224                // the interleaved order: row r's walk reads keys 0..len0+r only, and
6225                // rows > r land at slots it never touches; every written cache row is
6226                // the per-token appender's exact warp program (kernel-check pinned).
6227                let t_kv_first = len0 + 1;
6228                let t_kv_last = len0 + t;
6229                let rows_batched = t >= 2
6230                    && seqs_append
6231                    && batch_fa_on
6232                    && dspark_fa_rows_on()
6233                    // the z-batched twins read stacked rows at the CACHE's kv dims;
6234                    // the projection stack is [T, n_head_kv*head_dim] — they must be
6235                    // the same stride or row z misaligns (true for this family; the
6236                    // guard keeps any asymmetric-kv model on the per-row loop).
6237                    && kdk == kv_dim
6238                    && kdv == kv_dim
6239                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6240                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6241                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6242                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6243                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6244                // grid only — bytes proven equal above). Capture-time invariants refuse
6245                // loudly rather than bake a divergent body.
6246                let (size_kv_max, sp) = match graph_cap {
6247                    Some((_, _, rung)) => {
6248                        if !rows_batched {
6249                            return Err(format!(
6250                                "fa graph capture: layer {il} round is not batchable \
6251                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6252                                 must never be captured"
6253                            )
6254                            .into());
6255                        }
6256                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6257                        if t_kv_last > rung
6258                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6259                        {
6260                            return Err(format!(
6261                                "fa graph capture: rung {rung} does not cover round \
6262                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6263                            )
6264                            .into());
6265                        }
6266                        (rung, sp_r)
6267                    }
6268                    None => (
6269                        t_kv_last,
6270                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6271                    ),
6272                };
6273                if let Some((_, ctr)) = stream {
6274                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6275                    // — the generic stream arm's exact shape (rows kernels are pinned
6276                    // byte-identical to the per-row programs by kernel-check). Host len
6277                    // stays a stale lower bound; the burst drain reconciles it.
6278                    let kvl = cache.kv[il].as_mut().unwrap();
6279                    e.append_kv_quantized_rows_dc(
6280                        &k,
6281                        &v,
6282                        &mut kvl.k,
6283                        &mut kvl.v,
6284                        ctr,
6285                        t,
6286                        kdk,
6287                        kdv,
6288                        ktb,
6289                        vtb,
6290                        Engine::kv_fp8_on(),
6291                    )?;
6292                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
6293                    let k_view = e.view_u8(&kvl.k, upper * ktb);
6294                    let v_view = e.view_u8(&kvl.v, upper * vtb);
6295                    e.fa_decode_rows_dc(
6296                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6297                        t, scale, ktb, vtb, 0, false,
6298                    )?;
6299                } else if rows_batched {
6300                    e.append_kv_quantized_seqs(
6301                        &k,
6302                        &v,
6303                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6304                        pos_d,
6305                        t,
6306                        kdk,
6307                        kdv,
6308                        ktb,
6309                        vtb,
6310                    )?;
6311                    if graph_cap.is_none() {
6312                        cache.kv[il].as_mut().unwrap().len += t;
6313                    }
6314                    e.fa_decode_batch_seqs_v4(
6315                        &q,
6316                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6317                        pos_d,
6318                        &mut attn,
6319                        head_dim,
6320                        n_head,
6321                        n_head_kv,
6322                        t,
6323                        size_kv_max,
6324                        scale,
6325                        sp,
6326                        ktb,
6327                        vtb,
6328                    )?;
6329                } else {
6330                    if pos_rows.is_none() {
6331                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
6332                        // the dc rows kernels above and never reaches this fallback).
6333                        *pos_rows = Some(match stream {
6334                            Some((_, ctr)) => (0..t)
6335                                .map(|r| {
6336                                    let mut b = e.alloc_uninit::<i32>(1)?;
6337                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
6338                                    Ok(b)
6339                                })
6340                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6341                            None => (0..t)
6342                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6343                                .collect::<Result<_, _>>()?,
6344                        });
6345                    }
6346                    let pos_rows = pos_rows.as_ref().unwrap();
6347                    for r in 0..t {
6348                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6349                        // whose row 0 is this row (arithmetic-free materialization copies,
6350                        // same as decode's per-seq fallback arm).
6351                        let mut k_row = e.uninit(kv_dim)?;
6352                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6353                        let mut v_row = e.uninit(kv_dim)?;
6354                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6355                        let pos_row = &pos_rows[r];
6356                        let kvl = cache.kv[il].as_mut().unwrap();
6357                        if seqs_append {
6358                            e.append_kv_quantized_seqs(
6359                                &k_row,
6360                                &v_row,
6361                                &kv_tbl.slice(kv_off..kv_off + 2),
6362                                pos_row,
6363                                1,
6364                                kdk,
6365                                kdv,
6366                                ktb,
6367                                vtb,
6368                            )?;
6369                            kvl.len += 1;
6370                        } else {
6371                            e.append_kv_quantized_view(
6372                                &k_row.slice(0..kv_dim),
6373                                &v_row.slice(0..kv_dim),
6374                                &mut kvl.k,
6375                                &mut kvl.v,
6376                                kvl.len,
6377                                kvl.kv_dim_k,
6378                                kvl.kv_dim_v,
6379                                kvl.k_tok_bytes,
6380                                kvl.v_tok_bytes,
6381                                Engine::kv_fp8_on(),
6382                            )?;
6383                            kvl.len += 1;
6384                        }
6385                        let t_kv = kvl.len;
6386                        let mut q_row = e.uninit(q_dim)?;
6387                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6388                        let mut a_row = e.uninit(q_dim)?;
6389                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
6390                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6391                            e.fa_decode_batch_seqs_v4(
6392                                &q_row,
6393                                &kv_tbl.slice(kv_off..kv_off + 2),
6394                                pos_row,
6395                                &mut a_row,
6396                                head_dim,
6397                                n_head,
6398                                n_head_kv,
6399                                1,
6400                                t_kv,
6401                                scale,
6402                                sp0_r,
6403                                ktb,
6404                                vtb,
6405                            )?;
6406                        } else {
6407                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6408                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6409                            let mut a_view = a_row.slice_mut(0..q_dim);
6410                            e.fa_decode_kvmod_view(
6411                                &q_row.slice(0..q_dim),
6412                                &k_view,
6413                                &v_view,
6414                                &mut a_view,
6415                                head_dim,
6416                                n_head,
6417                                n_head_kv,
6418                                t_kv,
6419                                scale,
6420                                kvl.k_tok_bytes,
6421                                kvl.v_tok_bytes,
6422                                Engine::kv_fp8_on(),
6423                            )?;
6424                        }
6425                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
6426                    }
6427                }
6428
6429                // Output gate (element-wise) + o-proj at m=T.
6430                let attn_g = match &gate {
6431                    Some(g) => {
6432                        let n = t * q_dim;
6433                        let mut gsig = e.uninit(n)?;
6434                        e.sigmoid(g, &mut gsig, n)?;
6435                        let mut ag = e.uninit(n)?;
6436                        e.mul(&attn, &gsig, &mut ag, n)?;
6437                        ag
6438                    }
6439                    None => attn,
6440                };
6441                e.matmul(&fa.wo, &attn_g, t)?
6442            }
6443        };
6444
6445        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6446        let pnorm = layer.post_attn_norm.float_data();
6447        let mut x1 = e.uninit(t * n_embd)?;
6448        let mut zn = e.uninit(t * n_embd)?;
6449        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6450        let ffn_out = match &layer.ffn {
6451            crate::hybrid::Ffn::Dense {
6452                ffn_gate,
6453                ffn_up,
6454                ffn_down,
6455            } => {
6456                assert!(
6457                    self.cfg.m3.is_none(),
6458                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6459                );
6460                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6461            }
6462            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6463        };
6464        let mut x2 = e.uninit(t * n_embd)?;
6465        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6466        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6467        self.dflash_tap(e, cache, il, &x2, t)?;
6468        Ok(x2)
6469    }
6470
6471    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
6472    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
6473    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
6474    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
6475    /// bit-identical by construction:
6476    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
6477    ///   the device sequence is driven entirely by the 6-entry pointer table, which
6478    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
6479    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
6480    ///   legacy post-swap clone read.
6481    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
6482    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
6483    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
6484    /// None builds the per-verify table exactly as before.
6485    #[allow(clippy::too_many_arguments)]
6486    fn qwen35_tparallel_linear_layer(
6487        &self,
6488        e: &Engine,
6489        il: usize,
6490        x: &CudaSlice<f32>,
6491        t: usize,
6492        cache: &mut Cache,
6493        mut ckpt: Option<&mut VerifyCkpt>,
6494        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
6495        table_src: Option<(&CudaSlice<u64>, usize)>,
6496    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6497        use cudarc::driver::DevicePtr;
6498        let cfg = &self.cfg;
6499        let n_embd = cfg.n_embd as usize;
6500        let eps = cfg.rms_eps;
6501        let layer = &self.layers[il];
6502        let Mixer::Linear(la) = &layer.mixer else {
6503            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
6504        };
6505        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6506        let anorm = layer.attn_norm.float_data();
6507        let mut xn = e.uninit(t * n_embd)?;
6508        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6509        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6510
6511        let geometry = la.geometry;
6512        let d_state = geometry.key_head_dim as usize;
6513        let num_k = geometry.key_heads as usize;
6514        let num_v = geometry.value_heads as usize;
6515        let d_conv = geometry.conv_kernel as usize;
6516        let key_dim = d_state * num_k;
6517        let value_dim = geometry.value_head_dim as usize * num_v;
6518        let conv_dim = key_dim * 2 + value_dim;
6519        let gdn_scale = 1.0 / (d_state as f32).sqrt();
6520
6521        // ---- batched projections: one weight read for all T rows ----
6522        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
6523        // per (tensor, token, row) to the four singles; refused (layout/tier) or
6524        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
6525        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
6526            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
6527            &hq,
6528            &hd,
6529            t,
6530        )? {
6531            Some(mut g4) => {
6532                let alpha = g4.pop().unwrap();
6533                let beta_raw = g4.pop().unwrap();
6534                let z = g4.pop().unwrap();
6535                let qkv_mixed = g4.pop().unwrap();
6536                (qkv_mixed, z, beta_raw, alpha)
6537            }
6538            None => (
6539                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
6540                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
6541                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
6542                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
6543            ),
6544        };
6545        let beta_w = la.ssm_beta.out_features();
6546        let alpha_w = la.ssm_alpha.out_features();
6547        let qkv_w = la.wqkv.out_features();
6548
6549        // ---- per-row state chain through the b_n=1 serving kernels ----
6550        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
6551        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
6552        let table_local: Option<CudaSlice<u64>> = match table_src {
6553            Some(_) => None,
6554            None => {
6555                let rl = cache.recur[il].as_ref().unwrap();
6556                let s = &e.gpu.stream();
6557                let (pc, _g0) = rl.conv_state.device_ptr(s);
6558                let (p0, _g1) = rl.ssm_state.device_ptr(s);
6559                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
6560                Some(e.htod_u64(&[
6561                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
6562                ])?)
6563            }
6564        };
6565        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
6566            Some((tb, off)) => (tb, off),
6567            None => (table_local.as_ref().unwrap(), 0),
6568        };
6569        let mut o_all = e.uninit(t * value_dim)?;
6570        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6571            if ckpt.is_some() && stash.is_none() && t >= 2 {
6572                Some(Vec::with_capacity(t - 1))
6573            } else {
6574                None
6575            };
6576        let mut stash = stash;
6577        // Per-row scratch reused across rows (uninit is cheap but not free at
6578        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
6579        // [T, ...] buffers — zero arithmetic-free copies in this loop.
6580        let mut conv_out = e.uninit(conv_dim)?;
6581        let mut q_l2 = e.uninit(value_dim)?;
6582        let mut k_l2 = e.uninit(value_dim)?;
6583        let mut v_gd = e.uninit(value_dim)?;
6584        let mut beta_b = e.uninit(num_v)?;
6585        let mut g_log = e.uninit(num_v)?;
6586        for r in 0..t {
6587            let base = toff + if r % 2 == 0 { 0 } else { 3 };
6588            let conv_view = table.slice(base..base + 1);
6589            let in_view = table.slice(base + 1..base + 2);
6590            let out_view = table.slice(base + 2..base + 3);
6591            e.ssm_conv1d_fused_decode_b_view(
6592                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
6593                &conv_view,
6594                la.ssm_conv1d.float_data(),
6595                &mut conv_out,
6596                conv_dim,
6597                d_conv,
6598                1,
6599            )?;
6600            e.gdn_prep_decode_b_view(
6601                &conv_out,
6602                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
6603                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
6604                la.ssm_dt.float_data(),
6605                la.ssm_a.float_data(),
6606                &mut q_l2,
6607                &mut k_l2,
6608                &mut v_gd,
6609                &mut beta_b,
6610                &mut g_log,
6611                d_state,
6612                num_v,
6613                num_k,
6614                key_dim,
6615                eps,
6616                conv_dim,
6617                1,
6618            )?;
6619            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
6620            e.gdn_scan_s128_batched_view(
6621                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
6622                gdn_scale,
6623            )?;
6624            if r + 1 < t {
6625                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
6626                // odd rows write s0 — the same physical state the legacy post-swap
6627                // canonical clone read.
6628                let rl = cache.recur[il]
6629                    .as_ref()
6630                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
6631                let ssm_src = if r % 2 == 0 {
6632                    &rl.ssm_state_alt
6633                } else {
6634                    &rl.ssm_state
6635                };
6636                match stash.as_mut() {
6637                    Some((conv_slab, ssm_slab)) => {
6638                        // BOTH stash reads go through the pointer table at run time: the
6639                        // ssm handles ping-pong between rounds, and the ctx (with its
6640                        // captured graphs) outlives the Cache — a fresh generation's
6641                        // conv/ssm buffers land at new addresses that only the per-round
6642                        // table refresh knows. A baked direct copy would read freed
6643                        // memory (parity was the slice-3 smoke divergence; cache
6644                        // lifetime is the cross-generation twin).
6645                        e.copy_indirect_src_f32(
6646                            &conv_view,
6647                            conv_slab,
6648                            r * conv_dim * (d_conv - 1),
6649                            conv_dim * (d_conv - 1),
6650                        )?;
6651                        // The ssm handles PING-PONG between rounds: a captured direct
6652                        // copy would bake the capture-time physical buffer and read the
6653                        // wrong parity after any odd-vt round (the slice-3 smoke
6654                        // divergence). Read the src address from row r's OUT table
6655                        // entry at run time — the same entry the scan just wrote.
6656                        e.copy_indirect_src_f32(
6657                            &out_view,
6658                            ssm_slab,
6659                            r * d_state * d_state * num_v,
6660                            d_state * d_state * num_v,
6661                        )?;
6662                    }
6663                    None => {
6664                        if let Some(states) = col_states.as_mut() {
6665                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
6666                        }
6667                    }
6668                }
6669            }
6670        }
6671        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
6672        // handle motion is identical and the device sequence never read the handles.
6673        if t % 2 == 1 {
6674            let rl = cache.recur[il].as_mut().unwrap();
6675            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
6676        }
6677        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6678            checkpoint.cols[il] = Some(states);
6679        }
6680
6681        // ---- batched gated norm + out-projection at m=T ----
6682        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
6683            let (gq, gd) = e.gated_rmsnorm_q8_1(
6684                &o_all,
6685                la.ssm_norm.float_data(),
6686                &z,
6687                d_state,
6688                t * num_v,
6689                eps,
6690            )?;
6691            let g0 = e.zeros(0)?;
6692            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
6693        } else {
6694            let mut gn = e.uninit(t * value_dim)?;
6695            e.gated_rmsnorm(
6696                &o_all,
6697                la.ssm_norm.float_data(),
6698                &z,
6699                &mut gn,
6700                d_state,
6701                t * num_v,
6702                eps,
6703            )?;
6704            e.matmul(&la.ssm_out, &gn, t)?
6705        };
6706
6707        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6708        let pnorm = layer.post_attn_norm.float_data();
6709        let mut x1 = e.uninit(t * n_embd)?;
6710        let mut zn = e.uninit(t * n_embd)?;
6711        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6712        let ffn_out = match &layer.ffn {
6713            crate::hybrid::Ffn::Dense {
6714                ffn_gate,
6715                ffn_up,
6716                ffn_down,
6717            } => {
6718                assert!(
6719                    self.cfg.m3.is_none(),
6720                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6721                );
6722                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6723            }
6724            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6725        };
6726        let mut x2 = e.uninit(t * n_embd)?;
6727        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6728        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6729        self.dflash_tap(e, cache, il, &x2, t)?;
6730        Ok(x2)
6731    }
6732
6733    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
6734    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
6735    /// carried in from outside the range) and exits with the range's final residual materialized
6736    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
6737    /// instead of one.
6738    ///
6739    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
6740    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
6741    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
6742    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
6743    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
6744    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
6745    /// code — there is no "split version" of the verify math.
6746    ///
6747    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
6748    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
6749    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
6750    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
6751    #[allow(clippy::too_many_arguments)]
6752    fn verify_layers(
6753        &self,
6754        e: &Engine,
6755        mut x: CudaSlice<f32>,
6756        lo: usize,
6757        hi: usize,
6758        pos_d: &CudaSlice<i32>,
6759        pos0: usize,
6760        t: usize,
6761        cache: &mut Cache,
6762        mut ckpt: Option<&mut VerifyCkpt>,
6763        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6764        graphs: Option<&mut DsparkVerifyGraphs>,
6765    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6766        if self.sliding_gated_moe_batch_program() {
6767            if stream.is_some() {
6768                return Err(
6769                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
6770                            cannot express the SWA offset KV view)"
6771                        .into(),
6772                );
6773            }
6774            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
6775        }
6776        if self.batched_serving_numeric_class() {
6777            return self.qwen35_verify_batch_layers(
6778                e,
6779                x,
6780                lo,
6781                hi,
6782                pos0,
6783                t,
6784                cache,
6785                ckpt.take(),
6786                stream,
6787                graphs,
6788            );
6789        }
6790        let n_embd = self.cfg.n_embd as usize;
6791        let eps = self.cfg.rms_eps;
6792        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
6793        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
6794        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
6795        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
6796        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
6797        // residual the next layer needs) as its `res` output. Falls back to the separate add
6798        // when the next layer is off the fused-q8 path.
6799        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
6800        for il in lo..hi {
6801            let layer = &self.layers[il];
6802            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
6803            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
6804            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
6805            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
6806            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
6807            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
6808            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
6809            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
6810            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
6811            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
6812            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
6813            // projections only; Linear mixer: the batched arm — the per-column fallback needs
6814            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
6815            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
6816            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
6817            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
6818            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
6819            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
6820            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
6821            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
6822            let lin_q8_only = match &layer.mixer {
6823                Mixer::Linear(la) => {
6824                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
6825                }
6826                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
6827                _ => true,
6828            };
6829            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
6830            // a non-fused layer still performs the residual add.
6831            let taken = pending.take();
6832            let (h, h_q8) = if norm_fused && lin_q8_only {
6833                let pair = match taken {
6834                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
6835                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
6836                    Some((x1p, f1p)) => {
6837                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
6838                        let p = e.add_rms_norm_q8_1(
6839                            &x1p,
6840                            &f1p,
6841                            layer.attn_norm.float_data(),
6842                            &mut x2,
6843                            n_embd,
6844                            t,
6845                            eps,
6846                        )?;
6847                        x = x2;
6848                        p
6849                    }
6850                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
6851                };
6852                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
6853            } else {
6854                if let Some((x1p, f1p)) = taken {
6855                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
6856                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
6857                    x = x2;
6858                }
6859                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
6860                if norm_fused {
6861                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6862                } else {
6863                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
6864                }
6865                (h, None)
6866            };
6867            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
6868
6869            let mixed = match &layer.mixer {
6870                Mixer::Full(fa) => self.full_attn_verify(
6871                    e,
6872                    fa,
6873                    &h,
6874                    h_q8_ref,
6875                    pos_d,
6876                    t,
6877                    cache,
6878                    il,
6879                    stream.map(|(_, c)| c),
6880                )?,
6881                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6882                Mixer::Linear(la) => {
6883                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
6884                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
6885                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
6886                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
6887                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
6888                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
6889                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
6890                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
6891                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
6892                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
6893                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
6894                    if (t >= 3 || (t == 2 && spec_m2()))
6895                        && mixer_fast
6896                        && e.uses_q8_1_fast(&la.ssm_out)
6897                    {
6898                        let want = ckpt.is_some();
6899                        let (out, stash) =
6900                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
6901                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6902                            ck.gdn[il] = Some(st);
6903                        }
6904                        out
6905                    } else {
6906                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
6907                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6908                            if ckpt.is_some() && t >= 2 {
6909                                Some(Vec::with_capacity(t - 1))
6910                            } else {
6911                                None
6912                            };
6913                        for col in 0..t {
6914                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
6915                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
6916                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
6917                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
6918                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
6919                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
6920                            // (pure dtod — cannot change any computed value). Last column skipped:
6921                            // rebuild targets are j <= t-1 columns.
6922                            if let Some(cs) = col_states.as_mut() {
6923                                if col + 1 < t {
6924                                    let rl = cache.recur[il].as_ref().unwrap();
6925                                    cs.push((
6926                                        e.clone_dtod(&rl.conv_state)?,
6927                                        e.clone_dtod(&rl.ssm_state)?,
6928                                    ));
6929                                }
6930                            }
6931                        }
6932                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
6933                            // ReplaySSM-assessment instrumentation (2026-07-30): the
6934                            // per-column clones are the only true state snapshots left in
6935                            // the verify (the batched path stashes INPUTS and replays).
6936                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
6937                                static ONCE: std::sync::Once = std::sync::Once::new();
6938                                let bytes: usize =
6939                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
6940                                ONCE.call_once(|| eprintln!(
6941                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
6942                                    cs.len(), bytes as f64 / 1e6));
6943                            }
6944                            ck.cols[il] = Some(cs);
6945                        }
6946                        out
6947                    }
6948                }
6949            };
6950
6951            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
6952            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
6953            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
6954            let ffn_fuse = match &layer.ffn {
6955                crate::hybrid::Ffn::Dense {
6956                    ffn_gate, ffn_up, ..
6957                } => {
6958                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
6959                        && e.uses_q8_1_fast(ffn_gate)
6960                        && e.uses_q8_1_fast(ffn_up)
6961                }
6962                crate::hybrid::Ffn::Moe(_) => false,
6963            };
6964            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
6965            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
6966            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
6967            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
6968            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
6969            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
6970            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
6971            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
6972            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
6973            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
6974            // mirror decode's dispatch or spec self-consistency fails.
6975            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
6976            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
6977            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
6978            let mut z = e.zeros(0)?; // replaced below on the unfused arms
6979            let z_q8 = if fuse_q8 {
6980                Some(e.add_rms_norm_q8_1(
6981                    &x,
6982                    &mixed,
6983                    layer.post_attn_norm.float_data(),
6984                    &mut x1,
6985                    n_embd,
6986                    t,
6987                    eps,
6988                )?)
6989            } else {
6990                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
6991                if ffn_fuse {
6992                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
6993                    e.rms_norm_decode(
6994                        &x1,
6995                        layer.post_attn_norm.float_data(),
6996                        &mut zf,
6997                        n_embd,
6998                        t,
6999                        eps,
7000                    )?;
7001                } else {
7002                    e.add_rms_norm(
7003                        &x,
7004                        &mixed,
7005                        layer.post_attn_norm.float_data(),
7006                        &mut x1,
7007                        &mut zf,
7008                        n_embd,
7009                        t,
7010                        eps,
7011                    )?;
7012                }
7013                z = zf;
7014                None
7015            };
7016            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7017            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7018            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7019            let ffn_out = match &layer.ffn {
7020                crate::hybrid::Ffn::Dense {
7021                    ffn_gate,
7022                    ffn_up,
7023                    ffn_down,
7024                } => {
7025                    let n_ff = ffn_gate.out_features();
7026                    if let Some((zq, zd)) = z_q8.as_ref() {
7027                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7028                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7029                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7030                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7031                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7032                        // structure at nrows=t.
7033                        let pair =
7034                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
7035                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
7036                                None => None,
7037                            };
7038                        let (gate, gs, up, us) = match pair {
7039                            Some(x4) => x4,
7040                            None => (
7041                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7042                                1.0, // scale already applied inside _pre
7043                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7044                                1.0,
7045                            ),
7046                        };
7047                        if e.uses_q8_1_fast(ffn_down) {
7048                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7049                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7050                        } else {
7051                            let mut act = vbuf(e, t * n_ff)?;
7052                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7053                            e.matmul_decode_exact(ffn_down, &act, t)?
7054                        }
7055                    } else {
7056                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7057                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7058                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7059                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7060                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7061                        let (gate, up) =
7062                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7063                                Some(pair) => pair,
7064                                None => (
7065                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
7066                                    e.matmul_decode_exact(ffn_up, &z, t)?,
7067                                ),
7068                            };
7069                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7070                        Self::ffn_act_lim(
7071                            e,
7072                            &self.cfg,
7073                            &gate,
7074                            &up,
7075                            1.0,
7076                            1.0,
7077                            dense_lim,
7078                            &mut act,
7079                            t * n_ff,
7080                        )?;
7081                        e.matmul_decode_exact(ffn_down, &act, t)?
7082                    }
7083                }
7084                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7085            };
7086            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7087            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7088            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7089            pending = Some((x1, ffn_out));
7090        }
7091        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7092        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7093        if let Some((x1p, f1p)) = pending.take() {
7094            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7095            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7096            x = x2;
7097        }
7098        Ok(x)
7099    }
7100    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7101    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7102    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7103    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7104    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7105    /// ssm state exactly like T sequential decode steps.
7106    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7107    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7108    #[allow(clippy::too_many_arguments)]
7109    fn linear_attn_verify_t(
7110        &self,
7111        e: &Engine,
7112        la: &LinearAttnLayer,
7113        h: &CudaSlice<f32>,
7114        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7115        t: usize,
7116        cache: &mut Cache,
7117        il: usize,
7118        want_stash: bool,
7119    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7120        let cfg = &self.cfg;
7121        let geometry = la.geometry;
7122        let d_state = geometry.key_head_dim as usize;
7123        let num_k = geometry.key_heads as usize;
7124        let num_v = geometry.value_heads as usize;
7125        let d_conv = geometry.conv_kernel as usize;
7126        let key_dim = d_state * num_k;
7127        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7128        let eps = cfg.rms_eps;
7129        let scale = 1.0 / (d_state as f32).sqrt();
7130
7131        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7132        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7133        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7134        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7135        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7136        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7137        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7138        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7139        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7140        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7141        // Bit-identical per (tensor,token,row) — see spec_fused_t().
7142        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7143        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7144        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7145        // and feeds every projection; the caller guaranteed all four input projections are
7146        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7147        let h_q8_t = if h_q8.is_none()
7148            && spec_fused_t()
7149            && (2..=4).contains(&t)
7150            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7151                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7152        {
7153            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7154        } else {
7155            None
7156        };
7157        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7158        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7159            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7160        let (qkv_mixed, z) = {
7161            let mut fused = None;
7162            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7163                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7164                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7165            } else if let Some((hq, hd)) = hq8_any {
7166                if spec_fused_t() && (2..=4).contains(&t) {
7167                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7168                }
7169            }
7170            match (fused, hq8_any) {
7171                (Some(pair), _) => pair,
7172                (None, Some((hq, hd))) if h_q8.is_some() => (
7173                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7174                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7175                ),
7176                (None, _) => (
7177                    e.matmul_decode_exact(&la.wqkv, h, t)?,
7178                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7179                ),
7180            }
7181        };
7182        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7183        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7184        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7185        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7186        let (beta_raw, alpha) = if t == 1 {
7187            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7188            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7189                Some(((mut b, bs), (mut a, as_))) => {
7190                    if bs != 1.0 {
7191                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7192                    }
7193                    if as_ != 1.0 {
7194                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7195                    }
7196                    (b, a)
7197                }
7198                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7199                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7200                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7201                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7202                    Some((b, a)) => (b, a),
7203                    None => (
7204                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7205                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7206                    ),
7207                },
7208            }
7209        } else {
7210            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7211            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7212            let mut nvfp4_fused = None;
7213            let mut q8_fused = None;
7214            if let Some((hq, hd)) = hq8_any {
7215                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7216                    nvfp4_fused =
7217                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7218                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7219                        static ONCE: std::sync::Once = std::sync::Once::new();
7220                        ONCE.call_once(|| {
7221                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7222                        });
7223                    }
7224                }
7225                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7226                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7227                }
7228            }
7229            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7230                if bs != 1.0 {
7231                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7232                }
7233                if as_ != 1.0 {
7234                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7235                }
7236                (b, a)
7237            } else if let Some(pair) = q8_fused {
7238                pair
7239            } else {
7240                match hq8_any {
7241                    Some((hq, hd)) if h_q8.is_some() => (
7242                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7243                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7244                    ),
7245                    _ => (
7246                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7247                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7248                    ),
7249                }
7250            }
7251        };
7252
7253        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7254        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7255        let rl = cache.recur[il].as_mut().unwrap();
7256        let mut conv_out = e.uninit(conv_dim * t)?;
7257        e.ssm_conv1d_tm_state(
7258            &qkv_mixed,
7259            &mut rl.conv_state,
7260            la.ssm_conv1d.float_data(),
7261            &mut conv_out,
7262            conv_dim,
7263            t,
7264            d_conv,
7265        )?;
7266
7267        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7268        let mut q_g = e.uninit(d_state * num_v * t)?;
7269        let mut k_g = e.uninit(d_state * num_v * t)?;
7270        let mut v_g = e.uninit(d_state * num_v * t)?;
7271        e.qkv_to_gdn_repack(
7272            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7273        )?;
7274        let mut q_l2 = e.uninit(d_state * num_v * t)?;
7275        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7276        let mut k_l2 = e.uninit(d_state * num_v * t)?;
7277        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7278        let mut beta = e.uninit(t * num_v)?;
7279        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7280        let mut g_log = e.uninit(t * num_v)?;
7281        e.gdn_glog(
7282            &alpha,
7283            la.ssm_dt.float_data(),
7284            la.ssm_a.float_data(),
7285            &mut g_log,
7286            num_v,
7287            t,
7288        )?;
7289
7290        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7291        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7292        let mut o = e.uninit(d_state * num_v * t)?;
7293        {
7294            let crate::cache::RecurLayer {
7295                ssm_state,
7296                ssm_state_alt,
7297                ..
7298            } = rl;
7299            e.gdn_scan_s128(
7300                &q_l2,
7301                &k_l2,
7302                &v_g,
7303                &g_log,
7304                &beta,
7305                ssm_state,
7306                ssm_state_alt,
7307                &mut o,
7308                num_v,
7309                t,
7310                scale,
7311            )?;
7312        }
7313        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7314
7315        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7316        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7317        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7318        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7319        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7320        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7321        let out = if e.uses_q8_1_fast(&la.ssm_out) {
7322            let (gq, gd) =
7323                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7324            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7325        } else {
7326            let mut gn = e.uninit(d_state * num_v * t)?;
7327            e.gated_rmsnorm(
7328                &o,
7329                la.ssm_norm.float_data(),
7330                &z,
7331                &mut gn,
7332                d_state,
7333                num_v * t,
7334                eps,
7335            )?;
7336            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7337            // would fall to dp4a with a different FP reduction order — same class of bug as
7338            // the input projs).
7339            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7340        };
7341        let stash = if want_stash {
7342            Some(GdnStash {
7343                qkv_mixed,
7344                q_l2,
7345                k_l2,
7346                v_g,
7347                g_log,
7348                beta,
7349            })
7350        } else {
7351            None
7352        };
7353        Ok((out, stash))
7354    }
7355
7356    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7357    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7358    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7359    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
7360    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7361    ///   replaying them.
7362    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7363    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7364    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7365    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7366    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
7367    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7368    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7369    fn commit_verified_prefix(
7370        &self,
7371        e: &Engine,
7372        cache: &mut Cache,
7373        snap: &crate::cache::CacheSnapshot,
7374        ckpt: &VerifyCkpt,
7375        j: usize,
7376        kv_lens_done: bool,
7377        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
7378    ) -> Result<(), Box<dyn std::error::Error>> {
7379        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
7380        // recurrent state and must never be forced through a synthetic SSM geometry.
7381        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
7382        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
7383        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
7384        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
7385        // buffers and stream order are identical to the per-layer memcpy sequence; the
7386        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
7387        let mut batched_cols = false;
7388        if state_copy_batch_on() && dev_j.is_none() {
7389            use cudarc::driver::DevicePtr;
7390            let s = &e.gpu.stream();
7391            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
7392            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
7393            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
7394            let mut uniform = true;
7395            for il in 0..self.layers.len() {
7396                let Some(rl) = cache.recur[il].as_ref() else {
7397                    continue;
7398                };
7399                if ckpt.gdn[il].is_some() {
7400                    continue; // kernel-rebuild arm restores below, per layer
7401                }
7402                let Some(cols) = &ckpt.cols[il] else {
7403                    continue; // missing-ckpt error surfaces in the main loop
7404                };
7405                let (c, st) = &cols[j - 1];
7406                if conv_pairs.is_empty() {
7407                    conv_words = c.len();
7408                    ssm_words = st.len();
7409                } else if c.len() != conv_words || st.len() != ssm_words {
7410                    uniform = false;
7411                    break;
7412                }
7413                let (pc, _g0) = c.device_ptr(s);
7414                let (dc, _g1) = rl.conv_state.device_ptr(s);
7415                let (ps, _g2) = st.device_ptr(s);
7416                let (ds, _g3) = rl.ssm_state.device_ptr(s);
7417                conv_pairs.push((pc as u64, dc as u64));
7418                ssm_pairs.push((ps as u64, ds as u64));
7419            }
7420            if uniform && !conv_pairs.is_empty() {
7421                let n = conv_pairs.len();
7422                let mut t = vec![0u64; 2 * n];
7423                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
7424                    t[k] = src;
7425                    t[n + k] = dst;
7426                }
7427                let conv_t = e.htod_u64(&t)?;
7428                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
7429                    t[k] = src;
7430                    t[n + k] = dst;
7431                }
7432                let ssm_t = e.htod_u64(&t)?;
7433                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
7434                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
7435                batched_cols = true;
7436            }
7437        }
7438        rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
7439        for il in 0..self.layers.len() {
7440            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7441                kvl.len = saved + j;
7442                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
7443                if !kv_lens_done {
7444                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7445                }
7446            }
7447            if let Some(rl) = cache.recur[il].as_mut() {
7448                let Mixer::Linear(linear) = &self.layers[il].mixer else {
7449                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7450                };
7451                let geometry = linear.geometry;
7452                let d_state = geometry.key_head_dim as usize;
7453                let num_k = geometry.key_heads as usize;
7454                let num_v = geometry.value_heads as usize;
7455                let d_conv = geometry.conv_kernel as usize;
7456                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7457                let scale = 1.0 / (d_state as f32).sqrt();
7458                if let Some(st) = &ckpt.gdn[il] {
7459                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7460                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7461                    if let Some((acc, base, t_v)) = dev_j {
7462                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
7463                        e.ssm_conv_ring_rebuild_dc(
7464                            &st.qkv_mixed,
7465                            ring_old,
7466                            &mut rl.conv_state,
7467                            conv_dim,
7468                            acc,
7469                            base,
7470                            t_v,
7471                            d_conv,
7472                        )?;
7473                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
7474                        e.gdn_scan_s128_dc(
7475                            &st.q_l2,
7476                            &st.k_l2,
7477                            &st.v_g,
7478                            &st.g_log,
7479                            &st.beta,
7480                            state_in,
7481                            &mut rl.ssm_state,
7482                            &mut o,
7483                            num_v,
7484                            acc,
7485                            base,
7486                            t_v,
7487                            scale,
7488                        )?;
7489                    } else {
7490                        e.ssm_conv_ring_rebuild(
7491                            &st.qkv_mixed,
7492                            ring_old,
7493                            &mut rl.conv_state,
7494                            conv_dim,
7495                            j,
7496                            d_conv,
7497                        )?;
7498                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
7499                        e.gdn_scan_s128(
7500                            &st.q_l2,
7501                            &st.k_l2,
7502                            &st.v_g,
7503                            &st.g_log,
7504                            &st.beta,
7505                            state_in,
7506                            &mut rl.ssm_state,
7507                            &mut o,
7508                            num_v,
7509                            j,
7510                            scale,
7511                        )?;
7512                    }
7513                } else if let Some(cols) = &ckpt.cols[il] {
7514                    if !batched_cols {
7515                        let (c, s) = &cols[j - 1];
7516                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
7517                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
7518                    }
7519                } else {
7520                    return Err(
7521                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
7522                    );
7523                }
7524            }
7525        }
7526        cache.pos = snap.pos + j;
7527        Ok(())
7528    }
7529
7530    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
7531    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
7532    fn commit_verified_prefix_stream(
7533        &self,
7534        e: &Engine,
7535        cache: &mut Cache,
7536        snap: &crate::cache::CacheSnapshot,
7537        ckpt: &VerifyCkpt,
7538        acc: &CudaSlice<u32>,
7539        base: usize,
7540        t_v: usize,
7541    ) -> Result<(), Box<dyn std::error::Error>> {
7542        for il in 0..self.layers.len() {
7543            if let Some(rl) = cache.recur[il].as_mut() {
7544                let Mixer::Linear(linear) = &self.layers[il].mixer else {
7545                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7546                };
7547                let geometry = linear.geometry;
7548                let d_state = geometry.key_head_dim as usize;
7549                let num_k = geometry.key_heads as usize;
7550                let num_v = geometry.value_heads as usize;
7551                let d_conv = geometry.conv_kernel as usize;
7552                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7553                let scale = 1.0 / (d_state as f32).sqrt();
7554                let st = ckpt.gdn[il]
7555                    .as_ref()
7556                    .ok_or("stream restore: batched-linear stash missing")?;
7557                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7558                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7559                e.ssm_conv_ring_rebuild_dc(
7560                    &st.qkv_mixed,
7561                    ring_old,
7562                    &mut rl.conv_state,
7563                    conv_dim,
7564                    acc,
7565                    base,
7566                    t_v,
7567                    d_conv,
7568                )?;
7569                let mut o = e.uninit(d_state * num_v * t_v)?;
7570                e.gdn_scan_s128_dc(
7571                    &st.q_l2,
7572                    &st.k_l2,
7573                    &st.v_g,
7574                    &st.g_log,
7575                    &st.beta,
7576                    state_in,
7577                    &mut rl.ssm_state,
7578                    &mut o,
7579                    num_v,
7580                    acc,
7581                    base,
7582                    t_v,
7583                    scale,
7584                )?;
7585            }
7586        }
7587        Ok(())
7588    }
7589
7590    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
7591    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
7592    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
7593    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
7594    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
7595    pub fn decode_step_t_aux2(
7596        &self,
7597        e: &Engine,
7598        tokens: &[u32],
7599        pos0: usize,
7600        cache: &mut Cache,
7601        aux_layers: &[usize],
7602        pred_col: Option<usize>,
7603    ) -> Result<
7604        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
7605        Box<dyn std::error::Error>,
7606    > {
7607        let cfg = &self.cfg;
7608        let n_embd = cfg.n_embd as usize;
7609        let eps = cfg.rms_eps;
7610        let t = tokens.len();
7611        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7612        let pos_d = e.htod_i32(&pos_vec)?;
7613        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7614        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
7615        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
7616        let want_pred = pred_col.is_some();
7617
7618        for (il, layer) in self.layers.iter().enumerate() {
7619            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
7620            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7621            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7622            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7623            if norm_fused {
7624                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7625            } else {
7626                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7627            }
7628            let mixed = match &layer.mixer {
7629                Mixer::Full(fa) => {
7630                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
7631                }
7632                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7633                Mixer::Linear(la) => {
7634                    let mut out = e.zeros(t * n_embd)?;
7635                    for col in 0..t {
7636                        let mut h_col = e.zeros(n_embd)?;
7637                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
7638                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7639                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7640                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7641                    }
7642                    out
7643                }
7644            };
7645            let ffn_fuse = match &layer.ffn {
7646                crate::hybrid::Ffn::Dense {
7647                    ffn_gate, ffn_up, ..
7648                } => {
7649                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7650                        && e.uses_q8_1_fast(ffn_gate)
7651                        && e.uses_q8_1_fast(ffn_up)
7652                }
7653                crate::hybrid::Ffn::Moe(_) => false,
7654            };
7655            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
7656            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7657            if ffn_fuse {
7658                e.add(&x, &mixed, &mut x1, t * n_embd)?;
7659                e.rms_norm_decode(
7660                    &x1,
7661                    layer.post_attn_norm.float_data(),
7662                    &mut z,
7663                    n_embd,
7664                    t,
7665                    eps,
7666                )?;
7667            } else {
7668                e.add_rms_norm(
7669                    &x,
7670                    &mixed,
7671                    layer.post_attn_norm.float_data(),
7672                    &mut x1,
7673                    &mut z,
7674                    n_embd,
7675                    t,
7676                    eps,
7677                )?;
7678            }
7679            let ffn_out = match &layer.ffn {
7680                crate::hybrid::Ffn::Dense {
7681                    ffn_gate,
7682                    ffn_up,
7683                    ffn_down,
7684                } => {
7685                    let n_ff = ffn_gate.out_features();
7686                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
7687                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
7688                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7689                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
7690                    Self::ffn_act_lim(
7691                        e,
7692                        &self.cfg,
7693                        &gate,
7694                        &up,
7695                        1.0,
7696                        1.0,
7697                        self.cfg.clamp_shexp_at(il as u32),
7698                        &mut act,
7699                        t * n_ff,
7700                    )?;
7701                    e.matmul_decode_exact(ffn_down, &act, t)?
7702                }
7703                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7704            };
7705            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7706            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7707            if aux_layers.contains(&il) {
7708                let mut a = e.zeros(n_embd)?;
7709                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
7710                aux_last.push(a);
7711                if let Some(pc) = pred_col {
7712                    let mut ap = e.zeros(n_embd)?;
7713                    e.copy_view_into(
7714                        &mut ap,
7715                        0,
7716                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
7717                        n_embd,
7718                    )?;
7719                    aux_pred.push(ap);
7720                }
7721            }
7722            x = x2;
7723        }
7724        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
7725        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
7726        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
7727        let host = e.dtoh(&logits)?;
7728        cache.pos += t;
7729        Ok((
7730            host,
7731            aux_last,
7732            if want_pred { Some(aux_pred) } else { None },
7733        ))
7734    }
7735
7736    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
7737    /// `step35_decode_attn`.
7738    ///
7739    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
7740    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
7741    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
7742    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
7743    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
7744    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
7745    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
7746    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
7747    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
7748    /// position of each query row. A batched twin would have to reproduce all of that AND the
7749    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
7750    /// take one `base_len`, not a per-row offset).
7751    ///
7752    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
7753    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
7754    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
7755    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
7756    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
7757    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
7758    /// step35 twin is a perf lane's job and must be gated against this arm.
7759    ///
7760    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
7761    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
7762    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
7763    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
7764    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
7765    #[allow(clippy::too_many_arguments)]
7766    fn step35_verify(
7767        &self,
7768        e: &Engine,
7769        fa: &FullAttnLayer,
7770        h: &CudaSlice<f32>,
7771        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7772        t: usize,
7773        cache: &mut Cache,
7774        il: usize,
7775    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7776        let n_embd = self.cfg.n_embd as usize;
7777        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
7778        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
7779        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
7780        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
7781        // cannot regress it into silently reading an empty buffer.
7782        assert_eq!(
7783            h.len(),
7784            t * n_embd,
7785            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
7786             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
7787            h_q8.is_some()
7788        );
7789        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
7790        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
7791        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
7792        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
7793        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
7794        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
7795        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
7796        for r in 0..t {
7797            // Absolute position of this query row. `cache.pos` is the committed length at round
7798            // start and every row before r has already been appended by this loop, so the r-th
7799            // verify token sits at cache.pos + r — the same position eager decode would give it.
7800            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
7801            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
7802            e.copy_view_into(
7803                &mut h_row,
7804                0,
7805                &h.slice(r * n_embd..(r + 1) * n_embd),
7806                n_embd,
7807            )?;
7808            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
7809            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
7810            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
7811            debug_assert_eq!(
7812                o.len(),
7813                n_embd,
7814                "step35_decode_attn returns post-wo [n_embd]"
7815            );
7816            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
7817        }
7818        Ok(out)
7819    }
7820
7821    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
7822    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
7823    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
7824    #[allow(clippy::too_many_arguments)]
7825    fn full_attn_verify(
7826        &self,
7827        e: &Engine,
7828        fa: &FullAttnLayer,
7829        h: &CudaSlice<f32>,
7830        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7831        pos_d: &CudaSlice<i32>,
7832        t: usize,
7833        cache: &mut Cache,
7834        il: usize,
7835        stream_ctr: Option<&CudaSlice<i32>>,
7836    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7837        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
7838        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
7839        // its own arm. A verify that silently computes different attention than decode defeats the
7840        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
7841        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
7842        // shape and not laziness.
7843        if self.sliding_gated_moe_batch_program() {
7844            if stream_ctr.is_some() {
7845                return Err(
7846                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7847                            cannot express the SWA offset KV view; same root cause as the dc \
7848                            decode refusal) — run spec without the stream arm"
7849                        .into(),
7850                );
7851            }
7852            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
7853        }
7854        let cfg = &self.cfg;
7855        let geometry = cfg.full_attention_geometry_at(il as u32);
7856        let n_head = geometry.n_head as usize;
7857        let n_head_kv = geometry.n_head_kv as usize;
7858        let head_dim = geometry.head_dim_k as usize;
7859        let eps = cfg.rms_eps;
7860        let scale = geometry.attention_scale();
7861        let n_embd = cfg.n_embd as usize;
7862
7863        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
7864        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
7865        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
7866        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
7867        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
7868        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
7869        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
7870        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
7871        let (qf, mut k, v) = {
7872            let mut fused = None;
7873            let qkv_fast =
7874                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
7875            if t == 1 && qkv_fast {
7876                let (hq_o, hd_o);
7877                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7878                    Some(p) => p,
7879                    None => {
7880                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
7881                        (&hq_o, &hd_o)
7882                    }
7883                };
7884                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
7885            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
7886                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
7887                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
7888                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
7889                let (hq_o, hd_o);
7890                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
7891                    Some(p) => p,
7892                    None => {
7893                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
7894                        (&hq_o, &hd_o)
7895                    }
7896                };
7897                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
7898            }
7899            match (fused, h_q8) {
7900                (Some(triple), _) => triple,
7901                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
7902                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
7903                (None, Some((hq, hd))) if qkv_fast => (
7904                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
7905                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
7906                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
7907                ),
7908                (None, _) => (
7909                    e.matmul_decode_exact(&fa.wq, h, t)?,
7910                    e.matmul_decode_exact(&fa.wk, h, t)?,
7911                    e.matmul_decode_exact(&fa.wv, h, t)?,
7912                ),
7913            }
7914        };
7915        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
7916        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7917        let (mut q, gate) = if gated {
7918            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7919            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
7920            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
7921            (q, Some(gate))
7922        } else {
7923            (qf, None)
7924        };
7925
7926        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
7927        e.rms_norm(
7928            &q,
7929            fa.q_norm.float_data(),
7930            &mut qn,
7931            head_dim,
7932            n_head * t,
7933            eps,
7934        )?;
7935        q = qn;
7936        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
7937        e.rms_norm(
7938            &k,
7939            fa.k_norm.float_data(),
7940            &mut kn,
7941            head_dim,
7942            n_head_kv * t,
7943            eps,
7944        )?;
7945        k = kn;
7946        let rope_dims = geometry.n_rot as usize;
7947        e.rope_neox(
7948            &mut q,
7949            pos_d,
7950            head_dim,
7951            rope_dims,
7952            n_head,
7953            t,
7954            geometry.rope_base,
7955            1.0,
7956        )?;
7957        e.rope_neox(
7958            &mut k,
7959            pos_d,
7960            head_dim,
7961            rope_dims,
7962            n_head_kv,
7963            t,
7964            geometry.rope_base,
7965            1.0,
7966        )?;
7967
7968        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
7969        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
7970        let kvl = cache.kv[il].as_mut().unwrap();
7971        let (kv_dim_k, kv_dim_v, ktb, vtb) =
7972            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
7973        if let Some(ctr) = stream_ctr {
7974            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
7975            // math on a (block, token) grid, documented byte-identical); host len is a stale
7976            // LOWER BOUND under pre-issue (drain reconciles it).
7977            e.append_kv_quantized_rows_dc(
7978                &k,
7979                &v,
7980                &mut kvl.k,
7981                &mut kvl.v,
7982                ctr,
7983                t,
7984                kv_dim_k,
7985                kv_dim_v,
7986                ktb,
7987                vtb,
7988                crate::Engine::kv_fp8_on(),
7989            )?;
7990        } else {
7991            for i in 0..t {
7992                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
7993                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
7994                e.append_kv_quantized_view(
7995                    &k_row,
7996                    &v_row,
7997                    &mut kvl.k,
7998                    &mut kvl.v,
7999                    kvl.len + i,
8000                    kv_dim_k,
8001                    kv_dim_v,
8002                    ktb,
8003                    vtb,
8004                    crate::Engine::kv_fp8_on(),
8005                )?;
8006            }
8007            kvl.len += t;
8008        }
8009
8010        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8011        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8012        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8013        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8014        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8015        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8016        // keys. The verify appends all T tokens first but bounds the key range per row.
8017        //
8018        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8019        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8020        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8021        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8022        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8023        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8024        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8025        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8026        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8027        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8028        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8029        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8030        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8031        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8032        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8033        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8034        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8035        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8036        if let Some(ctr) = stream_ctr {
8037            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8038            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8039            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8040            let upper = kvl.len + t + 64;
8041            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8042            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8043            e.fa_decode_rows_dc(
8044                &q,
8045                &k_view,
8046                &v_view,
8047                &mut attn,
8048                head_dim,
8049                n_head,
8050                n_head_kv,
8051                ctr,
8052                upper.min(cache.max_ctx),
8053                t,
8054                scale,
8055                ktb,
8056                vtb,
8057                0,
8058                false,
8059            )?;
8060        } else if spec_lean() && t == 1 {
8061            let t_kv = base_len + 1;
8062            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8063            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8064            e.fa_decode_kvmod(
8065                &q,
8066                &k_view,
8067                &v_view,
8068                &mut attn,
8069                head_dim,
8070                n_head,
8071                n_head_kv,
8072                t_kv,
8073                scale,
8074                ktb,
8075                vtb,
8076                crate::Engine::kv_fp8_on(),
8077            )?;
8078        } else if e.fa_rows_eligible(base_len, head_dim) {
8079            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8080            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8081            e.fa_decode_rows(
8082                &q,
8083                &k_view,
8084                &v_view,
8085                &mut attn,
8086                head_dim,
8087                n_head,
8088                n_head_kv,
8089                base_len,
8090                t,
8091                scale,
8092                ktb,
8093                vtb,
8094                None,
8095                false,
8096                crate::Engine::kv_fp8_on(),
8097                None,
8098            )?;
8099        } else {
8100            for r in 0..t {
8101                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8102                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8103                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8104                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8105                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8106                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8107                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8108                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8109                e.fa_decode_kvmod(
8110                    &q_row,
8111                    &k_view_r,
8112                    &v_view_r,
8113                    &mut attn_row,
8114                    head_dim,
8115                    n_head,
8116                    n_head_kv,
8117                    t_kv_r,
8118                    scale,
8119                    ktb,
8120                    vtb,
8121                    crate::Engine::kv_fp8_on(),
8122                )?;
8123                e.copy_into(
8124                    &mut attn,
8125                    r * n_head * head_dim,
8126                    &attn_row,
8127                    n_head * head_dim,
8128                )?;
8129            }
8130        }
8131
8132        let attn_g = match &gate {
8133            Some(gate) => {
8134                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8135                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8136                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8137                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8138                ag
8139            }
8140            None => attn,
8141        };
8142        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8143        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8144        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
8145    }
8146
8147    /// Context-linear bytes for a plain serving session's trunk cache.
8148    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8149        crate::cache::cache_bytes_per_token_for_plan(
8150            &self.cfg,
8151            &self.plan,
8152            0,
8153            self.plan.layers.len(),
8154        )
8155    }
8156
8157    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8158    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8159        (
8160            self.plain_session_kv_bytes_per_token(),
8161            crate::cache::cache_ring_bytes_per_token_for_plan(
8162                &self.cfg,
8163                &self.plan,
8164                0,
8165                self.plan.layers.len(),
8166            ),
8167            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8168        )
8169    }
8170
8171    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8172    /// scratch. With no MTP head this equals the plain coefficient.
8173    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8174        let scratch = self
8175            .mtp
8176            .iter()
8177            .chain(self.mtp_extra.iter())
8178            .map(|mtp| {
8179                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8180                k + v
8181            })
8182            .sum::<usize>();
8183        self.plain_session_kv_bytes_per_token()
8184            .saturating_add(scratch)
8185    }
8186
8187    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8188    /// capped by the same SWA ring rows as the trunk.
8189    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8190        let total = self.spec_session_kv_bytes_per_token();
8191        let (_, mut ring, rows) = self.plain_session_kv_shape();
8192        if rows > 0 {
8193            ring = ring.saturating_add(
8194                self.mtp
8195                    .iter()
8196                    .chain(self.mtp_extra.iter())
8197                    .map(|mtp| {
8198                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8199                        k + v
8200                    })
8201                    .sum::<usize>(),
8202            );
8203        }
8204        (total, ring, rows)
8205    }
8206
8207    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8208    /// the NextN head to draft K tokens then verifies them in one batched target forward.
8209    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8210    /// acceptance rate. `k` = draft length per round.
8211    ///
8212    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8213    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8214    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8215    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8216    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8217    /// captured graph references is event-free; the spec loop is strictly single-stream.
8218    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8219    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8220    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8221    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8222    /// generate_spec_inner2.
8223    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8224    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8225    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8226    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8227    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8228    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8229    pub fn new_session(
8230        &self,
8231        e: &Engine,
8232        max_ctx: usize,
8233    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8234        Ok(SpecSession {
8235            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8236            // is the SERVING spec-session path, and with the ppN door open across two cards a
8237            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8238            // round — the wrong-card class already fixed on the two batched serving paths
8239            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8240            // branch, same allocations), so single-device behavior is byte-unchanged.
8241            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8242            scratch: self.new_mtp_scratch(e, max_ctx)?,
8243            committed: Vec::new(),
8244            last_h: None,
8245            next_pred: None,
8246            sctr: 0,
8247            uctr: 0,
8248            draft_ctx: None,
8249            pending_tok: None,
8250            turn_ckpt: None,
8251            telem: SpecTelemetryCounters::default(),
8252            capture_at: None,
8253            boundary_captures: Vec::new(),
8254            ckpt_at: None,
8255        })
8256    }
8257
8258    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8259    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8260    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8261    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8262    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8263    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8264    /// worker always receives a fully-warm continuation session (committed = whole
8265    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8266    /// boundary logits on the empty-suffix shape).
8267    ///
8268    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8269    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8270    /// request, and plain feeds a carried suffix via eager `decode_step` below
8271    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8272    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8273    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8274    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8275    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8276    /// burst prime.
8277    ///
8278    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8279    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8280    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8281    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8282    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8283    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8284    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8285    /// cold session draws from the identical row at counter 0 and then runs its rounds from
8286    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8287    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8288    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8289    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8290    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8291    ///
8292    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8293    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8294    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8295    /// and are never routed here.
8296    ///
8297    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8298    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8299    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8300    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8301    /// entry stays published for the next request.
8302    #[allow(clippy::too_many_arguments)]
8303    pub fn spec_session_from_restored(
8304        &self,
8305        e: &Engine,
8306        mut cache: Cache,
8307        prefix: Vec<u32>,
8308        suffix: &[u32],
8309        draft_k: &CudaSlice<u8>,
8310        draft_v: &CudaSlice<u8>,
8311        draft_k_tok_bytes: usize,
8312        draft_v_tok_bytes: usize,
8313        draft_len: usize,
8314        last_h: &[f32],
8315        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8316        // when a suffix follows — the feed's own logits are the boundary then.
8317        boundary_logits: &[f32],
8318        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8319        // ONE place instead of being half-applied by the worker.
8320        sampling: Option<SpecSampling>,
8321        require_anchor: bool,
8322        max_ctx: usize,
8323        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8324        // prompt position to split the suffix feed at and capture the extended-entry
8325        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8326        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8327        // WHY: the prompt-end capture below includes the template's live generation header
8328        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8329        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8330        // diverged from every future prompt and the hit boundary FROZE at the first
8331        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8332        republish_at: Option<usize>,
8333    ) -> Result<SpecSession, (Option<Cache>, String)> {
8334        let pos = prefix.len();
8335        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8336            Err((Some(cache), msg))
8337        };
8338        if self.mtp.is_none() {
8339            return fail(cache, "no MTP head attached (nothing to draft with)".into());
8340        }
8341        if pos == 0 {
8342            return fail(cache, "empty committed prefix".into());
8343        }
8344        if cache.pos != pos {
8345            let msg = format!(
8346                "restored cache pos {} != restored prefix len {pos}",
8347                cache.pos
8348            );
8349            return fail(cache, msg);
8350        }
8351        if draft_len != pos {
8352            return fail(
8353                cache,
8354                format!("draft plane len {draft_len} != restored prefix len {pos}"),
8355            );
8356        }
8357        if pos + suffix.len() >= max_ctx {
8358            return fail(
8359                cache,
8360                format!(
8361                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8362                    pos + suffix.len(),
8363                ),
8364            );
8365        }
8366        let mut scratch = match MtpScratch::new(
8367            e,
8368            &self.cfg,
8369            &self.plan,
8370            max_ctx,
8371            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8372        ) {
8373            Ok(s) => s,
8374            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8375        };
8376        if scratch.kv.ring.is_some() {
8377            return fail(
8378                cache,
8379                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8380            );
8381        }
8382        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8383            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8384        {
8385            return fail(
8386                cache,
8387                format!(
8388                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8389                     {}/{} bytes/token (stale entry across a format change)",
8390                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8391                ),
8392            );
8393        }
8394        if pos > scratch.cap {
8395            return fail(
8396                cache,
8397                format!(
8398                    "draft plane rows {pos} exceed scratch capacity {}",
8399                    scratch.cap
8400                ),
8401            );
8402        }
8403        let kb = pos * draft_k_tok_bytes;
8404        let vb = pos * draft_v_tok_bytes;
8405        if draft_k.len() < kb || draft_v.len() < vb {
8406            return fail(
8407                cache,
8408                format!(
8409                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8410                    draft_k.len(),
8411                    draft_v.len(),
8412                ),
8413            );
8414        }
8415        if kb > 0 {
8416            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
8417                return fail(cache, format!("draft K restore copy failed: {err}"));
8418            }
8419        }
8420        if vb > 0 {
8421            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
8422                return fail(cache, format!("draft V restore copy failed: {err}"));
8423            }
8424        }
8425        if let Err(err) = scratch.set_len(e, pos) {
8426            return fail(cache, format!("draft scratch len set failed: {err}"));
8427        }
8428        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
8429            // anchor upload failure is acceptance-only when a suffix feed follows (fill
8430            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
8431            // burst entry asserts committed + last_h + next_pred) — the caller says which.
8432            e.htod(last_h).ok()
8433        } else {
8434            None
8435        };
8436        if require_anchor && last_h_dev.is_none() {
8437            return fail(
8438                cache,
8439                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
8440            );
8441        }
8442        let mut committed = prefix;
8443        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
8444        // what the empty-suffix continuation assert in the burst entry requires.
8445        let next_pred;
8446        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
8447        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
8448        // drawing its own first token from the same row.
8449        let mut sctr = 0u32;
8450        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
8451        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
8452        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
8453        // after the suffix joins `committed` below.
8454        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
8455        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
8456        if !suffix.is_empty() {
8457            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
8458            // From here on the trunk cache mutates: failures return Err((None, _)) and
8459            // the worker serves the request cold-plain instead of reusing the carrier.
8460            let dirty =
8461                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
8462            let n_embd = self.cfg.n_embd as usize;
8463            let t = suffix.len();
8464            let mut h_rows = match e.uninit(t * n_embd) {
8465                Ok(b) => b,
8466                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
8467            };
8468            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
8469            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
8470            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
8471            let b_rel = republish_at
8472                .and_then(|abs| abs.checked_sub(pos))
8473                .filter(|&r| r > 0 && r < t);
8474            let mut feed_logits = Vec::new();
8475            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
8476                || e.frozen_cpu_experts_prefer_tokenwise_prime();
8477            let mut fed = 0usize;
8478            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
8479                if seg_end <= fed {
8480                    continue;
8481                }
8482                let seg = &suffix[fed..seg_end];
8483                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
8484                if batched {
8485                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
8486                    // queued after this segment ride `queued_after` so Step35 arm selection
8487                    // stays keyed to the request's end (tick-seg law).
8488                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
8489                        Ok((l, _h_seed, hiddens)) => {
8490                            if let Err(err) =
8491                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
8492                            {
8493                                return dirty(format!("suffix hidden copy: {err}"));
8494                            }
8495                            feed_logits = l;
8496                        }
8497                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
8498                    }
8499                } else {
8500                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
8501                    for (i, &tok) in seg.iter().enumerate() {
8502                        match self.decode_step_h(e, tok, &mut cache) {
8503                            Ok((l, h)) => {
8504                                if let Err(err) =
8505                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
8506                                {
8507                                    return dirty(format!("suffix hidden copy: {err}"));
8508                                }
8509                                feed_logits = l;
8510                            }
8511                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
8512                        }
8513                    }
8514                }
8515                fed = seg_end;
8516                if Some(seg_end) == b_rel {
8517                    // The stable pre-generation boundary: capture the extended-entry
8518                    // publication AND this session's own turn checkpoint here instead of at
8519                    // prompt-end (both would otherwise carry the volatile live-header tail
8520                    // the next re-render replaces). Failure silent, turn_ckpt convention.
8521                    debug_assert_eq!(
8522                        cache.pos,
8523                        pos + seg_end,
8524                        "stable-boundary capture off the feed split"
8525                    );
8526                    if spec_restore_republish_on() {
8527                        if let Ok(snap) = cache.snapshot(e) {
8528                            boundary_captures.push(SpecBoundaryCapture {
8529                                snap,
8530                                pos: pos + seg_end,
8531                                logits: feed_logits.clone(),
8532                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
8533                            });
8534                        }
8535                    }
8536                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8537                        e.uninit(n_embd).and_then(|mut a| {
8538                            e.copy_view_into(
8539                                &mut a,
8540                                0,
8541                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
8542                                n_embd,
8543                            )?;
8544                            Ok(a)
8545                        });
8546                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
8547                        restored_turn_ckpt = Some(SpecCheckpoint {
8548                            snap,
8549                            pos: pos + seg_end,
8550                            last_h,
8551                        });
8552                    }
8553                }
8554            }
8555            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
8556            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
8557            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
8558            // with T). Fill failures are acceptance-only — truncate to the restored rows
8559            // and continue; the burst's own set_len keeps the invariant.
8560            let mtp = self.mtp.as_ref().expect("mtp checked above");
8561            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8562            let embd_gpu = if spec_host_embd() {
8563                None
8564            } else {
8565                Some(
8566                    self.embd_gpu
8567                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8568                )
8569            };
8570            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8571            let fill_chunk = 4096usize;
8572            let mut filled = true;
8573            let mut start = 0usize;
8574            'fill: while start < t {
8575                let end = (start + fill_chunk).min(t);
8576                let tc = end - start;
8577                let Ok(mut phs) = e.zeros(tc * n_embd) else {
8578                    filled = false;
8579                    break 'fill;
8580                };
8581                let (src_lo, dst_off, n_copy) = if start == 0 {
8582                    (0, n_embd, (tc - 1) * n_embd)
8583                } else {
8584                    ((start - 1) * n_embd, 0, tc * n_embd)
8585                };
8586                if start == 0 {
8587                    if let Some(lh) = last_h_dev.as_ref() {
8588                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
8589                            filled = false;
8590                            break 'fill;
8591                        }
8592                    }
8593                }
8594                if n_copy > 0
8595                    && e.copy_view_into(
8596                        &mut phs,
8597                        dst_off,
8598                        &h_rows.slice(src_lo..src_lo + n_copy),
8599                        n_copy,
8600                    )
8601                    .is_err()
8602                {
8603                    filled = false;
8604                    break 'fill;
8605                }
8606                if self
8607                    .mtp_kv_fill_all(
8608                        e,
8609                        &suffix[start..end],
8610                        &phs,
8611                        pos + start,
8612                        &mut scratch,
8613                        embd_dev,
8614                    )
8615                    .is_err()
8616                {
8617                    filled = false;
8618                    break 'fill;
8619                }
8620                start = end;
8621            }
8622            if !filled {
8623                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
8624                // so keep only the restored rows resident and let verify arbitrate.
8625                if let Err(err) = scratch.set_len(e, pos) {
8626                    return dirty(format!("scratch truncation after failed fill: {err}"));
8627                }
8628            }
8629            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
8630            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
8631            // finding (d)). Pre-lane, publication was armed only for COLD sessions
8632            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
8633            // non-continuation burst — but a converted hit's first burst IS a continuation,
8634            // so a growing conversation learned exactly ONE boundary and turn 3 could never
8635            // hit a longer prefix than turn 2 did.
8636            //
8637            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
8638            // line — the trunk is primed over the whole prompt, nothing is generated, and the
8639            // draft plane rows [0..prompt) are filled just above. That is a complete
8640            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
8641            // publishes; the worker's existing publication sweep picks it up because it is
8642            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
8643            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
8644            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
8645            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
8646            // publication is an optimization, never a correctness dependency.
8647            //
8648            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
8649            // entry's tail is the live generation header the next re-render replaces, so on a
8650            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
8651            // the stable-boundary capture above IS this publication, minus the poisoned tail.
8652            if spec_restore_republish_on() && boundary_captures.is_empty() {
8653                debug_assert_eq!(
8654                    cache.pos,
8655                    pos + t,
8656                    "extended-entry capture must sit at the restored session's prompt end",
8657                );
8658                if let Ok(snap) = cache.snapshot(e) {
8659                    boundary_captures.push(SpecBoundaryCapture {
8660                        snap,
8661                        pos: pos + t,
8662                        logits: feed_logits.clone(),
8663                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
8664                    });
8665                }
8666            }
8667            // continuation seed: the feed's boundary logits ARE the plain path's boundary
8668            // logits (same program), so greedy's argmax here is plain's first emitted token,
8669            // and the sampled draw is the cold sampled session's own first token.
8670            next_pred = Some(if sampled {
8671                let sp = sampling.expect("sampled implies a sampler");
8672                // `committed` is still the restored prefix here; the suffix joins it below —
8673                // so this is the last-N window over the WHOLE prompt, exactly the cold
8674                // session's own window at its first token.
8675                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
8676                match sample_boundary_token(
8677                    e,
8678                    &feed_logits,
8679                    &sp,
8680                    &hist,
8681                    &mut sctr,
8682                    "restore-suffix-feed",
8683                ) {
8684                    Ok(t) => t,
8685                    // the trunk is already fed: hand nothing back, the worker serves the
8686                    // request cold-plain. Never fall back to an argmax — that would put a
8687                    // greedy token in a sampled stream to save a slow path.
8688                    Err(err) => {
8689                        return dirty(format!("boundary token draw failed: {err}"));
8690                    }
8691                }
8692            } else {
8693                argmax(&feed_logits) as u32
8694            });
8695            let mut lh = match e.uninit(n_embd) {
8696                Ok(b) => b,
8697                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
8698            };
8699            if let Err(err) = e.copy_view_into(
8700                &mut lh,
8701                0,
8702                &h_rows.slice((t - 1) * n_embd..t * n_embd),
8703                n_embd,
8704            ) {
8705                return dirty(format!("boundary hidden copy: {err}"));
8706            }
8707            last_h_dev = Some(lh);
8708            committed.extend_from_slice(suffix);
8709        } else {
8710            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
8711            // ENTRY's boundary logits are the boundary row, and this is the token the cold
8712            // session emits from that same row. Owned here rather than in the worker so the
8713            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
8714            if boundary_logits.is_empty() {
8715                return fail(
8716                    cache,
8717                    "full-cover restore without the entry's boundary logits".into(),
8718                );
8719            }
8720            next_pred = Some(if sampled {
8721                let sp = sampling.expect("sampled implies a sampler");
8722                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
8723                match sample_boundary_token(
8724                    e,
8725                    boundary_logits,
8726                    &sp,
8727                    &hist,
8728                    &mut sctr,
8729                    "restore-full-cover",
8730                ) {
8731                    Ok(t) => t,
8732                    // nothing has been mutated on this shape — hand the carrier back and let
8733                    // the hit serve PLAIN (the banked pre-lane path).
8734                    Err(err) => {
8735                        return fail(cache, format!("boundary token draw failed: {err}"));
8736                    }
8737                }
8738            } else {
8739                argmax(boundary_logits) as u32
8740            });
8741        }
8742        Ok(SpecSession {
8743            cache,
8744            scratch,
8745            committed,
8746            last_h: last_h_dev,
8747            next_pred,
8748            sctr,
8749            uctr: 0,
8750            draft_ctx: None,
8751            pending_tok: None,
8752            // Stable-boundary capture from the split feed above (None on the legacy shape):
8753            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
8754            // affinity probe declined ("no turn checkpoint retained") and the conversation
8755            // fell back to the frozen prefix entry forever.
8756            turn_ckpt: restored_turn_ckpt,
8757            telem: SpecTelemetryCounters::default(),
8758            capture_at: None,
8759            boundary_captures,
8760            ckpt_at: None,
8761        })
8762    }
8763
8764    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
8765    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
8766    /// snapshot, or draft-KV row that only corrupts the following round.
8767    pub fn optipipe_compare_session_state(
8768        &self,
8769        e: &Engine,
8770        reference: &SpecSession,
8771        candidate: &SpecSession,
8772    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
8773        fn fail(what: &str) -> Box<dyn std::error::Error> {
8774            format!("optipipe state mismatch: {what}").into()
8775        }
8776        fn same_f32(a: &[f32], b: &[f32]) -> bool {
8777            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
8778        }
8779        fn compare_layers(
8780            es: &Engine,
8781            range: std::ops::Range<usize>,
8782            reference: &SpecSession,
8783            candidate: &SpecSession,
8784            report: &mut OptiForkStateIdentity,
8785        ) -> Result<(), Box<dyn std::error::Error>> {
8786            for il in range {
8787                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
8788                    (Some(a), Some(b)) => {
8789                        if a.len != b.len {
8790                            return Err(fail(&format!(
8791                                "layer {il} host KV len {} != {}",
8792                                a.len, b.len
8793                            )));
8794                        }
8795                        let ad = es.dtoh_i32(&a.len_d)?;
8796                        let bd = es.dtoh_i32(&b.len_d)?;
8797                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
8798                            return Err(fail(&format!(
8799                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
8800                                a.len,
8801                            )));
8802                        }
8803                        let kb = a.len * a.k_tok_bytes;
8804                        let vb = a.len * a.v_tok_bytes;
8805                        if kb > 0 {
8806                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
8807                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
8808                            if ak != bk {
8809                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
8810                                return Err(fail(&format!(
8811                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
8812                                    at / a.k_tok_bytes,
8813                                    at % a.k_tok_bytes,
8814                                    ak[at],
8815                                    bk[at],
8816                                )));
8817                            }
8818                        }
8819                        if vb > 0 {
8820                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
8821                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
8822                            if av != bv {
8823                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
8824                                return Err(fail(&format!(
8825                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
8826                                    at / a.v_tok_bytes,
8827                                    at % a.v_tok_bytes,
8828                                    av[at],
8829                                    bv[at],
8830                                )));
8831                            }
8832                        }
8833                        report.trunk_kv_bytes += kb + vb;
8834                    }
8835                    (None, None) => {}
8836                    _ => return Err(fail(&format!("layer {il} KV presence"))),
8837                }
8838                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
8839                    (Some(a), Some(b)) => {
8840                        let ac = es.dtoh(&a.conv_state)?;
8841                        let bc = es.dtoh(&b.conv_state)?;
8842                        if !same_f32(&ac, &bc) {
8843                            return Err(fail(&format!("layer {il} conv state")));
8844                        }
8845                        let as_ = es.dtoh(&a.ssm_state)?;
8846                        let bs = es.dtoh(&b.ssm_state)?;
8847                        if !same_f32(&as_, &bs) {
8848                            return Err(fail(&format!("layer {il} SSM state")));
8849                        }
8850                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
8851                    }
8852                    (None, None) => {}
8853                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
8854                }
8855            }
8856            Ok(())
8857        }
8858
8859        if reference.committed != candidate.committed {
8860            return Err(fail("committed token ids"));
8861        }
8862        if reference.cache.pos != candidate.cache.pos
8863            || reference.cache.max_ctx != candidate.cache.max_ctx
8864        {
8865            return Err(fail("cache pos/capacity"));
8866        }
8867        if reference.pending_tok != candidate.pending_tok
8868            || reference.next_pred != candidate.next_pred
8869            || reference.sctr != candidate.sctr
8870            || reference.uctr != candidate.uctr
8871        {
8872            return Err(fail("pending/prediction/counter tail"));
8873        }
8874
8875        let mut report = OptiForkStateIdentity::default();
8876        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
8877            let rt = crate::pp::PpNRt::get(e)?;
8878            for stage in 0..rt.n_stages() {
8879                let _scope = rt.enter(stage);
8880                compare_layers(
8881                    rt.engine(stage, e),
8882                    fence[stage]..fence[stage + 1],
8883                    reference,
8884                    candidate,
8885                    &mut report,
8886                )?;
8887            }
8888        } else {
8889            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
8890        }
8891
8892        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
8893            return Err(fail("draft scratch plane count"));
8894        }
8895        for index in 0..reference.scratch.plane_count() {
8896            let (a, _) = reference.scratch.plane(index);
8897            let (b, _) = candidate.scratch.plane(index);
8898            if a.len != b.len
8899                || a.kv_dim_k != b.kv_dim_k
8900                || a.kv_dim_v != b.kv_dim_v
8901                || a.k_tok_bytes != b.k_tok_bytes
8902                || a.v_tok_bytes != b.v_tok_bytes
8903                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
8904            {
8905                return Err(fail(&format!("draft scratch plane {index} length/layout")));
8906            }
8907            let kb = a.len * a.k_tok_bytes;
8908            let vb = a.len * a.v_tok_bytes;
8909            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
8910                return Err(fail(&format!("draft scratch plane {index} K bytes")));
8911            }
8912            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
8913                return Err(fail(&format!("draft scratch plane {index} V bytes")));
8914            }
8915            report.scratch_kv_bytes += kb + vb;
8916        }
8917
8918        match (&reference.last_h, &candidate.last_h) {
8919            (Some(a), Some(b)) => {
8920                let ah = e.dtoh(a)?;
8921                let bh = e.dtoh(b)?;
8922                if !same_f32(&ah, &bh) {
8923                    return Err(fail("last hidden/seed bytes"));
8924                }
8925                report.hidden_bytes = ah.len() * 4;
8926            }
8927            (None, None) => {}
8928            _ => return Err(fail("last hidden/seed presence")),
8929        }
8930        Ok(report)
8931    }
8932
8933    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
8934    /// retained prompt-end checkpoint, so a request whose prompt matches
8935    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
8936    ///
8937    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
8938    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
8939    /// restored from the device copy taken there, draft scratch length reset, `committed`
8940    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
8941    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
8942    /// every burst after it are identical to a cold run of the same token stream — the
8943    /// committed-tokens-authoritative contract.
8944    ///
8945    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
8946    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
8947    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
8948    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
8949    /// (the scratch KV, the resident embedding), none of which the rewind moves.
8950    ///
8951    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
8952    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
8953    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
8954    pub fn spec_rewind_to_checkpoint(
8955        &self,
8956        e: &Engine,
8957        sess: &mut SpecSession,
8958    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
8959        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
8960            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
8961        }) {
8962            return Err(
8963                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
8964            );
8965        }
8966        let Some(ckpt) = sess.turn_ckpt.take() else {
8967            return Ok(None);
8968        };
8969        assert!(
8970            ckpt.pos <= sess.committed.len(),
8971            "checkpoint past committed ({} > {})",
8972            ckpt.pos,
8973            sess.committed.len()
8974        );
8975        // Restore through each layer's owning engine. A single primary-engine rollback is not
8976        // sufficient when the serving cache is stage-owned under cross-device PP.
8977        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
8978        debug_assert_eq!(
8979            sess.cache.pos, ckpt.pos,
8980            "rollback landed off the checkpoint"
8981        );
8982        sess.scratch.set_len(e, ckpt.pos)?;
8983        sess.committed.truncate(ckpt.pos);
8984        sess.last_h = Some(ckpt.last_h);
8985        sess.next_pred = None;
8986        sess.pending_tok = None;
8987        Ok(Some(ckpt.pos))
8988    }
8989
8990    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
8991    /// checkpoint without re-priming the checkpoint prefix.
8992    ///
8993    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
8994    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
8995    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
8996    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
8997    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
8998    ///
8999    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9000    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9001    pub fn spec_grow_and_rewind_to_checkpoint(
9002        &self,
9003        e: &Engine,
9004        sess: &mut SpecSession,
9005        target_cap: usize,
9006    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9007        if target_cap <= sess.cache.max_ctx {
9008            return self.spec_rewind_to_checkpoint(e, sess);
9009        }
9010        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9011            return Ok(None);
9012        };
9013        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9014            return Err(format!(
9015                "checkpoint pos {} outside committed length {}",
9016                ckpt.pos,
9017                sess.committed.len(),
9018            )
9019            .into());
9020        }
9021        if ckpt.pos > target_cap {
9022            return Err(format!(
9023                "checkpoint pos {} exceeds grown capacity {target_cap}",
9024                ckpt.pos,
9025            )
9026            .into());
9027        }
9028
9029        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9030        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9031        crate::pp::restore_cache_checkpoint(
9032            e,
9033            self,
9034            Some(&sess.cache),
9035            &mut grown_cache,
9036            &ckpt.snap,
9037        )?;
9038
9039        if sess.scratch.plane_count() != grown_scratch.plane_count() {
9040            return Err("checkpoint draft plane count mismatch".into());
9041        }
9042        for index in 0..sess.scratch.plane_count() {
9043            let (src, _) = sess.scratch.plane(index);
9044            let (dst, _) = grown_scratch.plane_mut(index);
9045            if ckpt.pos > src.len
9046                || src.kv_dim_k != dst.kv_dim_k
9047                || src.kv_dim_v != dst.kv_dim_v
9048                || src.k_tok_bytes != dst.k_tok_bytes
9049                || src.v_tok_bytes != dst.v_tok_bytes
9050            {
9051                return Err(format!(
9052                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9053                    ckpt.pos, src.len,
9054                )
9055                .into());
9056            }
9057            let kb = ckpt.pos * src.k_tok_bytes;
9058            let vb = ckpt.pos * src.v_tok_bytes;
9059            if kb > 0 {
9060                e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9061            }
9062            if vb > 0 {
9063                e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9064            }
9065        }
9066        grown_scratch.set_len(e, ckpt.pos)?;
9067        // The old scratch is dropped immediately after publication below. Bound its D2D reads
9068        // first; growth happens once per rewritten turn, outside the decode hot loop.
9069        e.stream().synchronize()?;
9070
9071        let ckpt = sess
9072            .turn_ckpt
9073            .take()
9074            .expect("checkpoint remained present through transactional grow");
9075        let pos = ckpt.pos;
9076        sess.cache = grown_cache;
9077        sess.scratch = grown_scratch;
9078        sess.committed.truncate(pos);
9079        sess.last_h = Some(ckpt.last_h);
9080        sess.next_pred = None;
9081        sess.pending_tok = None;
9082        sess.draft_ctx = None;
9083        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9084        debug_assert!(
9085            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9086            "grown draft rewind landed off checkpoint"
9087        );
9088        Ok(Some(pos))
9089    }
9090
9091    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9092    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9093    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9094    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9095    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9096    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9097    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9098    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9099    /// park-time flush is a future request whose sampler is not knowable here (residual
9100    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9101    pub fn spec_flush_pending(
9102        &self,
9103        e: &Engine,
9104        sess: &mut SpecSession,
9105        sampling: Option<SpecSampling>,
9106    ) -> Result<(), Box<dyn std::error::Error>> {
9107        let Some(b) = sess.pending_tok.take() else {
9108            return Ok(());
9109        };
9110        if self.mtp.is_none() {
9111            return Err("pending carry requires an MTP head".into());
9112        }
9113        let n_embd = self.cfg.n_embd as usize;
9114        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9115        let embd_gpu = if spec_host_embd() {
9116            None
9117        } else {
9118            Some(
9119                self.embd_gpu
9120                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9121            )
9122        };
9123        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9124        let pos_b = sess.cache.pos;
9125        sess.scratch.set_len(e, pos_b)?;
9126        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9127        sess.next_pred = Some(match sampling {
9128            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9129                // window includes `b` itself: it is committed by this pass, and the pre-lane
9130                // code never counted a boundary token in the penalty history at all.
9131                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9132                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9133            }
9134            _ => argmax(&lg_b) as u32,
9135        });
9136        let anchor = sess
9137            .last_h
9138            .as_ref()
9139            .expect("pending carry requires last_h (the predecessor-row anchor)");
9140        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9141        sess.last_h = Some(hb);
9142        sess.committed.push(b);
9143        Ok(())
9144    }
9145
9146    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9147    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9148    /// rounds through that same graph. Other model families keep their eager T=1 contract.
9149    fn spec_target_step_h(
9150        &self,
9151        e: &Engine,
9152        token: u32,
9153        cache: &mut Cache,
9154    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9155        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9156            return self.decode_step_h(e, token, cache);
9157        }
9158        let pos0 = cache.pos;
9159        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9160        Ok((e.dtoh(&logits)?, hidden))
9161    }
9162
9163    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9164    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9165    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9166    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9167    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9168    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9169    /// dispatch sites cannot drift apart again.
9170    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9171    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
9172    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
9173    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
9174    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
9175    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9176    fn mtp_graph_capturable(&self) -> bool {
9177        self.mtp
9178            .as_ref()
9179            .map(|m| match &m.ffn {
9180                crate::hybrid::Ffn::Dense { .. } => true,
9181                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
9182            })
9183            .unwrap_or(false)
9184    }
9185
9186    fn batched_serving_numeric_class(&self) -> bool {
9187        self.plan
9188            .trunk_operations()
9189            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9190    }
9191
9192    fn sliding_gated_moe_batch_program(&self) -> bool {
9193        self.uses_sliding_gated_moe_program()
9194    }
9195
9196    fn gemma_batch_program(&self) -> bool {
9197        self.uses_gemma_program()
9198    }
9199
9200    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
9201    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
9202    /// session already exist.
9203    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
9204        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
9205            || !spec_devacc()
9206            || spec_replay_env_enabled()
9207            || spec_stream()
9208            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
9209            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
9210            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
9211            || std::env::var("MEMRA_SPEC_PMIN")
9212                .ok()
9213                .and_then(|v| v.parse::<f32>().ok())
9214                .unwrap_or(0.0)
9215                > 0.0
9216            || self.is_gemma4_e4b()
9217            || self.gemma_batch_program()
9218            || self.mtp.is_none()
9219            || !self.mtp_extra.is_empty()
9220        {
9221            return false;
9222        }
9223        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
9224            return false;
9225        };
9226        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
9227            return false;
9228        }
9229        crate::pp::PpNRt::get(e)
9230            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
9231            .unwrap_or(false)
9232    }
9233
9234    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
9235    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
9236    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
9237    #[allow(clippy::too_many_arguments)]
9238    pub fn generate_spec_session_pair(
9239        &self,
9240        e: &Engine,
9241        sess_a: &mut SpecSession,
9242        max_new_a: usize,
9243        k_a: usize,
9244        sess_b: &mut SpecSession,
9245        max_new_b: usize,
9246        k_b: usize,
9247    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
9248    {
9249        if !self.spec_pipe_available(e) {
9250            return Err("two-session speculative pipeline is outside its reduced matrix".into());
9251        }
9252        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
9253            return Err(
9254                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
9255            );
9256        }
9257        for sess in [&*sess_a, &*sess_b] {
9258            if sess.committed.is_empty()
9259                || sess.last_h.is_none()
9260                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
9261            {
9262                return Err("two-session speculative pipeline requires warm continuations".into());
9263            }
9264        }
9265
9266        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9267            && !spec_host_embd()
9268            && self.mtp_graph_capturable()
9269            && self.mtp_extra.is_empty()
9270            && !crate::model::full_prec_enabled();
9271        let graph_a = graph_ok && k_a + 2 < 96;
9272        let graph_b = graph_ok && k_b + 2 < 96;
9273        let was_tracking = e.ctx().is_event_tracking();
9274        if (graph_a || graph_b) && was_tracking {
9275            unsafe {
9276                e.ctx().disable_event_tracking();
9277            }
9278        }
9279
9280        static LOGGED: std::sync::Once = std::sync::Once::new();
9281        LOGGED.call_once(|| {
9282            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
9283        });
9284        let sync = std::sync::Arc::new(SpecPipeSync::new());
9285        let lane_a = SpecPipeLane {
9286            sync: sync.clone(),
9287            lane: 0,
9288        };
9289        let lane_b = SpecPipeLane { sync, lane: 1 };
9290        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
9291        let (result_a, result_b) = std::thread::scope(|scope| {
9292            let b = scope.spawn(move || {
9293                let mut finish = SpecPipeFinish::new(&lane_b);
9294                let sess_b = unsafe { sess_b_ptr.get_mut() };
9295                let result = e
9296                    .ctx()
9297                    .bind_to_thread()
9298                    .map_err(|err| err.to_string())
9299                    .and_then(|_| {
9300                        self.generate_spec_inner2(
9301                            e,
9302                            &[],
9303                            max_new_b,
9304                            k_b,
9305                            graph_b,
9306                            Some(sess_b),
9307                            None,
9308                            None,
9309                            None,
9310                            None,
9311                            Some(&lane_b),
9312                        )
9313                        .map_err(|err| err.to_string())
9314                    });
9315                finish.close(result.is_err());
9316                result
9317            });
9318            let mut finish = SpecPipeFinish::new(&lane_a);
9319            let result_a = self.generate_spec_inner2(
9320                e,
9321                &[],
9322                max_new_a,
9323                k_a,
9324                graph_a,
9325                Some(sess_a),
9326                None,
9327                None,
9328                None,
9329                None,
9330                Some(&lane_a),
9331            );
9332            finish.close(result_a.is_err());
9333            let result_b = b
9334                .join()
9335                .map_err(|_| "paired speculative session B panicked".to_string())
9336                .and_then(|r| r);
9337            (result_a, result_b)
9338        });
9339
9340        if (graph_a || graph_b) && was_tracking {
9341            unsafe {
9342                e.ctx().enable_event_tracking();
9343            }
9344        }
9345        let result_a = result_a?;
9346        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
9347        Ok((result_a, result_b))
9348    }
9349
9350    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9351    /// message rendered through the chat template continuation). Returns (new tokens emitted,
9352    /// drafted, accepted); session.committed grows by suffix + emitted.
9353    pub fn generate_spec_session(
9354        &self,
9355        e: &Engine,
9356        sess: &mut SpecSession,
9357        suffix: &[u32],
9358        max_new: usize,
9359        k: usize,
9360    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9361        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9362    }
9363
9364    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9365    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9366    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9367    /// for the filtered target (feat/filtered-spec).
9368    ///
9369    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9370    /// output — once right after the prime's first token, then once per round commit — so a
9371    /// streaming caller can flush text at round cadence instead of once per burst. The slices
9372    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9373    /// timing only: token bytes, session state, and exactness are untouched.
9374    ///
9375    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9376    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9377    /// the caller's scheduler regains control without waiting the burst out. Burst size is
9378    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9379    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9380    /// drains and the defensive tail flush can land with nothing new committed).
9381    #[allow(clippy::too_many_arguments)]
9382    pub fn generate_spec_session_sampled(
9383        &self,
9384        e: &Engine,
9385        sess: &mut SpecSession,
9386        suffix: &[u32],
9387        max_new: usize,
9388        k: usize,
9389        sampling: Option<SpecSampling>,
9390        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9391    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9392        self.generate_spec_session_sampled_prime_split(
9393            e, sess, suffix, max_new, k, sampling, None, on_commit,
9394        )
9395    }
9396
9397    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9398    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9399    /// pass `None` and stay on the existing zero-prime path.
9400    #[allow(clippy::too_many_arguments)]
9401    pub fn generate_spec_session_sampled_prime_split(
9402        &self,
9403        e: &Engine,
9404        sess: &mut SpecSession,
9405        suffix: &[u32],
9406        max_new: usize,
9407        k: usize,
9408        sampling: Option<SpecSampling>,
9409        prime_split: Option<usize>,
9410        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9411    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9412        self.generate_spec_session_constrained_prime_split(
9413            e,
9414            sess,
9415            suffix,
9416            max_new,
9417            k,
9418            sampling,
9419            None,
9420            prime_split,
9421            on_commit,
9422        )
9423    }
9424
9425    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9426    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9427    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9428    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9429    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9430    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9431    /// may drop (drafter is unconstrained); that is measured, not hidden.
9432    #[allow(clippy::too_many_arguments)]
9433    pub fn generate_spec_session_constrained(
9434        &self,
9435        e: &Engine,
9436        sess: &mut SpecSession,
9437        suffix: &[u32],
9438        max_new: usize,
9439        k: usize,
9440        sampling: Option<SpecSampling>,
9441        constraint: Option<&mut dyn SpecConstraint>,
9442        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9443    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9444        self.generate_spec_session_constrained_prime_split(
9445            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9446        )
9447    }
9448
9449    #[allow(clippy::too_many_arguments)]
9450    pub fn generate_spec_session_constrained_prime_split(
9451        &self,
9452        e: &Engine,
9453        sess: &mut SpecSession,
9454        suffix: &[u32],
9455        max_new: usize,
9456        k: usize,
9457        sampling: Option<SpecSampling>,
9458        constraint: Option<&mut dyn SpecConstraint>,
9459        prime_split: Option<usize>,
9460        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9461    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9462        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9463            return Err(
9464                "constrained spec decode is greedy-only (worker routes sampled \
9465                        constrained to plain decode)"
9466                    .into(),
9467            );
9468        }
9469        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9470        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9471        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9472        // serve continuation case — consume the carry in-loop with zero solo passes.
9473        if sess.pending_tok.is_some()
9474            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
9475        {
9476            self.spec_flush_pending(e, sess, sampling)?;
9477        }
9478
9479        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9480        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9481        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9482        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9483            && !spec_host_embd()
9484            && self.mtp_graph_capturable()
9485            && self.mtp_extra.is_empty()
9486            && k + 2 < 96
9487            && !crate::model::full_prec_enabled();
9488        let was_tracking = e.ctx().is_event_tracking();
9489        if graph_draft && was_tracking {
9490            unsafe {
9491                e.ctx().disable_event_tracking();
9492            }
9493        }
9494        let r = self.generate_spec_inner2(
9495            e,
9496            suffix,
9497            max_new,
9498            k,
9499            graph_draft,
9500            Some(sess),
9501            sampling,
9502            constraint,
9503            on_commit,
9504            prime_split,
9505            None,
9506        );
9507        if graph_draft && was_tracking {
9508            unsafe {
9509                e.ctx().enable_event_tracking();
9510            }
9511        }
9512        let (out, d, a) = r?;
9513        Ok((out, d, a))
9514    }
9515
9516    pub fn generate_spec(
9517        &self,
9518        e: &Engine,
9519        prompt: &[u32],
9520        max_new: usize,
9521        k: usize,
9522    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9523        if crate::pp::pp_cuts(self.layers.len()).is_some()
9524            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9525        {
9526            return Err("pipeline rewrite is not qualified for speculative decode".into());
9527        }
9528        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9529            return Err("speculative rewrite is not qualified for this ModelPlan".into());
9530        }
9531        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9532        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9533        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9534            && !spec_host_embd()
9535            && self.mtp_graph_capturable()
9536            && self.mtp_extra.is_empty()
9537            && k + 2 < 96
9538            && !crate::model::full_prec_enabled();
9539        if !graph_draft {
9540            return self.generate_spec_inner2(
9541                e, prompt, max_new, k, false, None, None, None, None, None, None,
9542            );
9543        }
9544        let was_tracking = e.ctx().is_event_tracking();
9545        if was_tracking {
9546            unsafe {
9547                e.ctx().disable_event_tracking();
9548            }
9549        }
9550        let r = self.generate_spec_inner2(
9551            e, prompt, max_new, k, true, None, None, None, None, None, None,
9552        );
9553        if was_tracking {
9554            unsafe {
9555                e.ctx().enable_event_tracking();
9556            }
9557        }
9558        r
9559    }
9560
9561    fn generate_spec_inner2(
9562        &self,
9563        e: &Engine,
9564        prompt: &[u32],
9565        max_new: usize,
9566        k: usize,
9567        graph_draft: bool,
9568        mut sess: Option<&mut SpecSession>,
9569        sampling: Option<SpecSampling>,
9570        mut constraint: Option<&mut dyn SpecConstraint>,
9571        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9572        prime_split: Option<usize>,
9573        pipe: Option<&SpecPipeLane>,
9574    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9575        assert!(k >= 1, "k must be >= 1");
9576        if let Some(p) = pipe {
9577            p.setup_begin()?;
9578        }
9579        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9580        let mut flushed = 0usize;
9581        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9582        // at the next round boundary (same exit as max_new reached — the session tail runs).
9583        // Initialized by the unconditional post-prime flush below.
9584        let mut keep_going;
9585        let mtp = self
9586            .mtp
9587            .as_ref()
9588            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9589        let n_vocab = self.output.out_features();
9590        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9591        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9592        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9593        let d_vocab = mtp
9594            .shared_head_head
9595            .as_ref()
9596            .unwrap_or(&self.output)
9597            .out_features();
9598        if !self.mtp_extra.is_empty() {
9599            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
9600                || self.plan.mtp_blocks.len() != self.mtp_head_count()
9601                || mtp.d2t.is_some()
9602            {
9603                return Err(
9604                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
9605                );
9606            }
9607            for (offset, head) in self.mtp_extra.iter().enumerate() {
9608                if head.d2t.is_some()
9609                    || head
9610                        .shared_head_head
9611                        .as_ref()
9612                        .unwrap_or(&self.output)
9613                        .out_features()
9614                        != d_vocab
9615                {
9616                    return Err(format!(
9617                        "embedded MTP head {} has incompatible draft vocabulary",
9618                        offset + 1
9619                    )
9620                    .into());
9621                }
9622            }
9623            eprintln!(
9624                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
9625                self.mtp_head_count()
9626            );
9627        }
9628        let n_embd = self.cfg.n_embd as usize;
9629        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
9630        // already committed (their state is in the caches); 0 = fresh single-shot call.
9631        let session_mode = sess.is_some();
9632        let max_ctx = match sess.as_ref() {
9633            Some(s) => s.cache.max_ctx,
9634            None => prompt.len() + max_new + k + 8,
9635        };
9636        let mut own_cache;
9637        let mut own_scratch;
9638        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
9639        // (requested split, destination list). Single-shot per burst; fresh calls have none.
9640        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
9641        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
9642        // committed-length position; consumed one-shot like `capture_at`. None = legacy
9643        // prompt-end capture below.
9644        let mut ckpt_req: Option<usize> = None;
9645        let (
9646            cache,
9647            scratch,
9648            mut sess_tail,
9649            mut sess_draft_slot,
9650            mut sess_pending_slot,
9651            sess_ckpt_slot,
9652            sess_telem,
9653        ): (
9654            &mut Cache,
9655            &mut MtpScratch,
9656            Option<(
9657                &mut Vec<u32>,
9658                &mut Option<CudaSlice<f32>>,
9659                &mut Option<u32>,
9660                &mut u32,
9661                &mut u32,
9662            )>,
9663            Option<&mut Option<DraftGraphCtx>>,
9664            Option<&mut Option<u32>>,
9665            Option<&mut Option<SpecCheckpoint>>,
9666            Option<&SpecTelemetryCounters>,
9667        ) = match sess.take() {
9668            Some(sr) => {
9669                let SpecSession {
9670                    cache,
9671                    scratch,
9672                    committed,
9673                    last_h,
9674                    next_pred,
9675                    sctr: s_sctr,
9676                    uctr: s_uctr,
9677                    draft_ctx,
9678                    pending_tok,
9679                    turn_ckpt,
9680                    telem,
9681                    capture_at,
9682                    boundary_captures,
9683                    ckpt_at,
9684                } = sr;
9685                sess_capture = Some((capture_at.take(), boundary_captures));
9686                ckpt_req = ckpt_at.take();
9687                (
9688                    cache,
9689                    scratch,
9690                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
9691                    Some(draft_ctx),
9692                    Some(pending_tok),
9693                    Some(turn_ckpt),
9694                    Some(telem),
9695                )
9696            }
9697            None => {
9698                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
9699                // `Cache::new` verbatim.
9700                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
9701                // Persistent scratch = max_ctx rows (~2KB/token quantized).
9702                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
9703                (
9704                    &mut own_cache,
9705                    &mut own_scratch,
9706                    None,
9707                    None,
9708                    None,
9709                    None,
9710                    None,
9711                )
9712            }
9713        };
9714        if scratch.plane_count() != self.mtp_head_count() {
9715            return Err(format!(
9716                "MTP scratch/head count mismatch ({}/{})",
9717                scratch.plane_count(),
9718                self.mtp_head_count()
9719            )
9720            .into());
9721        }
9722        let base = cache.pos;
9723        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
9724        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
9725        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
9726        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
9727        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
9728        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
9729        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
9730        // acceptance-only — exactness is verify's job either way).
9731        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
9732        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
9733        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
9734        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
9735        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
9736        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
9737        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
9738        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
9739        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
9740        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
9741        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
9742        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
9743        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
9744        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
9745        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
9746        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
9747        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
9748        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
9749        // + fallback seam).
9750        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
9751        // bar — the retained verify-state commit proven equivalent to sequential serving —
9752        // was waiting on this arch running the serving batched verify class, which the
9753        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
9754        // replay-free commit consumes is now produced by the SAME serving-class verify that
9755        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
9756        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
9757        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
9758        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
9759        // rollback + A/B seam.
9760        let spec_replay = spec_replay_env_enabled();
9761        if constraint.is_some() && spec_replay {
9762            return Err(
9763                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
9764                        (legacy replay commits an unmasked bonus)"
9765                    .into(),
9766            );
9767        }
9768        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
9769        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
9770        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
9771        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
9772        if !refresh && !self.mtp_extra.is_empty() {
9773            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
9774        }
9775
9776        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
9777        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
9778        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
9779        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
9780        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
9781        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
9782        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
9783        // generation exactly where the last turn stopped — no prime at all. The stashed
9784        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
9785        // committed.last() by the same rule this entry applies to a cold prime's last row —
9786        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
9787        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
9788        // where the sampler and the session's Philox counters were live). `last_h` seeds the
9789        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
9790        let continuation = prompt.is_empty();
9791        if continuation {
9792            assert!(session_mode, "empty prompt requires a session");
9793            assert!(
9794                sess_tail
9795                    .as_ref()
9796                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
9797                        && lh.is_some()
9798                        && (np.is_some() || carried_pending.is_some())),
9799                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
9800            );
9801        }
9802        let mut prime_logits;
9803        let mut prompt_h: Option<CudaSlice<f32>> = None;
9804        let t_prime = std::time::Instant::now();
9805        let batched_prime = !continuation
9806            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
9807            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9808            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
9809        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
9810        if prime_split.is_some() && continuation {
9811            return Err("spec prime split requires a non-empty prime".into());
9812        }
9813        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
9814        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
9815        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
9816        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
9817        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
9818        // cannot honor (outside this prime's range) silently drops the capture — the
9819        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
9820        let ckpt_rel = if continuation {
9821            None
9822        } else {
9823            ckpt_req
9824                .and_then(|abs| abs.checked_sub(base))
9825                .filter(|&r| r > 0 && r < prompt.len())
9826        };
9827        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
9828        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
9829        // the legacy single-split program, byte-for-byte.
9830        let mut stops: Vec<usize> = Vec::new();
9831        for b in [prime_split, ckpt_rel].into_iter().flatten() {
9832            if !stops.contains(&b) {
9833                stops.push(b);
9834            }
9835        }
9836        stops.sort_unstable();
9837        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
9838        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
9839        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
9840        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
9841        if continuation {
9842            prime_logits = Vec::new();
9843        } else if !stops.is_empty() {
9844            if let Some(&first) = stops.first() {
9845                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
9846                    return Err(format!(
9847                        "spec prime split {first} is below PRIME_MIN_T {}",
9848                        crate::hybrid_forward::PRIME_MIN_T,
9849                    )
9850                    .into());
9851                }
9852            }
9853            // Mirror the plain worker's boundary stops exactly. Each segment is a
9854            // request-level prime (`queued_after` keeps Step35 arm selection independent of
9855            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
9856            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
9857            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
9858            // coherent prompt.
9859            let mut h_all = e.uninit(prompt.len() * n_embd)?;
9860            prime_logits = Vec::new();
9861            let mut prev = 0usize;
9862            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
9863                if seg_end <= prev {
9864                    continue;
9865                }
9866                let seg = &prompt[prev..seg_end];
9867                let is_final = seg_end == prompt.len();
9868                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
9869                    && (!is_final
9870                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
9871                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
9872                if batched_seg {
9873                    let (l, _, h_seg) =
9874                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
9875                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
9876                    prime_logits = l;
9877                } else {
9878                    for (i, &tok) in seg.iter().enumerate() {
9879                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
9880                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
9881                        prime_logits = l;
9882                    }
9883                }
9884                prev = seg_end;
9885                if is_final {
9886                    break;
9887                }
9888                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
9889                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
9890                // states are about to be advanced in place by the next segment, so this is
9891                // the ONLY moment the boundary's recurrent state exists. Capture iff the
9892                // worker requested exactly this stop (cold sessions only — `capture_at` is
9893                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
9894                // publication is an optimization, never a correctness dependency.
9895                if base == 0 {
9896                    if let Some((requested, slot)) = sess_capture.as_mut() {
9897                        // Publish at the requested miss-LCP stop (the shared-prefix class)
9898                        // AND at the stable-boundary stop (the next-turn re-render class,
9899                        // lane/frspec-multiturn-cache) — the same boundary set the plain
9900                        // prefill tick learns. Without the second entry, the turn after a
9901                        // cold re-park could only hit the OLDER lcp entry (the measured
9902                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
9903                        // rewound to 15222). Dedupe is the worker sweep's has_key.
9904                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
9905                            if let Ok(snap) = cache.snapshot(e) {
9906                                slot.push(SpecBoundaryCapture {
9907                                    snap,
9908                                    pos: seg_end,
9909                                    logits: prime_logits.clone(),
9910                                    // rows [0..seg_end) of h_all are primed — the following
9911                                    // segments append, never overwrite.
9912                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
9913                                });
9914                            }
9915                        }
9916                    }
9917                }
9918                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
9919                // same snapshot mechanics, installed post-prime in place of the prompt-end
9920                // capture the re-render class always diverged below.
9921                if ckpt_rel == Some(seg_end) {
9922                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9923                        e.uninit(n_embd).and_then(|mut a| {
9924                            e.copy_view_into(
9925                                &mut a,
9926                                0,
9927                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9928                                n_embd,
9929                            )?;
9930                            Ok(a)
9931                        });
9932                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
9933                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
9934                            snap,
9935                            pos: base + seg_end,
9936                            last_h,
9937                        }),
9938                        _ => None,
9939                    });
9940                }
9941            }
9942            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
9943                eprintln!(
9944                    "[spec-prime] stops={stops:?} tail={}",
9945                    prompt.len() - stops.last().copied().unwrap_or(0)
9946                );
9947            }
9948            prompt_h = Some(h_all);
9949        } else if batched_prime {
9950            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
9951            prime_logits = l;
9952            prompt_h = Some(hiddens);
9953        } else {
9954            prime_logits = Vec::new();
9955            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
9956            for (i, &tok) in prompt.iter().enumerate() {
9957                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
9958                if let Some(ph) = prompt_h.as_mut() {
9959                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
9960                }
9961                prime_logits = l;
9962            }
9963        }
9964        e.stream().synchronize()?;
9965        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
9966        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
9967        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
9968        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
9969        // prime_split. The mid-prompt capture above already consumed the request if it matched.
9970        if !continuation && base == 0 {
9971            if let Some((requested, slot)) = sess_capture.as_mut() {
9972                if *requested == Some(prompt.len()) && slot.is_empty() {
9973                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
9974                    if let Ok(snap) = cache.snapshot(e) {
9975                        slot.push(SpecBoundaryCapture {
9976                            snap,
9977                            pos: prompt.len(),
9978                            logits: prime_logits.clone(),
9979                            last_h: prompt_h
9980                                .as_ref()
9981                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
9982                                .unwrap_or_default(),
9983                        });
9984                    }
9985                }
9986            }
9987        }
9988        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
9989        // prime-subtraction hack.
9990        crate::PRIME_NANOS.store(
9991            t_prime.elapsed().as_nanos() as u64,
9992            std::sync::atomic::Ordering::Relaxed,
9993        );
9994
9995        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9996        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
9997        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
9998        let host_embd = spec_host_embd();
9999        let embd_gpu = if host_embd {
10000            None
10001        } else {
10002            Some(
10003                self.embd_gpu
10004                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10005            )
10006        };
10007        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10008        if host_embd {
10009            eprintln!(
10010                "[spec] host-row embedding: {} bytes kept off HBM",
10011                self.embd.raw.len()
10012            );
10013        }
10014        let mut out: Vec<u32> = Vec::with_capacity(max_new);
10015        let mut total_drafted = 0usize;
10016        let mut total_accepted = 0usize;
10017
10018        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10019        // The sampler config, the session's Philox counters and the penalty window are parsed
10020        // HERE, above the boundary-token selection, because the boundary token must be drawn
10021        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10022        // selection, which is the whole mechanical reason the boundary token was an argmax:
10023        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10024        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10025        // below takes the argmax path it always took).
10026        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10027        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10028        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10029        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10030        let sp = sampling.unwrap_or_else(|| SpecSampling {
10031            temp: std::env::var("MEMRA_SPEC_TEMP")
10032                .ok()
10033                .and_then(|v| v.parse().ok())
10034                .unwrap_or(0.0),
10035            seed: std::env::var("MEMRA_SEED")
10036                .ok()
10037                .and_then(|v| v.parse().ok())
10038                .unwrap_or(42),
10039            top_k: std::env::var("MEMRA_TOP_K")
10040                .ok()
10041                .and_then(|v| v.parse().ok())
10042                .unwrap_or(0),
10043            top_p: std::env::var("MEMRA_TOP_P")
10044                .ok()
10045                .and_then(|v| v.parse().ok())
10046                .unwrap_or(1.0),
10047            min_p: std::env::var("MEMRA_MIN_P")
10048                .ok()
10049                .and_then(|v| v.parse().ok())
10050                .unwrap_or(0.0),
10051            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10052                .ok()
10053                .and_then(|v| v.parse().ok())
10054                .unwrap_or(0),
10055            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10056                .ok()
10057                .and_then(|v| v.parse().ok())
10058                .unwrap_or(1.0),
10059            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10060                .ok()
10061                .and_then(|v| v.parse().ok())
10062                .unwrap_or(0.0),
10063            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10064                .ok()
10065                .and_then(|v| v.parse().ok())
10066                .unwrap_or(0.0),
10067        });
10068        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10069        let sampled = sp_temp > 0.0;
10070        // Counters resume from the session (burst continuity: randomness must never repeat
10071        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10072        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10073        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10074        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10075        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10076        // for the penalized+filtered target). History = generated tokens, host-tracked window.
10077        let pen_on = sampled
10078            && sp.penalty_last_n > 0
10079            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10080        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10081        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10082        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10083        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10084        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10085        // which is what the API contract says and what the plain sampler's own `history` does.
10086        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10087        let mut pen_hist: Vec<u32> = if pen_on {
10088            let sess_hist: &[u32] = if spec_pen_session_on() {
10089                sess_tail
10090                    .as_ref()
10091                    .map(|(c, ..)| c.as_slice())
10092                    .unwrap_or(&[])
10093            } else {
10094                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10095            };
10096            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10097        } else {
10098            Vec::new()
10099        };
10100        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10101        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10102        // request's own filtered/penalized target through the session's Philox stream
10103        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10104        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10105        // Emit it, then FEED it to establish the loop invariant below.
10106        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10107        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10108        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10109        // prompt's last logits (plain constrained-greedy identity); a continuation without
10110        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10111        // worker never resumes constrained sessions from the pool, so this cannot fire).
10112        if let Some(c) = constraint.as_deref_mut() {
10113            if continuation && carried_pending.is_none() {
10114                return Err("constrained spec continuation requires a carried pending \
10115                            (pool resume is unconstrained-only)"
10116                    .into());
10117            }
10118            if !continuation {
10119                c.mask_logits(&mut prime_logits)
10120                    .map_err(|e2| format!("constraint: {e2}"))?;
10121            }
10122        }
10123        let mut last_token = if let Some(b) = carried_pending {
10124            b
10125        } else if continuation {
10126            // A continuation's boundary token was DRAWN by the burst that stashed it (the
10127            // session tail below), or by `spec_session_from_restored` for a converted
10128            // prefix-cache hit — in both cases from the correct logits row with this same
10129            // session's Philox stream, which is why it can be consumed here as-is.
10130            sess_tail.as_ref().unwrap().2.unwrap()
10131        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10132            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10133        } else {
10134            // greedy (byte contract), the rollback door, or constrained (masked-argmax
10135            // identity — the worker routes sampled+constrained to the plain path, and this
10136            // function refuses the combination outright above).
10137            argmax(&prime_logits) as u32
10138        };
10139        if pen_on {
10140            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10141            // emitted token into its penalty history, and pre-lane the burst's first token
10142            // was invisible to penalties forever (never pushed, and never in `committed`
10143            // until this burst's tail). Covers the carry/continuation seeds too — neither is
10144            // in `committed` yet.
10145            pen_hist.push(last_token);
10146        }
10147        if carried_pending.is_none() {
10148            out.push(last_token);
10149            // grammar advances with every emitted token (carried pendings were consumed
10150            // by the burst that emitted them).
10151            if let Some(c) = constraint.as_deref_mut() {
10152                c.consume(last_token)
10153                    .map_err(|e2| format!("constraint: {e2}"))?;
10154            }
10155        }
10156        if continuation {
10157            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10158            // overhang so the chain's first append lands at slot base (== committed.len()).
10159            scratch.set_len(e, base)?;
10160        }
10161        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10162        // concatenating to the full `out`). Called after the prime's first token and after each
10163        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10164        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10165        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10166        fn flush_commit(
10167            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10168            out: &[u32],
10169            flushed: &mut usize,
10170        ) -> bool {
10171            if let Some(f) = cb.as_mut() {
10172                let keep = f(&out[*flushed..]);
10173                *flushed = out.len();
10174                keep
10175            } else {
10176                true
10177            }
10178        }
10179        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10180        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10181        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10182        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10183        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10184        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10185        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10186        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10187        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10188        // those, so their residual mass is p(x), correct by construction).
10189        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10190            match &mtp.d2t {
10191                Some(map) => Some(e.htod_u32_v(map)?),
10192                None => None,
10193            }
10194        } else {
10195            None
10196        };
10197        let mut q_full_buf: Option<CudaSlice<f32>> = None;
10198        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10199        // dspark sampled-admission walk); byte-identical to the closure it replaces.
10200        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10201        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10202        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10203        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10204        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10205        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10206        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10207        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10208        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10209        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10210        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10211        let t_ent = std::time::Instant::now();
10212
10213        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10214        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10215        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10216        // the one that matters (a history-rewriting client mutates what the session GENERATED,
10217        // so the next turn's prompt agrees with this one up to exactly here).
10218        //
10219        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10220        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10221        // hold exactly `base + prompt.len()` rows and nothing generated.
10222        //
10223        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10224        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10225        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10226        // `<think>` block the client strips, so every later turn's diff diverged exactly one
10227        // token below the checkpoint and affinity declined 100% of the time. Measured on the
10228        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10229        // whole mechanism inert while looking, from the outside, like a working
10230        // correctness-declines-safely path — hence the decline log carries the offsets.
10231        //
10232        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10233        // state (the reason a spec session could not rewind before). The draft scratch needs no
10234        // copy: rows below the boundary are rewritten by the next turn's own fill.
10235        //
10236        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10237        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10238        // checkpoint rather than replacing it with a strictly worse one.
10239        //
10240        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10241        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10242        // fail the burst that is already running — so the error is swallowed, loud only under
10243        // MEMRA_DEBUG_SPEC.
10244        //
10245        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10246        // posture above was DISPROVED for the think-posture template class — the prompt's own
10247        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10248        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10249        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10250        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10251        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10252        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10253        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10254        if let Some(slot) = sess_ckpt_slot {
10255            if let Some(early) = ckpt_early {
10256                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10257                    eprintln!(
10258                        "[spec] stable-boundary turn checkpoint skipped; \
10259                               next turn re-primes in full"
10260                    );
10261                }
10262                *slot = early;
10263            } else if !continuation {
10264                let pos = cache.pos;
10265                debug_assert_eq!(
10266                    pos,
10267                    base + prompt.len(),
10268                    "turn checkpoint must sit at the prompt end, before the init feed"
10269                );
10270                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10271                    if let Some(ph) = &prompt_h {
10272                        // hidden of the LAST primed row = the predecessor anchor at this
10273                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
10274                        // last_h, and what the next prime's fill reads for its first row).
10275                        let np = prompt.len();
10276                        e.uninit(n_embd).and_then(|mut a| {
10277                            e.copy_view_into(
10278                                &mut a,
10279                                0,
10280                                &ph.slice((np - 1) * n_embd..np * n_embd),
10281                                n_embd,
10282                            )?;
10283                            Ok(a)
10284                        })
10285                    } else {
10286                        Err("no prompt hiddens".into())
10287                    };
10288                match (cache.snapshot(e), anchor) {
10289                    (Ok(snap), Ok(last_h)) => {
10290                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
10291                    }
10292                    (s, a) => {
10293                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10294                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10295                            let err = s
10296                                .err()
10297                                .map(|e| e.to_string())
10298                                .or_else(|| a.err().map(|e| e.to_string()))
10299                                .unwrap_or_default();
10300                            eprintln!(
10301                                "[spec] turn checkpoint skipped ({err}); \
10302                                       next turn re-primes in full"
10303                            );
10304                        }
10305                    }
10306                }
10307            }
10308        }
10309        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10310        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10311        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10312        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10313        let mut last_pred = 0u32;
10314        let mut last_col_logits: Option<CudaSlice<f32>> = None;
10315        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10316        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10317        let mut init_logits_host: Option<Vec<f32>> = None;
10318        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10319            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10320            last_pred = argmax(&init_logits) as u32;
10321            if constraint.is_some() {
10322                init_logits_host = Some(init_logits.clone());
10323            }
10324            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10325            if sampled {
10326                last_col_logits = Some(e.htod(&init_logits)?);
10327            }
10328            h
10329        } else {
10330            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10331            let lh = sess_tail
10332                .as_ref()
10333                .unwrap()
10334                .1
10335                .as_ref()
10336                .expect("pending carry requires last_h");
10337            e.clone_dtod(lh)?
10338        };
10339        let t_init = t_ent.elapsed();
10340        let mut last_col_stats: Option<(f32, f32, f32)> = None;
10341        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10342        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10343        // stable pointer for the graph-draft round-start copy.
10344        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10345        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10346        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10347        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10348        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10349        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10350        // overwritten below).
10351        let mut fill_prev = e.clone_dtod(&h_seed0)?;
10352        {
10353            if let Some(ph) = &prompt_h {
10354                let np = prompt.len();
10355                e.copy_view_into(
10356                    &mut h_seed_buf,
10357                    0,
10358                    &ph.slice((np - 1) * n_embd..np * n_embd),
10359                    n_embd,
10360                )?;
10361            } else if continuation {
10362                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10363                    if let Some(lh) = lh.as_ref() {
10364                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10365                    }
10366                }
10367            }
10368        }
10369        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10370        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10371
10372        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10373        let fork_mode = OptiForkGateMode::configured();
10374        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10375        // the end. Metric normalization vs the reference engine: BOTH engines count
10376        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10377        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10378        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10379        let mut st_drafted = vec![0usize; k];
10380        let mut st_accepted = vec![0usize; k];
10381        let mut st_len_hist = vec![0usize; k + 1];
10382        let mut st_full = 0usize;
10383        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10384        // stop the draft chain early when the head's softmax confidence in its own pick drops
10385        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10386        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10387        let p_min = *PMIN.get_or_init(|| {
10388            std::env::var("MEMRA_SPEC_PMIN")
10389                .ok()
10390                .and_then(|v| v.parse().ok())
10391                .unwrap_or(0.0)
10392        });
10393        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10394        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10395        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10396        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10397        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10398        // verify batch is not); the j==0 exemption stays for pending-less rounds.
10399        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10400            .map(|v| v == "1")
10401            .unwrap_or(false);
10402
10403        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10404        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10405        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10406        // cuBLAS path in an exotic head) falls back to the eager draft chain.
10407        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10408        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10409        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10410        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10411        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10412        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10413        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10414        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10415        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10416            Some(c) => c,
10417            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10418        };
10419        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10420        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10421        if sampled && dctx.g_q.len() < d_vocab {
10422            dctx.g_q = e.zeros(d_vocab)?;
10423            dctx.g_perturb = e.zeros(d_vocab)?;
10424        }
10425        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10426        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10427        // truncation (the correctness backstop) stops cutting every tight-schema round.
10428        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10429        // shape, so a parked graph of the other shape is dropped and recaptured.
10430        let dmask_on = constraint
10431            .as_deref()
10432            .is_some_and(|c| c.draft_mask_enabled());
10433        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10434        if dmask_on && dctx.g_dmask.len() < dmask_words {
10435            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10436            dctx.graph = None; // the old capture baked the old (or no) mask pointer
10437            dctx.failed.clear_greedy();
10438            dctx.keeper.clear();
10439        }
10440        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
10441            dctx.graph = None;
10442            dctx.failed.clear_greedy();
10443            dctx.keeper.clear();
10444        }
10445        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
10446            let DraftGraphCtx {
10447                g_tok,
10448                g_pos,
10449                g_seed,
10450                g_p,
10451                g_dmask,
10452                ..
10453            } = &mut dctx;
10454            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
10455            // host uploads the position's real words, so the warmups stay grammar-free.
10456            if dmask_on {
10457                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10458            }
10459            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10460            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
10461            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
10462            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
10463            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
10464            // passes (and, in serve, other sessions) recycle those addresses and the replay then
10465            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
10466            let cap_res = e.capture_graph_retained(|e| {
10467                self.mtp_head_forward_cap(
10468                    e,
10469                    mtp,
10470                    g_tok,
10471                    g_pos,
10472                    g_seed,
10473                    g_p,
10474                    &mut *scratch,
10475                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
10476                    true,
10477                    embd_gpu.expect("graph draft requires resident embedding"),
10478                    embd_qt,
10479                    embd_rb,
10480                    d_vocab,
10481                    None,
10482                    None,
10483                    if dmask_on {
10484                        Some((g_dmask_ro, dmask_words))
10485                    } else {
10486                        None
10487                    },
10488                )
10489            });
10490            match cap_res {
10491                Ok((g, keep)) => {
10492                    scratch.set_len(e, base)?;
10493                    dctx.graph = Some(g);
10494                    dctx.graph_masked = dmask_on;
10495                    dctx.keeper = keep;
10496                }
10497                Err(err) => {
10498                    scratch.set_len(e, base)?;
10499                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
10500                    // silent. Once per flip — mark returns None on an already-failed ctx.
10501                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
10502                        eprintln!("{line}");
10503                    }
10504                }
10505            }
10506        }
10507        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
10508        // graph object, built only when sampled && graph-eligible — the greedy capture above is
10509        // untouched (and skipped when sampled: its graph would never be launched). Same head
10510        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
10511        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
10512        // once per round); the raw head logits land in the persistent g_q for the host's
10513        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
10514        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
10515        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
10516        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
10517        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
10518        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
10519        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
10520        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
10521        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
10522        // this compare misses at most ONCE per resumed request — the first burst recaptures
10523        // and every later burst in that request replays. A client that wants the parked graph
10524        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
10525        // stable across its whole conversation.
10526        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
10527        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
10528        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
10529        // force the eager draft (which computes stats/penalties per row).
10530        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
10531        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
10532        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
10533        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
10534        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
10535        // the request shape the vendor-default flip makes the majority).
10536        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
10537        let pure_temp = s_key.pure_temp();
10538        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
10539            dctx.graph_s = None;
10540            dctx.failed.clear_sampled();
10541            dctx.s_key = None;
10542            dctx.q_slots.clear();
10543            dctx.keeper_s.clear();
10544        }
10545        if graph_draft
10546            && sampled
10547            && pure_temp
10548            && dctx.graph_s.is_none()
10549            && !dctx.failed.sampled_failed()
10550        {
10551            let DraftGraphCtx {
10552                g_tok,
10553                g_pos,
10554                g_seed,
10555                g_p,
10556                g_ctr,
10557                g_perturb,
10558                g_q,
10559                ..
10560            } = &mut dctx;
10561            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
10562            let cap_res = e.capture_graph_retained(|e| {
10563                self.mtp_head_forward_cap(
10564                    e,
10565                    mtp,
10566                    g_tok,
10567                    g_pos,
10568                    g_seed,
10569                    g_p,
10570                    &mut *scratch,
10571                    p_min > 0.0,
10572                    true,
10573                    embd_gpu.expect("graph draft requires resident embedding"),
10574                    embd_qt,
10575                    embd_rb,
10576                    d_vocab,
10577                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
10578                    None,
10579                    None, // constrained spec is greedy-only — sampled never carries a hook
10580                )
10581            });
10582            match cap_res {
10583                Ok((g, keep)) => {
10584                    scratch.set_len(e, base)?;
10585                    for _ in 0..k {
10586                        dctx.q_slots.push(e.zeros(d_vocab)?);
10587                    }
10588                    dctx.graph_s = Some(g);
10589                    dctx.s_key = Some(s_key);
10590                    dctx.keeper_s = keep;
10591                }
10592                Err(err) => {
10593                    scratch.set_len(e, base)?;
10594                    // LOUD flip (audit Q2): same contract as the greedy capture above.
10595                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
10596                        eprintln!("{line}");
10597                    }
10598                }
10599            }
10600        }
10601        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
10602        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
10603        // captured under this request's exact regime, and capture requires `pure_temp` — so a
10604        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
10605        // the graph arm, so it is asserted here rather than assumed: a future change that widens
10606        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
10607        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
10608        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
10609        // rather than launching it; the launch site re-tests `pure_temp` independently.
10610        if sampled && !pure_temp && dctx.graph_s.is_some() {
10611            debug_assert!(
10612                false,
10613                "sampled draft graph parked under {:?} survived into a FILTERED request \
10614                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
10615                 softmax, so the verify's filtered q would test a distribution the draft was \
10616                 never sampled from",
10617                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10618            );
10619            eprintln!(
10620                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
10621                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
10622                 EAGER — the key must carry every field that shapes q",
10623                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
10624            );
10625            dctx.graph_s = None;
10626            dctx.s_key = None;
10627            dctx.q_slots.clear();
10628            dctx.keeper_s.clear();
10629        }
10630        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
10631        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
10632        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
10633        // arms below print which chain actually ran, so the probe never restates the condition.
10634        if skey_probe() {
10635            eprintln!(
10636                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
10637                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
10638                sampled as u8,
10639                pure_temp as u8,
10640                sp_temp,
10641                sp.top_k,
10642                sp.top_p,
10643                sp.min_p,
10644                pen_on as u8,
10645                k,
10646                graph_draft as u8,
10647                dctx.graph_s.is_some() as u8,
10648                dctx.s_key,
10649            );
10650        }
10651        let t_cap = t_ent.elapsed();
10652        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
10653        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
10654        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
10655        // fill: the first chain step processes it and appends its entry at slot prompt.len().
10656        if let Some(ph) = &prompt_h {
10657            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
10658            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
10659            // global positions [base..base+tp). Fresh call: base==0, identical to before.
10660            scratch.set_len(e, base)?;
10661            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
10662            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
10663            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
10664            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
10665            let tp = prompt.len();
10666            let fill_chunk: usize = if crate::cache::swa_ring_on() {
10667                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
10668            } else {
10669                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
10670                // meaning one monolithic fill.
10671                std::env::var("MEMRA_PRIME_CHUNK")
10672                    .ok()
10673                    .and_then(|v| v.parse().ok())
10674                    .unwrap_or(4096)
10675            };
10676            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
10677            let mut start = 0usize;
10678            while start < tp {
10679                let end = (start + fill_chunk).min(tp);
10680                let tc = end - start;
10681                {
10682                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
10683                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
10684                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
10685                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
10686                    let mut phs = e.zeros(tc * n_embd)?;
10687                    let (src_lo, dst_off) = if start == 0 {
10688                        (0, n_embd)
10689                    } else {
10690                        ((start - 1) * n_embd, 0)
10691                    };
10692                    let n_copy = if start == 0 {
10693                        (tc - 1) * n_embd
10694                    } else {
10695                        tc * n_embd
10696                    };
10697                    if start == 0 {
10698                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10699                            if let Some(lh) = lh.as_ref() {
10700                                e.copy_into(&mut phs, 0, lh, n_embd)?;
10701                            }
10702                        }
10703                    }
10704                    if n_copy > 0 {
10705                        e.copy_view_into(
10706                            &mut phs,
10707                            dst_off,
10708                            &ph.slice(src_lo..src_lo + n_copy),
10709                            n_copy,
10710                        )?;
10711                    }
10712                    self.mtp_kv_fill_all(
10713                        e,
10714                        &prompt[start..end],
10715                        &phs,
10716                        base + start,
10717                        &mut *scratch,
10718                        embd_dev,
10719                    )?;
10720                }
10721                start = end;
10722            }
10723        }
10724        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
10725        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
10726        // (=1 brackets the whole call in run_spec.rs, prime included.)
10727        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
10728            unsafe extern "C" {
10729                fn cudaProfilerStart() -> i32;
10730            }
10731            unsafe {
10732                cudaProfilerStart();
10733            }
10734        }
10735        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
10736        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
10737        // consume each other's device outputs; the host drains the ring every M rounds. v1
10738        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
10739        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
10740        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
10741        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
10742        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
10743        let stream_on = crate::spec::spec_stream()
10744            && !sampled
10745            && !spec_replay
10746            && self.mtp_extra.is_empty()
10747            && constraint.is_none()
10748            && !session_mode
10749            && embd_gpu.is_some()
10750            && !crate::model::full_prec_enabled()
10751            && k + 2 < 96;
10752        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
10753        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
10754        if stream_on {
10755            let cap = e.capture_graph(|e| {
10756                for j in 0..k.max(1) {
10757                    self.mtp_head_forward_cap(
10758                        e,
10759                        mtp,
10760                        &mut dctx.g_tok,
10761                        &mut dctx.g_pos,
10762                        &mut dctx.g_seed,
10763                        &mut dctx.g_p,
10764                        &mut *scratch,
10765                        true,
10766                        true,
10767                        embd_gpu.expect("round stream requires resident embedding"),
10768                        embd_qt,
10769                        embd_rb,
10770                        d_vocab,
10771                        None,
10772                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
10773                        None, // round-stream requires constraint.is_none() (see stream_on)
10774                    )?;
10775                }
10776                Ok(())
10777            });
10778            match cap {
10779                Ok(g) => {
10780                    scratch.set_len(e, 0)?;
10781                    stream_graph = Some(g);
10782                }
10783                Err(err) => {
10784                    scratch.set_len(e, 0)?;
10785                    if debug_spec {
10786                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
10787                    }
10788                }
10789            }
10790        }
10791        let stream_active = stream_on && stream_graph.is_some();
10792        if debug_spec {
10793            eprintln!(
10794                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
10795                crate::spec::spec_stream(),
10796                dctx.graph.is_some(),
10797                stream_graph.is_some()
10798            );
10799        }
10800        let t_v_s = k + 1;
10801        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
10802        // module (extracted 2026-07-12; the gemma burst reuses them).
10803        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
10804        let crate::round_stream::StreamBufs {
10805            mut vtok_d,
10806            mut brk_d,
10807            mut pend_d,
10808            last_pred_d,
10809            mut pos_ctr,
10810            mut pos_start_d,
10811            mut ring_d,
10812            acc_d: mut stream_acc,
10813            m_rounds,
10814            k: _,
10815        } = sb;
10816        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
10817            Some(crate::round_stream::kv_len_ptr_table(
10818                e,
10819                cache,
10820                Some(&pos_ctr),
10821            )?)
10822        } else {
10823            None
10824        };
10825
10826        let t_fill = t_ent.elapsed();
10827        let mut round = 0usize;
10828        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
10829        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
10830        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
10831        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
10832        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
10833        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
10834        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
10835        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
10836        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
10837        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
10838        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
10839        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
10840        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
10841        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
10842        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
10843        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
10844        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
10845        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
10846        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
10847        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
10848        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
10849        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
10850        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
10851        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
10852        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
10853        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
10854        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
10855        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
10856        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
10857        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
10858            .ok()
10859            .and_then(|v| v.parse().ok());
10860        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
10861            4
10862        } else if self.cfg.n_embd as usize >= 2500 {
10863            2
10864        } else {
10865            1
10866        };
10867        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
10868        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
10869        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
10870        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
10871        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
10872            .ok()
10873            .and_then(|v| v.parse().ok())
10874            .unwrap_or(1024);
10875        let floor_at = |pos: usize| -> usize {
10876            if adapt_floor_env.is_some() || pos < floor_ctx {
10877                adapt_floor
10878            } else if adapt_floor >= 4 {
10879                1
10880            } else {
10881                adapt_floor
10882            }
10883        };
10884        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
10885        // fixed-K default path is untouched by this whole block.
10886        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
10887            .ok()
10888            .and_then(|v| v.parse().ok())
10889            .unwrap_or(7);
10890        let k_cap = k.min(cap_max).max(1);
10891        let mut kc = k_cap;
10892        let mut opti_fork: Option<OptiForkState> = None;
10893        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
10894        if fork_mode != OptiForkGateMode::Disabled {
10895            let fence = crate::pp::pp_cuts(self.layers.len());
10896            let refusal = if !session_mode {
10897                Some("not-session")
10898            } else if k != 1 || adapt {
10899                Some("requires-fixed-k1")
10900            } else if sampled || constraint.is_some() || spec_replay {
10901                Some("sampled-constrained-or-replay")
10902            } else if pipe.is_some() {
10903                Some("two-session-pipeline")
10904            } else if !spec_devacc() {
10905                Some("requires-device-accept")
10906            } else if stream_active || crate::spec::spec_stream() {
10907                Some("round-stream")
10908            } else if !self.mtp_extra.is_empty() {
10909                Some("multi-head-mtp")
10910            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
10911                Some("swa-ring")
10912            } else if crate::pp::pp_host_bounce_active() {
10913                Some("host-bounce")
10914            } else if fork_mode == OptiForkGateMode::Controller
10915                && cache.recur.iter().any(Option::is_some)
10916            {
10917                Some("controller-requires-zero-recurrent-state")
10918            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
10919                Some("requires-pp2")
10920            } else {
10921                None
10922            };
10923            if let Some(reason) = refusal {
10924                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10925                eprintln!("[opti-fork] refused reason={reason}");
10926            } else {
10927                let fence = fence.expect("validated PP-2 fence");
10928                let rt = crate::pp::PpNRt::get(e)?;
10929                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
10930                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
10931                let primary_supported =
10932                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
10933                if !rt.cross_device() || !primary_supported {
10934                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10935                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
10936                } else {
10937                    // Both recurrent snapshots and both seed generations are allocated before
10938                    // the first fork, each through its owning PP stage. Allocation failure
10939                    // therefore happens before any optimistic state mutation can occur.
10940                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10941                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
10942                    let fork = OptiForkState::new(
10943                        e,
10944                        cache,
10945                        fork_mode,
10946                        alternate_snapshot,
10947                        &h_seed_buf,
10948                        &fill_prev,
10949                        rt,
10950                        fence[1],
10951                        self.layers.len(),
10952                    )?;
10953                    eprintln!(
10954                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
10955                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
10956                        fence[1],
10957                        fork.logical_payload_bytes[0],
10958                        fork.logical_payload_bytes[1],
10959                        fork.controller.map_or(0.0, |policy| policy.threshold),
10960                    );
10961                    fork_snapshot = Some(current_snapshot);
10962                    opti_fork = Some(fork);
10963                }
10964            }
10965        }
10966        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
10967        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
10968        let mut snap = match fork_snapshot {
10969            Some(snapshot) => snapshot,
10970            None => cache.snapshot(e)?,
10971        };
10972        let mut carried_opti: Option<OptiControllerTicket> = None;
10973        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
10974        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
10975        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
10976            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
10977        } else {
10978            None
10979        };
10980        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
10981        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
10982        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
10983        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
10984        // pass of any kind). Verify still
10985        // checks every emitted token against the target -> exactness holds by construction; only
10986        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
10987        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
10988        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
10989        let mut pending: Option<u32> = carried_pending;
10990        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
10991        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
10992        // the verify accept readback). Printed once at loop end via spec-stats.
10993        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
10994        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
10995        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
10996        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
10997        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
10998        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
10999        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11000        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11001        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11002        let mut ph_wait = 0f64;
11003        let mut ph_commit = 0f64;
11004        let mut ph_t = std::time::Instant::now();
11005        let mut ph_mark = |acc: &mut f64, on: bool| {
11006            if on {
11007                let now = std::time::Instant::now();
11008                *acc += (now - ph_t).as_secs_f64();
11009                ph_t = now;
11010            }
11011        };
11012        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11013        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11014        // arm holds it — the slab stash is live verify -> commit inside a round, and the
11015        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11016        // the model (rebuilding per call re-captures the pool per prompt, which is the
11017        // measured way to lose more than the launches cost); the captured bodies are
11018        // cache-independent, every state read going through per-round refreshed pointer
11019        // tables. None = the eager walk, byte-identical.
11020        //
11021        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11022        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11023        // whenever the stream is live rather than relying on that refusal.
11024        // The lock is taken ONLY when the door is armed: with the flag off this whole block
11025        // is inert, so the default path cannot serialize two spec generations behind a mutex
11026        // it never reads.
11027        let mut vg_guard = if crate::spec::spec_verify_graph_on() && !stream_active {
11028            let mut g = self.dspark_vgraphs.lock().unwrap();
11029            if g.is_none() {
11030                // Size by the WIDEST verify this run can present, which is k+1 and NOT
11031                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11032                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11033                // panic in the sampled ON arm, measured before this line said k+1).
11034                let vt_cap = (k.max(k_cap) + 1).max(2);
11035                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11036                if g.is_some() {
11037                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11038                    // than trusting that a flag set means a pool built.
11039                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11040                } else {
11041                    eprintln!(
11042                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11043                         non-uniform state, or vt_cap < 2) — eager walk"
11044                    );
11045                }
11046            }
11047            Some(g)
11048        } else {
11049            None
11050        };
11051        // Capacity fail-safe: a round wider than the pool was built for must take the eager
11052        // walk, not slice the stash past its rows. The sizing above already covers every
11053        // round this run can present; this keeps a future caller (or a k that grows behind
11054        // the pool's back) on the byte-identical fallback instead of a panic.
11055        let vg_t_cap = vg_guard
11056            .as_ref()
11057            .and_then(|g| g.as_ref())
11058            .map(|g| g.t_capacity())
11059            .unwrap_or(0);
11060        if let Some(p) = pipe {
11061            p.setup_end();
11062        }
11063        while keep_going && out.len() < max_new {
11064            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11065            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11066            if let (true, Some(sg), Some(ptrs)) = (
11067                stream_active && round >= 1 && pending.is_some(),
11068                &stream_graph,
11069                &stream_ptrs,
11070            ) {
11071                if debug_spec {
11072                    static ONCE: std::sync::Once = std::sync::Once::new();
11073                    ONCE.call_once(|| {
11074                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11075                    });
11076                }
11077                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11078                e.set_u32_one(&mut pend_d, pending.unwrap())?;
11079                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11080                for _mi in 0..m_rounds {
11081                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11082                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11083                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11084                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11085                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11086                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11087                    sg.launch()?;
11088                    e.spec_assemble_verify(
11089                        &g_tokp2k,
11090                        &pend_d,
11091                        d2t_dev.as_ref(),
11092                        &mut vtok_d,
11093                        &mut brk_d,
11094                        p_min,
11095                        k,
11096                        pmin0,
11097                    )?;
11098                    let mut ck = VerifyCkpt::new(self.layers.len());
11099                    let dummy = vec![0u32; t_v_s];
11100                    let (tl_d, vx) = self.decode_step_t_core_stream(
11101                        e,
11102                        &dummy,
11103                        0,
11104                        &mut *cache,
11105                        embd_dev,
11106                        Some(&mut ck),
11107                        Some((&vtok_d, &pos_ctr)),
11108                        None,
11109                        None,
11110                        None,
11111                    )?;
11112                    for j in 0..t_v_s {
11113                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11114                    }
11115                    e.spec_accept_greedy_dc(
11116                        &preds_d,
11117                        &vtok_d,
11118                        &last_pred_d,
11119                        &brk_d,
11120                        &mut stream_acc,
11121                    )?;
11122                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
11123                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11124                    self.commit_verified_prefix_stream(
11125                        e,
11126                        &mut *cache,
11127                        &snap,
11128                        &ck,
11129                        &stream_acc,
11130                        1,
11131                        t_v_s,
11132                    )?;
11133                    e.spec_rollback_stream(
11134                        ptrs,
11135                        &pos_start_d,
11136                        &stream_acc,
11137                        1,
11138                        self.layers.len() + 1,
11139                    )?;
11140                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
11141                }
11142                e.stream().synchronize()?;
11143                let ring_h = e.dtoh_u32(&ring_d)?;
11144                let cnt = ring_h[0] as usize;
11145                for i in 0..cnt {
11146                    if out.len() < max_new {
11147                        out.push(ring_h[1 + i]);
11148                    }
11149                }
11150                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
11151                for il in 0..self.layers.len() {
11152                    if let Some(kvl) = cache.kv[il].as_mut() {
11153                        kvl.len = pos_h;
11154                    }
11155                }
11156                cache.pos = pos_h;
11157                scratch.kv.len = pos_h;
11158                pending = Some(ring_h[cnt]); // last drained token = the live bonus
11159                last_token = ring_h[cnt];
11160                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
11161                total_accepted += cnt.saturating_sub(m_rounds);
11162                if let Some(t) = sess_telem {
11163                    // totals only — the burst's per-round accept counts stayed on device
11164                    // (that is the point of the round-stream arm). pos_* untouched.
11165                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
11166                }
11167                round += m_rounds;
11168                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
11169                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11170                continue;
11171            }
11172            let pipe_draft = match pipe {
11173                Some(p) => Some(p.draft_begin(round)?),
11174                None => None,
11175            };
11176            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
11177            let mut current_opti = carried_opti.take();
11178            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
11179                match opti_fork.as_mut() {
11180                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
11181                    None => None,
11182                    Some(_) => None,
11183                }
11184            } else {
11185                None
11186            };
11187            if current_opti.is_none() {
11188                if let Some(fork) = opti_fork.as_ref() {
11189                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
11190                } else {
11191                    cache.snapshot_into(e, &mut snap)?;
11192                }
11193            } else if snap.pos != pos {
11194                return Err(format!(
11195                    "optipipe carried snapshot pos {} != current pos {pos}",
11196                    snap.pos
11197                )
11198                .into());
11199            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
11200            ph_mark(&mut ph_rest, phase_on);
11201
11202            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
11203            // p-min semantics (both paths): stop the chain early when the head's confidence in
11204            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
11205            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
11206            let base0 = if pending.is_some() { 1usize } else { 0usize };
11207            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
11208            // accepted run + 1 (the gemma law — see the setup block above the loop).
11209            let k_this = if adapt { kc } else { k };
11210            let mut draft: Vec<u32> = Vec::with_capacity(k);
11211            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
11212            let mut controller_draft_prob: Option<f32> = None;
11213            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
11214            if let Some(ticket) = current_opti.as_mut() {
11215                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
11216                if ticket.verify_tokens[0] != carried_pending {
11217                    return Err(format!(
11218                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
11219                        ticket.verify_tokens[0],
11220                    )
11221                    .into());
11222                }
11223                draft.push(ticket.verify_tokens[1]);
11224                controller_draft_prob = Some(ticket.draft_prob);
11225                controller_eager_state = ticket
11226                    .take_eager_seed()
11227                    .map(|seed| (ticket.verify_tokens[1], seed));
11228            } else {
11229                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
11230                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
11231                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
11232                // rejected drafts and p-min extras via the len mechanism).
11233                scratch.set_len(e, pos + base0 - 1)?;
11234                if pen_on {
11235                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
11236                    // device dedup: `penalty_last_n` is usize::MAX for any serve request with
11237                    // a penalty, so without the cap this grew with the whole session.
11238                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
11239                    let w0 = pen_hist.len().saturating_sub(win);
11240                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
11241                }
11242                if sampled {
11243                    draft_logits.clear();
11244                    draft_stats.clear();
11245                }
11246                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
11247                // position's mask is computed on that clone and advanced by the PROPOSED token. The
11248                // real state moves only on emission (verify's job), so the emitted stream is
11249                // unchanged — the mask only removes tokens the verify would have truncated anyway.
11250                let mut dmask_live = dmask_on;
11251                if dmask_live {
11252                    let t_c = std::time::Instant::now();
11253                    constraint
11254                        .as_deref_mut()
11255                        .unwrap()
11256                        .draft_begin()
11257                        .map_err(|e2| format!("constraint: {e2}"))?;
11258                    dm_clone_ns += t_c.elapsed().as_nanos();
11259                    dm_rounds += 1;
11260                }
11261                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
11262                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
11263                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
11264                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
11265                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11266                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
11267                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11268                    for j in 0..k_this {
11269                        // per-position mask upload (contents only — the graph's baked pointer is
11270                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
11271                        // mask node degrades to a no-op ban instead of needing a second graph.
11272                        if dmask_live
11273                            && !upload_draft_mask(
11274                                e,
11275                                constraint.as_deref_mut().unwrap(),
11276                                &mut dctx.g_dmask,
11277                                mtp.d2t.as_ref(),
11278                                d_vocab,
11279                                dmask_words,
11280                            )?
11281                        {
11282                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
11283                            // genuinely miss the legal set): neutralize the captured mask node and
11284                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
11285                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11286                            dmask_live = false;
11287                        }
11288                        gr.launch()?;
11289                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11290                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11291                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
11292                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
11293                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
11294                        // replay's embed node, and the MMU fault kills the CUDA context for the
11295                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
11296                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
11297                        // buffer (g_seed = the verify-side handoff vs head-side compute).
11298                        if (idx as usize) >= d_vocab {
11299                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
11300                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
11301                            // seed, untouched since the round-start copy — the pair discriminates
11302                            // "seed arrived poisoned" from "head forward produced NaN".
11303                            let seed_h = e.dtoh(&dctx.g_seed)?;
11304                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11305                            let in_h = e.dtoh(&h_seed_buf)?;
11306                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
11307                            return Err(format!(
11308                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11309                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
11310                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
11311                             the embed row (#87 trap)"
11312                            )
11313                            .into());
11314                        }
11315                        // trimmed draft vocab -> target token id (identity when no d2t map)
11316                        let d = match &mtp.d2t {
11317                            Some(map) => map[idx as usize],
11318                            None => idx,
11319                        };
11320                        let draft_p = if p_min > 0.0
11321                            || opti_fork
11322                                .as_ref()
11323                                .is_some_and(|fork| fork.controller.is_some())
11324                        {
11325                            Some(e.dtoh(&dctx.g_p)?[0])
11326                        } else {
11327                            None
11328                        };
11329                        if j == 0 {
11330                            controller_draft_prob = draft_p;
11331                        }
11332                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11333                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11334                                break;
11335                            }
11336                        }
11337                        draft.push(d);
11338                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
11339                        // index the argmax wrote — patch the persistent token buffer (4B htod).
11340                        if d != idx {
11341                            e.set_u32_one(&mut dctx.g_tok, d)?;
11342                        }
11343                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
11344                        // unmasked drafting for the remaining positions (verify still arbitrates).
11345                        // speculative advance; a chain the grammar can no longer follow (EOS
11346                        // proposed) ends here. The captured mask node always runs, so a dead chain
11347                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
11348                        if dmask_live
11349                            && !constraint
11350                                .as_deref_mut()
11351                                .unwrap()
11352                                .draft_advance(d)
11353                                .map_err(|e2| format!("constraint: {e2}"))?
11354                        {
11355                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11356                            break;
11357                        }
11358                    }
11359                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
11360                // legal ONLY in the regime it was captured in. The condition used to read
11361                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
11362                // which it could not, because the key omitted the filters. Both halves are now
11363                // enforced: the key drops a stale graph, and this site refuses to launch one.
11364                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
11365                    if skey_probe() {
11366                        eprintln!(
11367                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
11368                             top_p={} min_p={} s_key_parked={:?}",
11369                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11370                        );
11371                    }
11372                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
11373                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
11374                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
11375                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
11376                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
11377                    // stream. Host sctr advances in lockstep (computed, no readback needed).
11378                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11379                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
11380                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11381                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
11382                    for j in 0..k_this {
11383                        gr.launch()?;
11384                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11385                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
11386                        // counts the p-min-discarded token too)
11387                        // q retention: ONE async D2D of the persistent head-logits buffer into this
11388                        // round's slot j (stream-ordered after the replay, before the next one).
11389                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
11390                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11391                        // #87 SENTINEL TRAP (see the greedy graph arm above).
11392                        if (idx as usize) >= d_vocab {
11393                            let seed_h = e.dtoh(&dctx.g_seed)?;
11394                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11395                            return Err(format!(
11396                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
11397                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
11398                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
11399                             (#87 trap)"
11400                            )
11401                            .into());
11402                        }
11403                        let d = match &mtp.d2t {
11404                            Some(map) => map[idx as usize],
11405                            None => idx,
11406                        };
11407                        draft_idx.push(idx);
11408                        if p_min > 0.0 {
11409                            let p = e.dtoh(&dctx.g_p)?[0];
11410                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11411                                break;
11412                            }
11413                        }
11414                        draft.push(d);
11415                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
11416                        if d != idx {
11417                            e.set_u32_one(&mut dctx.g_tok, d)?;
11418                        }
11419                    }
11420                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
11421                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
11422                    for j in 0..draft.len().max(draft_idx.len()) {
11423                        let rows0 = e.htod_i32(&[0])?;
11424                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11425                        e.filter_stats(
11426                            &dctx.q_slots[j],
11427                            d_vocab,
11428                            &rows0,
11429                            &mut th_d,
11430                            &mut z_d,
11431                            &mut mx_d,
11432                            d_vocab,
11433                            1,
11434                            sp_temp,
11435                            sp.top_k,
11436                            sp.top_p,
11437                            sp.min_p,
11438                        )?;
11439                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
11440                    }
11441                } else {
11442                    if skey_probe() && sampled {
11443                        eprintln!(
11444                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
11445                             top_p={} min_p={} s_key_parked={:?}",
11446                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11447                        );
11448                    }
11449                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
11450                    let chain_heads = !self.mtp_extra.is_empty();
11451                    let mut e_tok = last_token;
11452                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
11453                    let mut chain_tokens = if chain_heads {
11454                        vec![last_token]
11455                    } else {
11456                        Vec::new()
11457                    };
11458                    let mut chain_seeds = if chain_heads {
11459                        vec![e.clone_dtod(&h_seed_buf)?]
11460                    } else {
11461                        Vec::new()
11462                    };
11463                    for j in 0..k_this {
11464                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
11465                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
11466                        let mtp_pos = pos + base0 + j;
11467                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
11468                        // A position with no legal draft-vocab row drops to unmasked drafting for
11469                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
11470                        if dmask_live {
11471                            dmask_live = upload_draft_mask(
11472                                e,
11473                                constraint.as_deref_mut().unwrap(),
11474                                &mut dctx.g_dmask,
11475                                mtp.d2t.as_ref(),
11476                                d_vocab,
11477                                dmask_words,
11478                            )?;
11479                        }
11480                        let mask = if dmask_live {
11481                            Some((&dctx.g_dmask, dmask_words))
11482                        } else {
11483                            None
11484                        };
11485                        let (dl_d, h_nextn) = if chain_heads {
11486                            if debug_spec {
11487                                eprintln!(
11488                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
11489                                    mtp_chain_head_index(j, self.mtp_head_count()),
11490                                    chain_tokens.len(),
11491                                );
11492                            }
11493                            self.mtp_chain_forward_dev(
11494                                e,
11495                                &chain_tokens,
11496                                &chain_seeds,
11497                                &mut *scratch,
11498                                pos + base0 - 1,
11499                                embd_dev,
11500                                mask,
11501                            )?
11502                        } else {
11503                            self.mtp_head_forward_dev(
11504                                e,
11505                                mtp,
11506                                e_tok,
11507                                &d_seed,
11508                                &mut *scratch,
11509                                mtp_pos,
11510                                embd_dev,
11511                                mask,
11512                            )?
11513                        };
11514                        let tok_d = if sampled {
11515                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
11516                            // the filtered softmax (filters off => th=0, exact v1 semantics).
11517                            if perturb_buf.is_none() {
11518                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
11519                            }
11520                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
11521                            if pen_on {
11522                                let h = pen_hist_d.as_ref().unwrap();
11523                                let nh = h.len();
11524                                e.penalize_logits(
11525                                    &mut q_row,
11526                                    h,
11527                                    nh,
11528                                    sp.penalty_repeat,
11529                                    sp.penalty_freq,
11530                                    sp.penalty_present,
11531                                    d_vocab,
11532                                )?;
11533                            }
11534                            let rows0 = e.htod_i32(&[0])?;
11535                            let (mut th_d, mut z_d, mut mx_d) =
11536                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11537                            e.filter_stats(
11538                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
11539                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
11540                            )?;
11541                            let (th, z, mx) =
11542                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
11543                            let pb = perturb_buf.as_mut().unwrap();
11544                            e.gumbel_perturb_filtered(
11545                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
11546                            )?;
11547                            sctr += 1;
11548                            draft_logits.push(q_row);
11549                            draft_stats.push((mx, th, z));
11550                            e.argmax_token_device(pb, d_vocab)?
11551                        } else {
11552                            e.argmax_token_device(&dl_d, d_vocab)?
11553                        };
11554                        let idx = e.dtoh_u32_one(&tok_d)?;
11555                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
11556                        // here because the eager chain's operands are all readable: dl_d (the head
11557                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
11558                        if (idx as usize) >= d_vocab {
11559                            let dl_h = e.dtoh(&dl_d)?;
11560                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
11561                            let seed_h = if chain_heads {
11562                                e.dtoh(chain_seeds.last().unwrap())?
11563                            } else {
11564                                e.dtoh(&d_seed)?
11565                            };
11566                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11567                            return Err(format!(
11568                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11569                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
11570                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
11571                             embed row (#87 trap)"
11572                            )
11573                            .into());
11574                        }
11575                        let d = match &mtp.d2t {
11576                            Some(map) => map[idx as usize],
11577                            None => idx,
11578                        };
11579                        if sampled {
11580                            draft_idx.push(idx);
11581                        }
11582                        let draft_p = if p_min > 0.0
11583                            || opti_fork
11584                                .as_ref()
11585                                .is_some_and(|fork| fork.controller.is_some())
11586                        {
11587                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
11588                            Some(e.dtoh(&p_d)?[0])
11589                        } else {
11590                            None
11591                        };
11592                        if j == 0 {
11593                            controller_draft_prob = draft_p;
11594                        }
11595                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11596                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11597                                break;
11598                            }
11599                        }
11600                        draft.push(d);
11601                        if chain_heads {
11602                            chain_tokens.push(d);
11603                            chain_seeds.push(h_nextn);
11604                        } else {
11605                            e_tok = d;
11606                            d_seed = h_nextn;
11607                        }
11608                        // speculative advance; a chain the grammar can no longer follow (EOS
11609                        // proposed) ends here — the prefix already proposed still rides verify.
11610                        if dmask_live
11611                            && !constraint
11612                                .as_deref_mut()
11613                                .unwrap()
11614                                .draft_advance(d)
11615                                .map_err(|e2| format!("constraint: {e2}"))?
11616                        {
11617                            break;
11618                        }
11619                    }
11620                    if !chain_heads
11621                        && opti_fork
11622                            .as_ref()
11623                            .is_some_and(|fork| fork.controller.is_some())
11624                    {
11625                        controller_eager_state = Some((e_tok, d_seed));
11626                    }
11627                }
11628            }
11629            let k_round = draft.len();
11630            if let Some(p) = pipe {
11631                p.draft_end(round);
11632            }
11633            drop(pipe_draft);
11634
11635            ph_mark(&mut ph_draft, phase_on);
11636            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
11637            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
11638            let verify_tokens: Vec<u32> = match pending {
11639                Some(b) => {
11640                    let mut v = Vec::with_capacity(k_round + 1);
11641                    v.push(b);
11642                    v.extend_from_slice(&draft);
11643                    v
11644                }
11645                None => draft.clone(),
11646            };
11647            let base = if pending.is_some() { 1 } else { 0 };
11648            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
11649            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
11650            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
11651                Some(ticket.take_ckpt())
11652            } else if spec_replay {
11653                None
11654            } else {
11655                Some(VerifyCkpt::new(self.layers.len()))
11656            };
11657            let controller_can_probe = base == 1
11658                && k_round == 1
11659                && out.len().saturating_add(2) < max_new
11660                && controller_draft_prob.is_some()
11661                && opti_fork
11662                    .as_ref()
11663                    .and_then(|fork| fork.controller.as_ref())
11664                    .is_some_and(|policy| !policy.breaker_tripped);
11665            let mut successor_attempt: Option<OptiControllerTicket> = None;
11666            let mut rejected_probe: Option<(f32, u32)> = None;
11667            let mut controller_prepared: Option<OptiControllerPrepared> = None;
11668            if controller_can_probe {
11669                // Prepare d2/q and, on admission, d3 before either current verify half is
11670                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
11671                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
11672                // the primary stream after N stage 1 would serialize the supposed pipeline.
11673                let eager_pos = scratch.kv.len + 1;
11674                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
11675                    e,
11676                    mtp,
11677                    &mut dctx,
11678                    &mut *scratch,
11679                    d_vocab,
11680                    &mut controller_eager_state,
11681                    eager_pos,
11682                    embd_dev,
11683                )?;
11684                let first_probability = controller_draft_prob
11685                    .ok_or("optipipe controller probe lost first-token probability")?;
11686                let q_proxy = first_probability * pending_probability;
11687                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11688                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11689                let admitted = opti_fork
11690                    .as_ref()
11691                    .and_then(|fork| fork.controller.as_ref())
11692                    .ok_or("optipipe controller policy disappeared")?
11693                    .admit(q_proxy);
11694                if admitted {
11695                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11696                    let eager_pos = scratch.kv.len + 1;
11697                    let (optimistic_draft, optimistic_draft_probability) = self
11698                        .opti_controller_draft_step(
11699                            e,
11700                            mtp,
11701                            &mut dctx,
11702                            &mut *scratch,
11703                            d_vocab,
11704                            &mut controller_eager_state,
11705                            eager_pos,
11706                            embd_dev,
11707                        )?;
11708                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11709                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
11710                        debug_assert_eq!(token, optimistic_draft);
11711                        seed
11712                    });
11713                    controller_prepared = Some(OptiControllerPrepared {
11714                        verify_tokens: [optimistic_pending, optimistic_draft],
11715                        draft_prob: optimistic_draft_probability,
11716                        eager_seed,
11717                        q_proxy,
11718                        scratch_len: scratch.kv.len,
11719                    });
11720                } else {
11721                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11722                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11723                    rejected_probe = Some((q_proxy, optimistic_pending));
11724                    eprintln!(
11725                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
11726                        opti_fork
11727                            .as_ref()
11728                            .and_then(|fork| fork.controller.as_ref())
11729                            .expect("controller policy")
11730                            .threshold,
11731                    );
11732                }
11733            }
11734            let fork_attempt = match fork_generation.take() {
11735                Some(generation) if base == 1 && k_round == 1 => Some(generation),
11736                Some(generation) => {
11737                    opti_fork
11738                        .as_mut()
11739                        .expect("fork generation without fork state")
11740                        .retire(generation)?;
11741                    None
11742                }
11743                None => None,
11744            };
11745            let (tlogits_d, vx) = if let Some(p) = pipe {
11746                self.decode_step_t_core_pipelined(
11747                    e,
11748                    &verify_tokens,
11749                    pos,
11750                    &mut *cache,
11751                    embd_dev,
11752                    ckpt.as_mut(),
11753                    p,
11754                    round,
11755                )?
11756            } else if controller_can_probe {
11757                let fence = opti_fork
11758                    .as_ref()
11759                    .ok_or("optipipe controller probe lost fork state")?
11760                    .fence;
11761                let boundary = match current_opti.as_mut() {
11762                    Some(ticket) => ticket.take_boundary(),
11763                    None => self.verify_stage0_issue(
11764                        e,
11765                        &verify_tokens,
11766                        pos,
11767                        &mut *cache,
11768                        embd_dev,
11769                        ckpt.as_mut(),
11770                        None,
11771                        &fence,
11772                        Some(true),
11773                        None,
11774                    )?,
11775                };
11776                if let Some(prepared) = controller_prepared.take() {
11777                    let generation = {
11778                        let fork = opti_fork
11779                            .as_mut()
11780                            .ok_or("optipipe controller admission lost fork state")?;
11781                        let generation = fork.reserve_successor()?;
11782                        let rt = fork.rt;
11783                        let snapshot_fence = fork.fence;
11784                        opti_snapshot_one_stage_owned_into(
11785                            e,
11786                            cache,
11787                            rt,
11788                            &snapshot_fence,
11789                            0,
11790                            fork.successor_snapshot_mut(),
11791                        )?;
11792                        generation
11793                    };
11794                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
11795                    let successor_boundary = self.verify_stage0_issue(
11796                        e,
11797                        &prepared.verify_tokens,
11798                        pos + verify_tokens.len(),
11799                        &mut *cache,
11800                        embd_dev,
11801                        Some(&mut successor_ckpt),
11802                        None,
11803                        &fence,
11804                        Some(false),
11805                        None,
11806                    )?;
11807                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11808                    let fork = opti_fork
11809                        .as_ref()
11810                        .ok_or("optipipe controller ticket lost fork state")?;
11811                    successor_attempt = Some(fork.controller_ticket(
11812                        generation,
11813                        successor_boundary,
11814                        successor_ckpt,
11815                        prepared.verify_tokens,
11816                        prepared.draft_prob,
11817                        prepared.eager_seed,
11818                        prepared.q_proxy,
11819                        prepared.scratch_len,
11820                    ));
11821                    eprintln!(
11822                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
11823                         verify={:?}",
11824                        generation.id,
11825                        prepared.q_proxy,
11826                        fork.controller.expect("controller policy").threshold,
11827                        prepared.verify_tokens,
11828                    );
11829                }
11830                let result = self.verify_stage1_finish(
11831                    e,
11832                    boundary,
11833                    &mut *cache,
11834                    ckpt.as_mut(),
11835                    None,
11836                    &fence,
11837                    successor_attempt.is_none(),
11838                )?;
11839                if let Some(ticket) = current_opti.as_mut() {
11840                    ticket.settle();
11841                }
11842                if successor_attempt.is_some() {
11843                    let fork = opti_fork
11844                        .as_mut()
11845                        .ok_or("optipipe successor snapshot lost fork state")?;
11846                    let rt = fork.rt;
11847                    let snapshot_fence = fork.fence;
11848                    opti_snapshot_one_stage_owned_into(
11849                        e,
11850                        cache,
11851                        rt,
11852                        &snapshot_fence,
11853                        1,
11854                        fork.successor_snapshot_mut(),
11855                    )?;
11856                    // Publish N only after both independent successor-state queues are complete.
11857                    fork.rt.publish_to(1, &e.stream())?;
11858                }
11859                result
11860            } else if let Some(ticket) = current_opti.as_mut() {
11861                let fork = opti_fork
11862                    .as_mut()
11863                    .ok_or("optipipe carried controller ticket lost fork state")?;
11864                let boundary = ticket.take_boundary();
11865                let result = self.verify_stage1_finish(
11866                    e,
11867                    boundary,
11868                    &mut *cache,
11869                    ckpt.as_mut(),
11870                    None,
11871                    &fork.fence,
11872                    true,
11873                )?;
11874                ticket.settle();
11875                result
11876            } else if let Some(generation) = fork_attempt {
11877                let fork = opti_fork
11878                    .as_mut()
11879                    .expect("fork generation without fork state");
11880                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
11881                let action = fork.mode.action(generation.id);
11882                let boundary = self.verify_stage0_issue(
11883                    e,
11884                    &verify_tokens,
11885                    pos,
11886                    &mut *cache,
11887                    embd_dev,
11888                    ckpt.as_mut(),
11889                    None,
11890                    &fork.fence,
11891                    Some(true),
11892                    None,
11893                )?;
11894                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11895                let mut ticket = fork.ticket(generation, boundary);
11896                if action == OptiForkAction::Abort {
11897                    return Err(format!(
11898                        "optipipe forced abort with generation {} stage0 in flight",
11899                        generation.id,
11900                    )
11901                    .into());
11902                }
11903                fork.reconcile(
11904                    e,
11905                    &mut *cache,
11906                    &mut *scratch,
11907                    &snap,
11908                    &mut h_seed_buf,
11909                    &mut fill_prev,
11910                    generation,
11911                    action,
11912                    verify_tokens[0],
11913                )?;
11914                let result = if action == OptiForkAction::Hit {
11915                    let boundary = ticket.take_boundary();
11916                    self.verify_stage1_finish(
11917                        e,
11918                        boundary,
11919                        &mut *cache,
11920                        ckpt.as_mut(),
11921                        None,
11922                        &fork.fence,
11923                        true,
11924                    )?
11925                } else {
11926                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
11927                    // verify only after E_restart published the restored stage-0 state.
11928                    self.decode_step_t_core(
11929                        e,
11930                        &verify_tokens,
11931                        pos,
11932                        &mut *cache,
11933                        embd_dev,
11934                        ckpt.as_mut(),
11935                    )?
11936                };
11937                ticket.settle();
11938                debug_assert_eq!(ticket.generation, generation);
11939                fork.retire(generation)?;
11940                result
11941            } else {
11942                // The serial verify every non-fork round takes — the MTP route's
11943                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
11944                // a pool above, and then the walk replays the captured trunk instead of
11945                // re-issuing it launch by launch.
11946                let vg_round = if verify_tokens.len() <= vg_t_cap {
11947                    vg_guard.as_mut().and_then(|g| g.as_mut())
11948                } else {
11949                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
11950                        // The commit reads this flag to pick its arm; a round that declines
11951                        // the pool must not inherit a stale `true` from the round before it.
11952                        g.round_slab = false;
11953                    }
11954                    None
11955                };
11956                self.decode_step_t_core_vg(
11957                    e,
11958                    &verify_tokens,
11959                    pos,
11960                    &mut *cache,
11961                    embd_dev,
11962                    ckpt.as_mut(),
11963                    vg_round,
11964                )?
11965            };
11966            let pipe_accept = match pipe {
11967                Some(p) => Some(p.accept_begin(round)?),
11968                None => None,
11969            };
11970
11971            ph_mark(&mut ph_verify, phase_on);
11972            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
11973            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
11974            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
11975            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
11976            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
11977            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
11978            // (== the bonus), so every index shifts by `base` and last_pred is unused.
11979            let t_v = verify_tokens.len();
11980            let mut preds: Vec<u32> = Vec::new();
11981            if !sampled {
11982                for j in 0..t_v {
11983                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
11984                }
11985                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
11986                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
11987                // next round's last_token = the next chain's embed lookup. Catch it at the
11988                // source with the column named — an all-NaN VERIFY column implicates the
11989                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
11990                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
11991                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
11992                    let mut probe = e.zeros(n_vocab)?;
11993                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
11994                    let col_h = e.dtoh(&probe)?;
11995                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
11996                    return Err(format!(
11997                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
11998                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
11999                         — the stage-split verify produced a poisoned column (#87 trap)",
12000                        preds[bad]
12001                    )
12002                    .into());
12003                }
12004            }
12005            ph_mark(&mut ph_wait, phase_on);
12006            let t_pred = |j: usize| -> u32 {
12007                if j == 0 && base == 0 {
12008                    last_pred
12009                } else {
12010                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12011                    // used to call this from the sampled arm and panicked the worker; it now goes
12012                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12013                    // out-of-range pred is a real bug, not something to paper over.
12014                    debug_assert!(
12015                        !sampled,
12016                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
12017                    );
12018                    preds[base + j - 1]
12019                }
12020            };
12021            let mut devacc_seeded = false;
12022            let mut devacc_acc: Option<CudaSlice<u32>> = None;
12023            let (n_acc, bonus) = if !sampled {
12024                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
12025                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
12026                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
12027                // gated on token identity vs the host walk (the arms below are bit-equal rules).
12028                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
12029                {
12030                    let draft_d = e.htod_u32_v(&draft)?;
12031                    let mut acc_out = e.alloc_u32_zeroed(2)?;
12032                    e.spec_accept_greedy(
12033                        &preds_d,
12034                        &draft_d,
12035                        last_pred,
12036                        base,
12037                        k_round,
12038                        &mut acc_out,
12039                    )?;
12040                    devacc_acc = Some(acc_out.clone());
12041                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
12042                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
12043                    // non-replay commit arms skip their host-offset seed copies (guarded below);
12044                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
12045                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
12046                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
12047                    // the update lands after the arms (devacc_seeded guard below).
12048                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
12049                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
12050                    // unified rule; full accept rewrites the verify-left value). Host mirrors
12051                    // update after the readback; commit_verified_prefix skips its len_d writes.
12052                    if let Some(successor) = successor_attempt.as_ref() {
12053                        opti_fork
12054                            .as_mut()
12055                            .ok_or("optipipe successor reconcile lost fork state")?
12056                            .queue_actual_reconcile(
12057                                e,
12058                                &snap,
12059                                &acc_out,
12060                                successor.verify_tokens[0],
12061                                base,
12062                            )?;
12063                    } else if let Some(ptrs) = &kv_len_ptrs {
12064                        let saved: Vec<i32> = (0..self.layers.len())
12065                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
12066                            .collect();
12067                        let saved_d = e.htod_i32(&saved)?;
12068                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
12069                    }
12070                    devacc_seeded = true;
12071                    let ab = e.dtoh_u32(&acc_out)?;
12072                    (ab[0] as usize, ab[1])
12073                } else {
12074                    let mut n_acc = 0usize;
12075                    for j in 0..k_round {
12076                        if t_pred(j) == draft[j] {
12077                            n_acc += 1;
12078                        } else {
12079                            break;
12080                        }
12081                    }
12082                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12083                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12084                    (n_acc, t_pred(n_acc))
12085                }
12086            } else {
12087                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12088                if col_buf.is_none() {
12089                    col_buf = Some(e.zeros(n_vocab)?);
12090                }
12091                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12092                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12093                let mut pj = vec![0f32; k_round.max(1)];
12094                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12095                if k_round > 0 {
12096                    let mut ids: Vec<u32> = Vec::new();
12097                    let mut rows: Vec<i32> = Vec::new();
12098                    for j in 0..k_round {
12099                        if j > 0 || base == 1 {
12100                            ids.push(draft[j]);
12101                            rows.push((base + j) as i32 - 1);
12102                        }
12103                    }
12104                    if !ids.is_empty() {
12105                        let nr = rows.len();
12106                        // penalties: materialize the used columns into one contiguous penalized
12107                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12108                        // penalties: materialize used columns contiguously, penalize all rows in
12109                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12110                        let p_rows: Vec<i32> = if pen_on {
12111                            (0..nr as i32).collect()
12112                        } else {
12113                            rows.clone()
12114                        };
12115                        if pen_on {
12116                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12117                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
12118                            }
12119                            let pc = pcol_buf.as_mut().unwrap();
12120                            for (i2, &r) in rows.iter().enumerate() {
12121                                let c = r as usize;
12122                                e.copy_view_into(
12123                                    pc,
12124                                    i2 * n_vocab,
12125                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12126                                    n_vocab,
12127                                )?;
12128                            }
12129                            let h = pen_hist_d.as_ref().unwrap();
12130                            let nh = h.len();
12131                            e.penalize_logits_rows(
12132                                pc,
12133                                h,
12134                                nh,
12135                                sp.penalty_repeat,
12136                                sp.penalty_freq,
12137                                sp.penalty_present,
12138                                n_vocab,
12139                                nr,
12140                            )?;
12141                        }
12142                        let p_src: &CudaSlice<f32> = if pen_on {
12143                            pcol_buf.as_ref().unwrap()
12144                        } else {
12145                            &tlogits_d
12146                        };
12147                        let rowsd = e.htod_i32(&p_rows)?;
12148                        let (mut th_d, mut z_d, mut mx_d) =
12149                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12150                        e.filter_stats(
12151                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12152                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12153                        )?;
12154                        let idsd = e.htod_u32_v(&ids)?;
12155                        let mut outd = e.zeros(nr)?;
12156                        e.softmax_gather_filtered(
12157                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12158                            sp_temp,
12159                        )?;
12160                        let outv = e.dtoh(&outd)?;
12161                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12162                        let mut oi = 0usize;
12163                        for j in 0..k_round {
12164                            if j > 0 || base == 1 {
12165                                pj[j] = outv[oi];
12166                                oi += 1;
12167                            }
12168                        }
12169                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12170                    }
12171                    if base == 0 {
12172                        let lc: &CudaSlice<f32> = if pen_on {
12173                            if col_buf.is_none() {
12174                                col_buf = Some(e.zeros(n_vocab)?);
12175                            }
12176                            let cb = col_buf.as_mut().unwrap();
12177                            e.copy_into(
12178                                cb,
12179                                0,
12180                                last_col_logits
12181                                    .as_ref()
12182                                    .expect("sampled: last_col_logits unset"),
12183                                n_vocab,
12184                            )?;
12185                            let h = pen_hist_d.as_ref().unwrap();
12186                            let nh = h.len();
12187                            e.penalize_logits(
12188                                cb,
12189                                h,
12190                                nh,
12191                                sp.penalty_repeat,
12192                                sp.penalty_freq,
12193                                sp.penalty_present,
12194                                n_vocab,
12195                            )?;
12196                            col_buf.as_ref().unwrap()
12197                        } else {
12198                            last_col_logits
12199                                .as_ref()
12200                                .expect("sampled: last_col_logits unset")
12201                        };
12202                        let rows0 = e.htod_i32(&[0])?;
12203                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12204                        e.filter_stats(
12205                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12206                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12207                        )?;
12208                        let idsd = e.htod_u32_v(&[draft[0]])?;
12209                        let mut outd = e.zeros(1)?;
12210                        e.softmax_gather_filtered(
12211                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12212                        )?;
12213                        pj[0] = e.dtoh(&outd)?[0];
12214                        last_col_stats =
12215                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12216                    }
12217                }
12218                // q source: the graph arm retained the head logits in the persistent q_slots;
12219                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
12220                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
12221                // computes them post-replay — graph engages only filter/penalty-free, so the
12222                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
12223                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
12224                    &dctx.q_slots
12225                } else {
12226                    &draft_logits
12227                };
12228                let mut n_acc = 0usize;
12229                for j in 0..k_round {
12230                    let (qmx, qth, qz) = draft_stats[j];
12231                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
12232                    let rowsd = e.htod_i32(&[0])?;
12233                    let thd = e.htod(&[qth])?;
12234                    let zd = e.htod(&[qz])?;
12235                    let _ = qmx;
12236                    let mut outd = e.zeros(1)?;
12237                    e.softmax_gather_filtered(
12238                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
12239                        sp_temp,
12240                    )?;
12241                    let qj = e.dtoh(&outd)?[0];
12242                    let u = host_u01(sp_seed, uctr);
12243                    uctr += 1;
12244                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
12245                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
12246                    // exactness signature (see `skey_probe`). Impossible when the draft was
12247                    // drawn from the same filtered distribution the verify reconstructs here;
12248                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
12249                    if skey_probe() && qj == 0.0 {
12250                        eprintln!(
12251                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
12252                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
12253                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
12254                        );
12255                    }
12256                    if accept {
12257                        n_acc += 1;
12258                    } else {
12259                        break;
12260                    }
12261                }
12262                let bonus = if n_acc == k_round {
12263                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
12264                    let col = base + k_round - 1;
12265                    let cb = col_buf.as_mut().unwrap();
12266                    e.copy_view_into(
12267                        cb,
12268                        0,
12269                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12270                        n_vocab,
12271                    )?;
12272                    if pen_on {
12273                        let h = pen_hist_d.as_ref().unwrap();
12274                        let nh = h.len();
12275                        e.penalize_logits(
12276                            cb,
12277                            h,
12278                            nh,
12279                            sp.penalty_repeat,
12280                            sp.penalty_freq,
12281                            sp.penalty_present,
12282                            n_vocab,
12283                        )?;
12284                    }
12285                    if perturb_buf.is_none() {
12286                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12287                    }
12288                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
12289                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
12290                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
12291                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
12292                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
12293                    // last gathered column, in both base arms. `th` is a threshold in e-units of
12294                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
12295                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
12296                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
12297                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
12298                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
12299                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
12300                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
12301                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
12302                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
12303                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
12304                    // and row_max is unused once nothing is masked), so this fix is a byte-level
12305                    // no-op for the untruncated serve default. One extra one-block filter_stats
12306                    // per full-accept round is the whole cost.
12307                    let (mx, th) = {
12308                        let rows0 = e.htod_i32(&[0])?;
12309                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12310                        let cb0 = col_buf.as_ref().unwrap();
12311                        e.filter_stats(
12312                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12313                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12314                        )?;
12315                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
12316                    };
12317                    let pb = perturb_buf.as_mut().unwrap();
12318                    let cb2 = col_buf.as_ref().unwrap();
12319                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
12320                    sctr += 1;
12321                    let td = e.argmax_token_device(pb, n_vocab)?;
12322                    e.dtoh_u32_one(&td)?
12323                } else {
12324                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
12325                    let cb = col_buf.as_mut().unwrap();
12326                    if n_acc > 0 || base == 1 {
12327                        let col = base + n_acc - 1;
12328                        e.copy_view_into(
12329                            cb,
12330                            0,
12331                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12332                            n_vocab,
12333                        )?;
12334                    } else {
12335                        let lc = last_col_logits.as_ref().unwrap();
12336                        e.copy_into(cb, 0, lc, n_vocab)?;
12337                    }
12338                    if pen_on {
12339                        let h = pen_hist_d.as_ref().unwrap();
12340                        let nh = h.len();
12341                        e.penalize_logits(
12342                            cb,
12343                            h,
12344                            nh,
12345                            sp.penalty_repeat,
12346                            sp.penalty_freq,
12347                            sp.penalty_present,
12348                            n_vocab,
12349                        )?;
12350                    }
12351                    let cb2 = col_buf.as_ref().unwrap();
12352                    let sc = sctr;
12353                    sctr += 1;
12354                    // p-stats for the reject column: from col_stats when the col was gathered,
12355                    // else (j==0&&base==0) from last_col_stats.
12356                    let p_stats = if n_acc > 0 || base == 1 {
12357                        // col index within the gathered set == number of gathered cols before n_acc
12358                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
12359                        col_stats.get(gi).copied().unwrap_or_else(|| {
12360                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
12361                        })
12362                    } else {
12363                        last_col_stats.expect("sampled: last_col_stats unset at reject")
12364                    };
12365                    let q_stats = draft_stats[n_acc];
12366                    if let Some(map) = &d2t_dev {
12367                        if q_full_buf.is_none() {
12368                            q_full_buf = Some(e.zeros(n_vocab)?);
12369                        }
12370                        let qf = q_full_buf.as_mut().unwrap();
12371                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
12372                        let qf2 = q_full_buf.as_ref().unwrap();
12373                        e.residual_sample_filtered(
12374                            cb2,
12375                            Some(qf2),
12376                            n_vocab,
12377                            sp_temp,
12378                            sp_seed,
12379                            sc,
12380                            p_stats,
12381                            q_stats,
12382                            &mut sample_tok,
12383                        )?;
12384                    } else {
12385                        e.residual_sample_filtered(
12386                            cb2,
12387                            Some(&q_bufs[n_acc]),
12388                            n_vocab,
12389                            sp_temp,
12390                            sp_seed,
12391                            sc,
12392                            p_stats,
12393                            q_stats,
12394                            &mut sample_tok,
12395                        )?;
12396                    }
12397                    e.dtoh_u32(&sample_tok)?[0]
12398                };
12399                (n_acc, bonus)
12400            };
12401            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
12402            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
12403            // ordering). Walk the accepted drafts through the grammar in commit order; the
12404            // first illegal token truncates acceptance at its slot, and that slot's emission
12405            // is recomputed as the MASKED argmax of the target's own verify column — token-
12406            // identical to constrained plain greedy decode (an unmasked argmax that is
12407            // grammar-legal IS the masked argmax: masking only removes competitors). The
12408            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
12409            // measured in acceptance numbers, never hidden.
12410            let (n_acc, bonus) = match constraint.as_deref_mut() {
12411                None => (n_acc, bonus),
12412                Some(c) => {
12413                    fn ce(e2: String) -> Box<dyn std::error::Error> {
12414                        format!("constraint: {e2}").into()
12415                    }
12416                    let mut na = n_acc;
12417                    let mut cut = false;
12418                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
12419                        if c.is_allowed(d).map_err(ce)? {
12420                            c.consume(d).map_err(ce)?;
12421                        } else {
12422                            na = j;
12423                            cut = true;
12424                            dm_cut_tokens += n_acc - j;
12425                            break;
12426                        }
12427                    }
12428                    if cut {
12429                        dm_cuts += 1;
12430                    }
12431                    let mut bo = bonus;
12432                    if cut || !c.is_allowed(bo).map_err(ce)? {
12433                        let mut row = if na == 0 && base == 0 {
12434                            init_logits_host
12435                                .clone()
12436                                .ok_or("constraint: init logits missing (round-0 cut)")?
12437                        } else {
12438                            e.dtoh_view(
12439                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
12440                            )?
12441                        };
12442                        c.mask_logits(&mut row).map_err(ce)?;
12443                        bo = argmax(&row) as u32;
12444                    }
12445                    c.consume(bo).map_err(ce)?;
12446                    (na, bo)
12447                }
12448            };
12449            let mut successor_valid = false;
12450            if let Some((q_proxy, expected_d2)) = rejected_probe {
12451                let v_n = n_acc == 1 && bonus == expected_d2;
12452                eprintln!(
12453                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
12454                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
12455                );
12456            }
12457            if let Some(successor) = successor_attempt.as_ref() {
12458                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
12459                let generation = successor.generation;
12460                let q_proxy = successor.q_proxy;
12461                let expected_pending = successor.verify_tokens[0];
12462                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
12463                let fork = opti_fork
12464                    .as_mut()
12465                    .ok_or("optipipe successor resolution lost fork state")?;
12466                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
12467                if successor_valid {
12468                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12469                } else {
12470                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12471                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12472                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
12473                }
12474                let breaker_tripped = fork
12475                    .controller
12476                    .as_mut()
12477                    .expect("controller policy")
12478                    .resolve(successor_valid);
12479                if breaker_tripped {
12480                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12481                }
12482                eprintln!(
12483                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
12484                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
12485                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
12486                    generation.id, successor_valid, !successor_valid, breaker_tripped,
12487                );
12488                if !successor_valid {
12489                    let mut successor = successor_attempt
12490                        .take()
12491                        .expect("controller successor disappeared on miss");
12492                    successor.settle();
12493                    fork.retire(generation)?;
12494                }
12495            }
12496            total_drafted += k_round;
12497            total_accepted += n_acc;
12498            if let Some(t) = sess_telem {
12499                // Greedy, rejection-sampling, and grammar truncation all converge here after
12500                // the accept decision is already on host. Fixed-size relaxed atomics only.
12501                t.record_round(k_round, n_acc);
12502            }
12503            if spec_stats {
12504                st_len_hist[k_round] += 1;
12505                for j in 0..k_round {
12506                    st_drafted[j] += 1;
12507                }
12508                for j in 0..n_acc {
12509                    st_accepted[j] += 1;
12510                }
12511                if n_acc == k_round {
12512                    st_full += 1;
12513                }
12514            }
12515
12516            if debug_spec {
12517                eprintln!(
12518                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
12519                    out.len(),
12520                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
12521                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
12522                    // the GPU worker thread — a debug flag that killed the exact regime you would
12523                    // set it to investigate. See `debug_t_pred0`.
12524                    debug_t_pred0(sampled, base, last_pred, &preds)
12525                );
12526            }
12527
12528            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
12529            let commit_started = std::time::Instant::now();
12530            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
12531            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
12532            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
12533            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
12534            for j in 0..n_acc {
12535                if !session_mode && out.len() >= max_new {
12536                    break;
12537                }
12538                out.push(draft[j]);
12539            }
12540            if pen_on {
12541                pen_hist.extend_from_slice(&draft[0..n_acc]);
12542                pen_hist.push(bonus);
12543            }
12544            let bonus_emitted = session_mode || out.len() < max_new;
12545            if bonus_emitted {
12546                out.push(bonus);
12547            }
12548            last_token = bonus;
12549
12550            // --- 5. ROLLBACK + advance (§C) ---
12551            if n_acc == k_round && !spec_replay {
12552                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
12553                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
12554                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
12555                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
12556                // last_pred is dead in the pending path (t_pred reads verify col 0).
12557                //
12558                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
12559                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
12560                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
12561                // trunk hidden (the last verify column). set_len first: a p-min break may have
12562                // left one extra chain append at that slot. Partial accepts need NO fill (the
12563                // chain already covered every accepted position; round-start set_len truncates).
12564                let mut vh_seed = e.zeros(n_embd)?;
12565                e.copy_view_into(
12566                    &mut vh_seed,
12567                    0,
12568                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
12569                    n_embd,
12570                )?;
12571                if refresh {
12572                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
12573                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
12574                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
12575                    // the full stack (vx) is already resident from the verify. Replaces both the
12576                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
12577                    // (draft attention quality); exactness stays the verify's job.
12578                    scratch.set_len(e, pos)?;
12579                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
12580                    // (hidden of the last committed row before this verify batch).
12581                    let mut vxs = e.zeros(t_v * n_embd)?;
12582                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12583                    if t_v > 1 {
12584                        e.copy_view_into(
12585                            &mut vxs,
12586                            n_embd,
12587                            &vx.slice(0..(t_v - 1) * n_embd),
12588                            (t_v - 1) * n_embd,
12589                        )?;
12590                    }
12591                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
12592                } else {
12593                    scratch.set_len(e, pos + base + k_round - 1)?;
12594                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
12595                    let mut hp = e.zeros(n_embd)?;
12596                    if t_v >= 2 {
12597                        e.copy_view_into(
12598                            &mut hp,
12599                            0,
12600                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
12601                            n_embd,
12602                        )?;
12603                    } else {
12604                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
12605                    }
12606                    self.mtp_kv_fill_all(
12607                        e,
12608                        &[draft[k_round - 1]],
12609                        &hp,
12610                        pos + base + k_round - 1,
12611                        &mut *scratch,
12612                        embd_dev,
12613                    )?;
12614                }
12615                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
12616                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
12617                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
12618                // col). Saves one MTP-block pass per round on top of the pairing fix.
12619                if !devacc_seeded {
12620                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
12621                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
12622                }
12623                pending = Some(bonus);
12624                if debug_spec {
12625                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
12626                }
12627            } else if !spec_replay && base + n_acc >= 1 {
12628                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
12629                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
12630                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
12631                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
12632                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
12633                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
12634                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
12635                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
12636                // accept (never compounds: the next verify recomputes true hiddens for all
12637                // committed columns).
12638                let j = base + n_acc;
12639                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
12640                // column stash was written into the graphs ctx's persistent slabs as in-graph
12641                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
12642                // commit must take the slab twin (same semantics, slab-addressed sources). The
12643                // ctx states which of the two this round produced via `round_slab`; trusting the
12644                // flag rather than the env keeps a round that fell back to the eager walk (a
12645                // capture that declined, a t the pool never captured) on the cols arm.
12646                let slab_commit = vg_guard
12647                    .as_ref()
12648                    .and_then(|g| g.as_ref())
12649                    .map(|g| g.round_slab)
12650                    .unwrap_or(false);
12651                if slab_commit {
12652                    self.dspark_commit_prefix_slab(
12653                        e,
12654                        &mut *cache,
12655                        &snap,
12656                        vg_guard
12657                            .as_ref()
12658                            .and_then(|g| g.as_ref())
12659                            .expect("slab_commit implies a graphs ctx"),
12660                        j,
12661                    )?;
12662                } else {
12663                    self.commit_verified_prefix(
12664                        e,
12665                        &mut *cache,
12666                        &snap,
12667                        ckpt.as_ref().unwrap(),
12668                        j,
12669                        devacc_seeded,
12670                        if devacc_seeded {
12671                            devacc_acc.as_ref().map(|a| (a, base, t_v))
12672                        } else {
12673                            None
12674                        },
12675                    )?;
12676                }
12677                let mut seed = e.zeros(n_embd)?;
12678                e.copy_view_into(
12679                    &mut seed,
12680                    0,
12681                    &vx.slice((j - 1) * n_embd..j * n_embd),
12682                    n_embd,
12683                )?;
12684                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
12685                // branch); without it the chain entries stand and only the tail truncates. Either
12686                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
12687                // (persistent mode), rope pos+j+1 (chain convention).
12688                if refresh {
12689                    scratch.set_len(e, pos)?;
12690                    let mut vxs = e.zeros(j * n_embd)?;
12691                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12692                    if j > 1 {
12693                        e.copy_view_into(
12694                            &mut vxs,
12695                            n_embd,
12696                            &vx.slice(0..(j - 1) * n_embd),
12697                            (j - 1) * n_embd,
12698                        )?;
12699                    }
12700                    self.mtp_kv_fill_all(
12701                        e,
12702                        &verify_tokens[0..j],
12703                        &vxs,
12704                        pos,
12705                        &mut *scratch,
12706                        embd_dev,
12707                    )?;
12708                } else {
12709                    scratch.set_len(e, pos + j)?;
12710                }
12711                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
12712                // bonus's predecessor (verify col j-1); no pseudo pass.
12713                if !devacc_seeded {
12714                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
12715                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
12716                }
12717                pending = Some(bonus);
12718                if debug_spec {
12719                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
12720                }
12721            } else if !spec_replay {
12722                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
12723                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
12724                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
12725                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
12726                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
12727                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
12728                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
12729                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
12730                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
12731                cache.rollback(e, &snap, 0)?;
12732                scratch.set_len(e, pos)?;
12733                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12734                pending = Some(bonus);
12735                if debug_spec {
12736                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
12737                }
12738            } else {
12739                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
12740                // this round survives, only possible before the first pending exists, ~round 0):
12741                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
12742                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
12743                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
12744                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
12745                // trunk hidden.
12746                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
12747                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
12748                if let Some(b) = pending.take() {
12749                    replay.push(b);
12750                }
12751                replay.extend_from_slice(&draft[0..n_acc]);
12752                replay.push(bonus);
12753                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
12754                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
12755                // last col exactly as before (byte-identical to the old _h_emb_dev call).
12756                let (rl_d, rx) = if self.batched_serving_numeric_class() {
12757                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
12758                    let mut hidden = e.uninit(replay.len() * n_embd)?;
12759                    for (row, &token) in replay.iter().enumerate() {
12760                        let (row_logits, row_hidden) =
12761                            self.spec_target_step_h(e, token, &mut *cache)?;
12762                        logits.extend_from_slice(&row_logits);
12763                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
12764                    }
12765                    (e.htod(&logits)?, hidden)
12766                } else {
12767                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
12768                };
12769                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
12770                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
12771                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
12772                last_pred = e.dtoh_u32(&preds_d)?[0];
12773                if sampled {
12774                    let lr0 = replay.len();
12775                    let lc = last_col_logits
12776                        .as_mut()
12777                        .expect("sampled: last_col_logits unset");
12778                    e.copy_view_into(
12779                        lc,
12780                        0,
12781                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
12782                        n_vocab,
12783                    )?;
12784                }
12785                let lr = replay.len();
12786                if lr >= 2 {
12787                    e.copy_view_into(
12788                        &mut h_seed_buf,
12789                        0,
12790                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
12791                        n_embd,
12792                    )?;
12793                } else {
12794                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
12795                    // last_token, whose own-row hidden fill_prev still holds.
12796                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
12797                }
12798                // the bonus is COMMITTED here — it becomes the last committed row.
12799                let mut rh_last = e.zeros(n_embd)?;
12800                e.copy_view_into(
12801                    &mut rh_last,
12802                    0,
12803                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
12804                    n_embd,
12805                )?;
12806                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
12807                if debug_spec {
12808                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
12809                }
12810            }
12811            if devacc_seeded {
12812                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
12813                // consumed the old value (both slots carry the same value in every non-replay arm).
12814                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12815            }
12816            if successor_valid {
12817                let optimistic_scratch_len = successor_attempt
12818                    .as_ref()
12819                    .expect("valid controller successor disappeared")
12820                    .scratch_len;
12821                // The normal current-round commit refreshed/truncated the logical scratch tail.
12822                // Its optimistic successor row was already written physically, so restoring only
12823                // the retained logical length makes that row live for the carried round.
12824                scratch.set_len(e, optimistic_scratch_len)?;
12825            }
12826            if let Some(current) = current_opti.take() {
12827                opti_fork
12828                    .as_mut()
12829                    .ok_or("optipipe current retirement lost fork state")?
12830                    .retire(current.generation)?;
12831            }
12832            if successor_valid {
12833                let successor = successor_attempt
12834                    .take()
12835                    .expect("valid controller successor disappeared before promotion");
12836                let generation = successor.generation;
12837                opti_fork
12838                    .as_mut()
12839                    .ok_or("optipipe successor promotion lost fork state")?
12840                    .promote_successor_snapshot(&mut snap, generation);
12841                carried_opti = Some(successor);
12842            }
12843            if anatomy_on {
12844                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
12845                // only for this diagnostic so it does not disappear into the following draft's
12846                // first token readback.
12847                e.stream().synchronize()?;
12848                ph_commit += commit_started.elapsed().as_secs_f64();
12849            }
12850            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
12851            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
12852            // final position — the floor's position key reads the committed depth). Burst
12853            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
12854            // like gemma's burst arm.
12855            if adapt {
12856                let fl_now = floor_at(cache.pos);
12857                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
12858            }
12859            ph_mark(&mut ph_rest, phase_on);
12860            if let Some(p) = pipe {
12861                p.accept_end(round);
12862            }
12863            drop(pipe_accept);
12864            round += 1;
12865            // sse-cadence: this round's accepted drafts + bonus are committed (out is
12866            // append-only past step 4) — flush at round cadence.
12867            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12868        }
12869        if let Some(mut ticket) = carried_opti.take() {
12870            opti_fork
12871                .as_mut()
12872                .ok_or("optipipe tail drain lost fork state")?
12873                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
12874        }
12875        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
12876        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
12877        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
12878
12879        if spec_stats {
12880            let per_slot: Vec<String> = (0..k)
12881                .map(|j| {
12882                    if st_drafted[j] > 0 {
12883                        format!(
12884                            "{}/{}={:.3}",
12885                            st_accepted[j],
12886                            st_drafted[j],
12887                            st_accepted[j] as f64 / st_drafted[j] as f64
12888                        )
12889                    } else {
12890                        "0/0".into()
12891                    }
12892                })
12893                .collect();
12894            let acc = if total_drafted > 0 {
12895                total_accepted as f64 / total_drafted as f64
12896            } else {
12897                0.0
12898            };
12899            eprintln!(
12900                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
12901                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
12902                       tok_per_round={:.3}",
12903                per_slot.join(" "),
12904                (total_accepted + round) as f64 / round.max(1) as f64
12905            );
12906        }
12907        if constraint.is_some() {
12908            eprintln!(
12909                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
12910                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
12911                dm_clone_ns as f64 / 1e6,
12912                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
12913            );
12914        }
12915        if phase_on {
12916            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
12917            eprintln!(
12918                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
12919                ph_draft * 1e3,
12920                ph_draft / tot * 100.0,
12921                ph_verify * 1e3,
12922                ph_verify / tot * 100.0,
12923                ph_wait * 1e3,
12924                ph_wait / tot * 100.0,
12925                ph_rest * 1e3,
12926                ph_rest / tot * 100.0
12927            );
12928        }
12929        if anatomy_on {
12930            let rounds_f = round.max(1) as f64;
12931            let other = (ph_rest - ph_commit).max(0.0);
12932            eprintln!(
12933                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
12934                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
12935                ph_draft * 1e3 / rounds_f,
12936                ph_verify * 1e3 / rounds_f,
12937                ph_wait * 1e3 / rounds_f,
12938                ph_commit * 1e3 / rounds_f,
12939                other * 1e3 / rounds_f,
12940            );
12941        }
12942        let _pipe_tail = pipe.map(|p| p.primary());
12943        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
12944        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
12945        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
12946        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
12947        if let Some(slot) = sess_draft_slot.take() {
12948            *slot = Some(dctx);
12949        }
12950        let t_rounds = t_ent.elapsed();
12951        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
12952            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
12953            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
12954            // HERE, where the sampler, the session Philox counters and the penalty window are
12955            // all live and the boundary logits row still exists — that is the "make the state
12956            // available" half of the fix; the consuming burst then just emits it. `sctr` is
12957            // written to the session BELOW the draws so the advance is never lost.
12958            *next_pred_slot = Some(last_pred);
12959            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
12960            let mut stashed_pending = false;
12961            if let Some(b) = pending.take() {
12962                if !sampled {
12963                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
12964                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
12965                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
12966                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
12967                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
12968                    // OUT of `committed` (cache rows == committed); the consuming call
12969                    // prepends it once its verify commits the row. next_pred is unknowable
12970                    // without the commit pass — None; callers gate on pending_tok too.
12971                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
12972                    if let Some(slot) = sess_pending_slot.take() {
12973                        *slot = Some(b);
12974                    }
12975                    *next_pred_slot = None;
12976                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
12977                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
12978                    *last_h = Some(e.clone_dtod(&fill_prev)?);
12979                    stashed_pending = true;
12980                } else {
12981                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
12982                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
12983                    let pos_b = cache.pos;
12984                    scratch.set_len(e, pos_b)?;
12985                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
12986                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
12987                    // itself — the prediction AFTER the bonus never materialized; it would have
12988                    // been the next round's verify col 0). The commit's logits ARE that
12989                    // prediction — so they are also the row the next burst's boundary token
12990                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
12991                    *next_pred_slot = Some(if sample_boundary {
12992                        sample_boundary_token(
12993                            e,
12994                            &lg_b,
12995                            &sp,
12996                            &pen_hist,
12997                            &mut sctr,
12998                            "burst-tail-commit",
12999                        )?
13000                    } else {
13001                        argmax(&lg_b) as u32
13002                    });
13003                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13004                    *last_h = Some(hb);
13005                }
13006            } else {
13007                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13008                *last_h = Some(e.clone_dtod(&fill_prev)?);
13009                if sample_boundary {
13010                    // No pending to commit, so the boundary row is the one `last_pred` was
13011                    // argmaxed from and the sampled path keeps it on device: the init feed's
13012                    // logits when the burst ran zero rounds, else the legacy-replay path's
13013                    // last verify column (both predict the token AFTER the last committed
13014                    // row). It is retained precisely because round 0's accept test needs it,
13015                    // so the draw costs no extra D2H of the [n_vocab] row.
13016                    match last_col_logits.as_ref() {
13017                        Some(lc) => {
13018                            *next_pred_slot = Some(sample_boundary_token_dev(
13019                                e,
13020                                lc,
13021                                n_vocab,
13022                                &sp,
13023                                &pen_hist,
13024                                &mut sctr,
13025                                "burst-tail-nopending",
13026                            )?);
13027                        }
13028                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
13029                        // burst always feeds or replays, so the row exists — but if it ever
13030                        // is, the stream takes a greedy token and SAYS so rather than
13031                        // silently regressing to the pre-lane behaviour.
13032                        None => eprintln!(
13033                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13034                             (reason: no retained boundary logits row)"
13035                        ),
13036                    }
13037                }
13038            }
13039            *sctr_slot = sctr;
13040            *uctr_slot = uctr;
13041            committed.extend_from_slice(prompt);
13042            if let Some(cb) = carried_pending {
13043                // the consumed carry's cache row landed in round 0's verify (every pending
13044                // round commits col 0) — it joins `committed` here, in sequence order.
13045                committed.push(cb);
13046            }
13047            if stashed_pending {
13048                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13049                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13050                // 18446744073709551615 out of range for slice of length 0", killing the
13051                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13052                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13053                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13054                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13055                // did). So a burst that stashes a pending without emitting anything of its own —
13056                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13057                // guard skipping every token under a tight budget — arrives here with
13058                // out.len() == 0 and stashed_pending == true.
13059                //
13060                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13061                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13062                // just above is already accounted. Saturating, not a min/assert: an empty `out`
13063                // here is a legitimate burst shape, not a corrupt state.
13064                let emitted = out.len().saturating_sub(1);
13065                committed.extend_from_slice(&out[..emitted]);
13066            } else {
13067                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13068            }
13069            debug_assert_eq!(
13070                cache.pos,
13071                committed.len(),
13072                "session invariant: cache rows == committed tokens"
13073            );
13074            if setup_trace {
13075                e.stream().synchronize()?; // bound the async tail fill in the trace
13076                let t_tail = t_ent.elapsed();
13077                eprintln!(
13078                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13079                    t_init.as_secs_f64() * 1e3,
13080                    (t_cap - t_init).as_secs_f64() * 1e3,
13081                    (t_fill - t_cap).as_secs_f64() * 1e3,
13082                    (t_rounds - t_fill).as_secs_f64() * 1e3,
13083                    (t_tail - t_rounds).as_secs_f64() * 1e3,
13084                    t_tail.as_secs_f64() * 1e3,
13085                    out.len(),
13086                    continuation
13087                );
13088            }
13089            return Ok((out, total_drafted, total_accepted));
13090        }
13091        out.truncate(max_new);
13092        Ok((out, total_drafted, total_accepted))
13093    }
13094
13095    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13096    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13097    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13098    pub fn extract_dspark_anchors(
13099        &self,
13100        e: &Engine,
13101        tokens: &[u32],
13102        anchor_positions: &[usize],
13103        gamma: usize,
13104        top_k: usize,
13105        chunk: usize,
13106        temperature: f32,
13107    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13108        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13109            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13110        }
13111        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13112            return Err("DSpark anchor positions must be sorted and unique".into());
13113        }
13114        for &position in anchor_positions {
13115            if position == 0 || position + gamma >= tokens.len() {
13116                return Err(format!(
13117                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13118                    tokens.len()
13119                )
13120                .into());
13121            }
13122        }
13123
13124        let n_vocab = self.output.out_features();
13125        let n_embd = self.cfg.n_embd as usize;
13126        let mut cache =
13127            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13128        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13129        let embd_gpu = if spec_host_embd() {
13130            None
13131        } else {
13132            Some(
13133                self.embd_gpu
13134                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13135            )
13136        };
13137        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13138
13139        struct PendingRecord {
13140            position: usize,
13141            hidden: Option<Vec<f32>>,
13142            tokens: Vec<u32>,
13143            target_top_ids: Vec<Option<Vec<u32>>>,
13144            target_top_logits: Vec<Option<Vec<f32>>>,
13145            target_top_probs: Vec<Option<Vec<f32>>>,
13146            target_tail_probs: Vec<Option<f32>>,
13147        }
13148
13149        let mut pending: Vec<PendingRecord> = anchor_positions
13150            .iter()
13151            .map(|&position| PendingRecord {
13152                position,
13153                hidden: None,
13154                tokens: tokens[position..=position + gamma].to_vec(),
13155                target_top_ids: vec![None; gamma],
13156                target_top_logits: vec![None; gamma],
13157                target_top_probs: vec![None; gamma],
13158                target_tail_probs: vec![None; gamma],
13159            })
13160            .collect();
13161
13162        let mut start = 0usize;
13163        while start < tokens.len() {
13164            let end = (start + chunk).min(tokens.len());
13165            let chunk_tokens = &tokens[start..end];
13166            let (target_logits, hidden_rows) =
13167                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13168            for record in &mut pending {
13169                let hidden_position = record.position - 1;
13170                if hidden_position >= start && hidden_position < end {
13171                    let local = hidden_position - start;
13172                    record.hidden = Some(
13173                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13174                    );
13175                }
13176                for slot in 0..gamma {
13177                    let target_row = record.position + slot;
13178                    if target_row < start || target_row >= end {
13179                        continue;
13180                    }
13181                    let local = target_row - start;
13182                    let logits =
13183                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13184                    let (ids, top_logits, probs, tail) =
13185                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13186                    record.target_top_ids[slot] = Some(ids);
13187                    record.target_top_logits[slot] = Some(top_logits);
13188                    record.target_top_probs[slot] = Some(probs);
13189                    record.target_tail_probs[slot] = Some(tail);
13190                }
13191            }
13192            start = end;
13193        }
13194
13195        pending
13196            .into_iter()
13197            .map(|record| {
13198                let hidden = record
13199                    .hidden
13200                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13201                let target_top_ids =
13202                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13203                let target_top_logits = flatten_dspark_rows(
13204                    record.target_top_logits,
13205                    record.position,
13206                    "target logits",
13207                )?;
13208                let target_top_probs =
13209                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13210                let target_tail_probs = record
13211                    .target_tail_probs
13212                    .into_iter()
13213                    .enumerate()
13214                    .map(|(slot, value)| {
13215                        value.ok_or_else(|| {
13216                            format!("missing DSpark tail at {} slot {slot}", record.position)
13217                        })
13218                    })
13219                    .collect::<Result<Vec<_>, _>>()?;
13220                Ok(DsparkAnchorRecord {
13221                    position: record.position,
13222                    hidden,
13223                    tokens: record.tokens,
13224                    target_top_ids,
13225                    target_top_logits,
13226                    target_top_probs,
13227                    target_tail_probs,
13228                })
13229            })
13230            .collect()
13231    }
13232
13233    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13234    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13235    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13236    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13237    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13238    /// quant-induced head/hidden-state mismatch from text drift.
13239    ///
13240    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13241    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13242    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13243    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13244    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
13245    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
13246    ///              conditions on the corpus — deterministic and arm-comparable by design.
13247    ///
13248    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13249    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13250    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13251    ///
13252    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13253    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13254    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13255    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13256    /// agreement vs this path — not usable as a training-data source).
13257    pub fn replay_acceptance(
13258        &self,
13259        e: &Engine,
13260        tokens: &[u32],
13261        k: usize,
13262        stride: usize,
13263        chunk: usize,
13264        mut hdump: Option<&mut std::fs::File>,
13265    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13266        assert!(k >= 1 && stride >= 1 && chunk >= 2);
13267        let mtp = self
13268            .mtp
13269            .as_ref()
13270            .expect("replay_acceptance requires an MTP head");
13271        let n_vocab = self.output.out_features();
13272        let d_vocab = mtp
13273            .shared_head_head
13274            .as_ref()
13275            .unwrap_or(&self.output)
13276            .out_features();
13277        let n_embd = self.cfg.n_embd as usize;
13278        let t_total = tokens.len();
13279        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13280        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13281        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13282        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13283        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13284        let embd_gpu = if spec_host_embd() {
13285            None
13286        } else {
13287            Some(
13288                self.embd_gpu
13289                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13290            )
13291        };
13292        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13293
13294        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13295        let mut bg: Vec<u32> = vec![0; t_total + 1];
13296        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
13297        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
13298        let mut seed_buf = e.zeros(n_embd)?;
13299        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
13300        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
13301        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
13302        let mut s = 0usize;
13303        while s < t_total {
13304            let cend = (s + chunk).min(t_total);
13305            let tc = cend - s;
13306            let ch = &tokens[s..cend];
13307            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
13308            //    the chunk's true hiddens.
13309            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
13310            for j in 0..tc {
13311                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13312            }
13313            let preds = e.dtoh_u32(&preds_d)?;
13314            for j in 0..tc {
13315                bg[s + j + 1] = preds[j];
13316            }
13317            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
13318            // checkpoint-quality metric (position j's logits score the GOLD next token).
13319            if nll_on {
13320                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
13321                if jmax > 0 {
13322                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
13323                    let rows: Vec<i32> = (0..jmax as i32).collect();
13324                    let idsd = e.htod_u32_v(&ids)?;
13325                    let rowsd = e.htod_i32(&rows)?;
13326                    let mut outd = e.zeros(jmax)?;
13327                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
13328                    for pr in e.dtoh(&outd)? {
13329                        nll_sum += -((pr.max(1e-30)) as f64).ln();
13330                        nll_cnt += 1;
13331                    }
13332                }
13333            }
13334            if let Some(f) = hdump.as_deref_mut() {
13335                use std::io::Write;
13336                let host: Vec<f32> = e.dtoh(&vx)?;
13337                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
13338                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
13339                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
13340                for v in &host[..tc * n_embd] {
13341                    let b = v.to_bits();
13342                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
13343                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
13344                }
13345                f.write_all(&bytes)?;
13346            }
13347            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
13348            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
13349            // per token saved; the forced trunk pass + hdump is all the mode needs).
13350            let chainless = stride > t_total;
13351            if chainless {
13352                e.copy_view_into(
13353                    &mut prev_last_h,
13354                    0,
13355                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
13356                    n_embd,
13357                )?;
13358                s = cend;
13359                continue;
13360            }
13361            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
13362            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
13363            let mut vxs = e.zeros(tc * n_embd)?;
13364            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
13365            if tc > 1 {
13366                e.copy_view_into(
13367                    &mut vxs,
13368                    n_embd,
13369                    &vx.slice(0..(tc - 1) * n_embd),
13370                    (tc - 1) * n_embd,
13371                )?;
13372            }
13373            scratch.set_len(e, s)?;
13374            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13375            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
13376            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
13377            //    truncates those approximate appends before they can ever be read.
13378            let ps: Vec<usize> = (s..cend)
13379                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
13380                .collect();
13381            for &p in ps.iter().rev() {
13382                scratch.set_len(e, p)?;
13383                if p == s {
13384                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
13385                } else {
13386                    e.copy_view_into(
13387                        &mut seed_buf,
13388                        0,
13389                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
13390                        n_embd,
13391                    )?;
13392                }
13393                let mut e_tok = tokens[p];
13394                let mut d_seed = e.clone_dtod(&seed_buf)?;
13395                let chain_heads = !self.mtp_extra.is_empty();
13396                let mut chain_tokens = if chain_heads {
13397                    vec![tokens[p]]
13398                } else {
13399                    Vec::new()
13400                };
13401                let mut chain_seeds = if chain_heads {
13402                    vec![e.clone_dtod(&seed_buf)?]
13403                } else {
13404                    Vec::new()
13405                };
13406                let mut drafts: Vec<u32> = Vec::with_capacity(k);
13407                for j in 0..k {
13408                    let (dl_d, h_nextn) = if chain_heads {
13409                        self.mtp_chain_forward_dev(
13410                            e,
13411                            &chain_tokens,
13412                            &chain_seeds,
13413                            &mut scratch,
13414                            p,
13415                            embd_dev,
13416                            None,
13417                        )?
13418                    } else {
13419                        self.mtp_head_forward_dev(
13420                            e,
13421                            mtp,
13422                            e_tok,
13423                            &d_seed,
13424                            &mut scratch,
13425                            p + 1 + j,
13426                            embd_dev,
13427                            None,
13428                        )?
13429                    };
13430                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
13431                    let idx = e.dtoh_u32_one(&tok_d)?;
13432                    let d = match &mtp.d2t {
13433                        Some(map) => map[idx as usize],
13434                        None => idx,
13435                    };
13436                    drafts.push(d);
13437                    if chain_heads {
13438                        chain_tokens.push(d);
13439                        chain_seeds.push(h_nextn);
13440                    } else {
13441                        e_tok = d;
13442                        d_seed = h_nextn;
13443                    }
13444                }
13445                // targets may live in a LATER chunk's bg — resolved after the walk.
13446                rows.push((p, drafts, Vec::new()));
13447            }
13448            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
13449            //    expect scratch.len == cend with exact rows).
13450            scratch.set_len(e, s)?;
13451            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13452            e.copy_view_into(
13453                &mut prev_last_h,
13454                0,
13455                &vx.slice((tc - 1) * n_embd..tc * n_embd),
13456                n_embd,
13457            )?;
13458            s = cend;
13459        }
13460        for (p, drafts, targets) in rows.iter_mut() {
13461            for j in 0..drafts.len() {
13462                targets.push(bg[*p + 1 + j]);
13463            }
13464        }
13465        rows.sort_by_key(|r| r.0);
13466        if nll_cnt > 0 {
13467            let mean = nll_sum / nll_cnt as f64;
13468            println!(
13469                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
13470                mean.exp()
13471            );
13472        }
13473        Ok((rows, bg))
13474    }
13475}
13476
13477#[cfg(test)]
13478mod mtp_chain_tests {
13479    use super::mtp_chain_head_index;
13480
13481    #[test]
13482    fn embedded_step_heads_cycle_in_declared_order() {
13483        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
13484        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
13485    }
13486
13487    #[test]
13488    fn standalone_draft_remains_single_head() {
13489        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
13490    }
13491}
13492
13493#[cfg(test)]
13494mod tp_verified_prefix_tests {
13495    use super::rewind_tp_kv_verified_prefix;
13496    use crate::tp::ResidentTpKvCache;
13497
13498    fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
13499        let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
13500        let transaction = cache.begin_transaction().unwrap();
13501        let target = cache.append_target(transaction, committed).unwrap();
13502        cache.publish_append(transaction, target).unwrap();
13503        let target = cache.commit_target(transaction, committed).unwrap();
13504        cache.publish_finalize(transaction, target).unwrap();
13505        cache
13506    }
13507
13508    #[test]
13509    fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
13510        let mut layers = vec![Some(cache_with_committed_len(5)), None];
13511        rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
13512        let cache = layers[0].as_ref().unwrap();
13513        assert_eq!(cache.committed_len(), 3);
13514        assert_eq!(cache.staged_len(), 3);
13515    }
13516
13517    #[test]
13518    fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
13519        let mut layers = vec![Some(cache_with_committed_len(1))];
13520        let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
13521            .unwrap_err()
13522            .to_string();
13523        assert!(error.contains("changed shape"), "unexpected error: {error}");
13524    }
13525}
13526
13527#[cfg(test)]
13528mod dspark_sparse_tests {
13529    use super::dspark_sparse_softmax_topk;
13530
13531    #[test]
13532    fn topk_keeps_full_softmax_mass_and_stable_ties() {
13533        let logits = [1.0f32, 3.0, 3.0, -2.0];
13534        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
13535        assert_eq!(ids, vec![1, 2]);
13536        assert_eq!(top_logits, vec![3.0, 3.0]);
13537        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
13538        let expected = 1.0 / denominator;
13539        assert!((probs[0] - expected).abs() < 1.0e-6);
13540        assert!((probs[1] - expected).abs() < 1.0e-6);
13541        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
13542        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
13543    }
13544}
13545
13546#[cfg(test)]
13547mod spec_replay_env_tests {
13548    use super::spec_replay_env_on;
13549
13550    #[test]
13551    fn replay_requires_literal_one() {
13552        assert!(!spec_replay_env_on(None));
13553        assert!(!spec_replay_env_on(Some("")));
13554        assert!(!spec_replay_env_on(Some("0")));
13555        assert!(!spec_replay_env_on(Some("true")));
13556        assert!(!spec_replay_env_on(Some("2")));
13557        assert!(spec_replay_env_on(Some("1")));
13558    }
13559}
13560
13561#[cfg(test)]
13562mod telem_tests {
13563    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
13564
13565    #[test]
13566    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
13567        let counters = SpecTelemetryCounters::default();
13568        for mask in [
13569            [true, true, true],
13570            [true, true, false],
13571            [true, false, false],
13572            [false, false, false],
13573        ] {
13574            let accepted = mask.iter().take_while(|&&value| value).count();
13575            counters.record_round(mask.len(), accepted);
13576        }
13577
13578        let snapshot = counters.snapshot();
13579        assert_eq!(
13580            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
13581            (4, 12, 6)
13582        );
13583        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
13584        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
13585        assert_eq!(snapshot.tau(), 1.5);
13586        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
13587        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
13588    }
13589
13590    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
13591    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
13592    #[test]
13593    fn delta_isolates_burst_contribution() {
13594        let mut t = SpecTelemetry::default();
13595        // "previous request": 2 rounds of k=3, accepts 3 then 1.
13596        for (kr, na) in [(3usize, 3usize), (3, 1)] {
13597            t.rounds += 1;
13598            t.drafted += kr as u64;
13599            t.accepted += na as u64;
13600            for j in 0..kr {
13601                t.pos_drafted[j] += 1;
13602            }
13603            for j in 0..na {
13604                t.pos_accepted[j] += 1;
13605            }
13606        }
13607        let before = t;
13608        // "this burst": 1 round k=3, accepts 2.
13609        t.rounds += 1;
13610        t.drafted += 3;
13611        t.accepted += 2;
13612        for j in 0..3 {
13613            t.pos_drafted[j] += 1;
13614        }
13615        for j in 0..2 {
13616            t.pos_accepted[j] += 1;
13617        }
13618        let d = t.delta_since(&before);
13619        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
13620        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
13621        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
13622        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
13623    }
13624
13625    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
13626    /// aggregation invariant.
13627    #[test]
13628    fn merge_accumulates_fieldwise() {
13629        let mut agg = SpecTelemetry::default();
13630        let mut d1 = SpecTelemetry {
13631            rounds: 2,
13632            drafted: 6,
13633            accepted: 4,
13634            ..Default::default()
13635        };
13636        d1.pos_drafted[0] = 2;
13637        d1.pos_accepted[0] = 2;
13638        let mut d2 = SpecTelemetry {
13639            rounds: 1,
13640            drafted: 3,
13641            accepted: 1,
13642            ..Default::default()
13643        };
13644        d2.pos_drafted[0] = 1;
13645        d2.pos_accepted[0] = 1;
13646        d2.pos_drafted[1] = 1;
13647        agg.merge(&d1);
13648        agg.merge(&d2);
13649        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
13650        assert_eq!(agg.pos_drafted[0], 3);
13651        assert_eq!(agg.pos_accepted[0], 3);
13652        assert_eq!(agg.pos_drafted[1], 1);
13653        assert_eq!(agg.pos_accepted[1], 0);
13654    }
13655
13656    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
13657    /// public metrics surface and must never publish a u64-wrapped garbage value.
13658    #[test]
13659    fn delta_saturates_never_wraps() {
13660        let small = SpecTelemetry {
13661            rounds: 1,
13662            drafted: 2,
13663            accepted: 1,
13664            ..Default::default()
13665        };
13666        let big = SpecTelemetry {
13667            rounds: 5,
13668            drafted: 15,
13669            accepted: 9,
13670            ..Default::default()
13671        };
13672        let d = small.delta_since(&big);
13673        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
13674    }
13675}
13676
13677#[cfg(test)]
13678mod opti_fork_tests {
13679    use super::{
13680        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
13681    };
13682
13683    #[test]
13684    fn controller_threshold_and_three_miss_breaker_are_exact() {
13685        let mut policy = OptiControllerPolicy {
13686            threshold: 0.7,
13687            consecutive_misses: 0,
13688            breaker_tripped: false,
13689        };
13690        assert!(!policy.admit(0.699_999));
13691        assert!(policy.admit(0.7));
13692        assert!(!policy.resolve(false));
13693        assert!(!policy.resolve(false));
13694        assert!(policy.resolve(false));
13695        assert!(policy.breaker_tripped);
13696        assert!(!policy.admit(1.0));
13697        assert!(
13698            !policy.resolve(true),
13699            "a resolved hit cannot re-arm a tripped request"
13700        );
13701        assert!(policy.breaker_tripped);
13702    }
13703
13704    #[test]
13705    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
13706        let mut policy = OptiControllerPolicy {
13707            threshold: 0.0,
13708            consecutive_misses: 0,
13709            breaker_tripped: false,
13710        };
13711        for _ in 0..16 {
13712            assert!(policy.admit(0.0));
13713            assert!(!policy.resolve(false));
13714        }
13715        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
13716            assert!(
13717                !policy.admit(invalid),
13718                "invalid q proxy must fail closed: {invalid}"
13719            );
13720        }
13721        assert!(!policy.breaker_tripped);
13722        assert_eq!(policy.consecutive_misses, 0);
13723    }
13724
13725    #[test]
13726    fn alternating_mode_flips_by_generation_not_round_parity() {
13727        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
13728        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
13729        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
13730        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
13731    }
13732
13733    #[test]
13734    fn live_generation_cannot_be_overwritten() {
13735        let mut tracker = OptiForkGenerationTracker::default();
13736        let g0 = tracker.reserve().unwrap();
13737        let g1 = tracker.reserve().unwrap();
13738        let err = tracker.reserve().unwrap_err().to_string();
13739        assert!(
13740            err.contains("still owns generation 0"),
13741            "unexpected error: {err}"
13742        );
13743        tracker.retire(g0).unwrap();
13744        let g2 = tracker.reserve().unwrap();
13745        assert_eq!((g2.id, g2.slot), (2, 0));
13746        tracker.retire(g1).unwrap();
13747        tracker.retire(g2).unwrap();
13748    }
13749
13750    #[test]
13751    fn teardown_rejects_a_stale_generation_tag() {
13752        let mut tracker = OptiForkGenerationTracker::default();
13753        let g0 = tracker.reserve().unwrap();
13754        tracker.retire(g0).unwrap();
13755        let err = tracker.retire(g0).unwrap_err().to_string();
13756        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
13757    }
13758}
13759
13760#[cfg(test)]
13761mod draft_graph_fallback_tests {
13762    use super::DraftGraphFallback;
13763
13764    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
13765    #[test]
13766    fn flip_is_loud_once_and_memoized_after() {
13767        let mut f = DraftGraphFallback::default();
13768        let line = f
13769            .mark_greedy("out of memory")
13770            .expect("first flip must return the warn line");
13771        assert!(
13772            line.contains("WARN"),
13773            "flip line must be warn-level: {line}"
13774        );
13775        assert!(
13776            line.contains("out of memory"),
13777            "flip line must carry the reason: {line}"
13778        );
13779        assert!(f.greedy_failed());
13780        // re-marking an already-failed graph is the memoization: quiet, still failed.
13781        assert!(f.mark_greedy("out of memory").is_none());
13782        assert!(f.greedy_failed());
13783        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
13784        assert!(!f.sampled_failed());
13785        let line_s = f
13786            .mark_sampled("capture unsupported")
13787            .expect("sampled flip is its own flip");
13788        assert!(
13789            line_s.contains("sampled"),
13790            "sampled flip names itself: {line_s}"
13791        );
13792        assert!(f.mark_sampled("capture unsupported").is_none());
13793    }
13794
13795    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
13796    /// and says so exactly when there was something to reset.
13797    #[test]
13798    fn reset_on_resume_clears_flags_and_logs_once() {
13799        let mut f = DraftGraphFallback::default();
13800        // clean session: resume is silent, nothing to reset.
13801        assert!(f.reset_on_resume().is_none());
13802        f.mark_greedy("oom").unwrap();
13803        f.mark_sampled("oom").unwrap();
13804        let note = f
13805            .reset_on_resume()
13806            .expect("a set flag must produce the reset note");
13807        assert!(
13808            note.contains("greedy+sampled"),
13809            "note names what was reset: {note}"
13810        );
13811        assert!(
13812            !f.greedy_failed() && !f.sampled_failed(),
13813            "both flags cleared"
13814        );
13815        // and the NEXT failure after a reset is a fresh flip — loud again.
13816        assert!(f.mark_greedy("oom again").is_some());
13817        let note2 = f.reset_on_resume().expect("greedy-only reset");
13818        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
13819    }
13820
13821    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
13822    /// they precede a fresh capture attempt whose own failure re-flips loudly.
13823    #[test]
13824    fn shape_change_clears_are_silent() {
13825        let mut f = DraftGraphFallback::default();
13826        f.mark_greedy("oom").unwrap();
13827        f.clear_greedy();
13828        assert!(!f.greedy_failed());
13829        f.mark_sampled("oom").unwrap();
13830        f.clear_sampled();
13831        assert!(!f.sampled_failed());
13832        // after a silent clear there is nothing left for resume to report.
13833        assert!(f.reset_on_resume().is_none());
13834    }
13835}
13836
13837/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
13838///
13839/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
13840/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
13841/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
13842/// than remembered.
13843#[cfg(test)]
13844mod sampled_graph_key_tests {
13845    use super::{SampledGraphKey, debug_t_pred0};
13846
13847    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
13848    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
13849        (k.seed, k.temp_bits, k.k)
13850    }
13851
13852    fn pure_temp_key() -> SampledGraphKey {
13853        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
13854        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
13855    }
13856
13857    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
13858    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
13859    #[test]
13860    fn vendor_filters_change_the_key() {
13861        let parked = pure_temp_key();
13862        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
13863        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
13864        assert_eq!(
13865            legacy_key(&parked),
13866            legacy_key(&vendor),
13867            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
13868        );
13869        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
13870        assert!(parked.pure_temp());
13871        assert!(!vendor.pure_temp());
13872    }
13873
13874    /// Each distribution-shaping field alone is enough to drop the parked graph.
13875    #[test]
13876    fn every_filter_field_is_keyed() {
13877        let base = pure_temp_key();
13878        for (what, other) in [
13879            (
13880                "top_k",
13881                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
13882            ),
13883            (
13884                "top_p",
13885                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
13886            ),
13887            (
13888                "min_p",
13889                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
13890            ),
13891            (
13892                "penalties",
13893                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
13894            ),
13895        ] {
13896            assert_ne!(base, other, "{what} must be part of the key");
13897            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
13898            assert_eq!(
13899                legacy_key(&base),
13900                legacy_key(&other),
13901                "{what} was invisible to the pre-fix key",
13902            );
13903        }
13904    }
13905
13906    /// The baked constants stay keyed (this half was always right — regression cover for it).
13907    #[test]
13908    fn baked_constants_stay_keyed() {
13909        let base = pure_temp_key();
13910        assert_ne!(
13911            base,
13912            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
13913            "seed"
13914        );
13915        assert_ne!(
13916            base,
13917            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
13918            "temp"
13919        );
13920        assert_ne!(
13921            base,
13922            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
13923            "k"
13924        );
13925        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
13926        assert_eq!(
13927            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
13928            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
13929        );
13930    }
13931
13932    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
13933    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
13934    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
13935    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
13936    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
13937    ///
13938    /// This test is the other end of that argument, asserted here rather than remembered in a
13939    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
13940    /// would silently become the unsound thing it is documented not to be.
13941    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
13942    #[test]
13943    fn seed_alone_still_rekeys_the_draft_graph() {
13944        let parked = pure_temp_key();
13945        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
13946        assert_ne!(
13947            parked, reseeded,
13948            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
13949             decision not to compare seed rests on exactly this",
13950        );
13951        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
13952        // because of a filter difference.
13953        assert!(parked.pure_temp() && reseeded.pure_temp());
13954    }
13955
13956    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
13957    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
13958    /// agree on the regime, so a graph that survives the drop is legal to launch.
13959    #[test]
13960    fn equal_keys_agree_on_the_regime() {
13961        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
13962        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
13963        assert_eq!(a, b);
13964        assert_eq!(a.pure_temp(), b.pure_temp());
13965        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
13966        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
13967        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
13968        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
13969    }
13970
13971    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
13972    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
13973    #[test]
13974    fn debug_print_survives_the_sampled_arm() {
13975        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
13976        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
13977        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
13978        // round 0 without a pending bonus still reports last_pred, in both arms.
13979        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
13980        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
13981        // greedy keeps the real prediction it always printed.
13982        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
13983        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
13984    }
13985}