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, DEFAULT ON for the GDN+MoE family since 2026-08-23
240/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
241///
242/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
243/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
244/// on this route. The MTP spec round is that caller.
245///
246/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
247/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
248/// the host is never waiting for the device, it is spending its own time launching the trunk.
249/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
250/// 8-10 ms per burst).
251///
252/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
253///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
254///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
255///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
256///     comes from per-round phase totals, which are internal to each boot).
257/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
258/// the round off the host and onto the device, which is the whole point.
259///
260/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
261/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
262/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
263/// at every K, kernel-check ALL GREEN.
264///
265/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
266/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
267/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
268/// opt in with `=1` once it has its own interleave. Also never armed together with
269/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
270pub(crate) fn spec_verify_graph_env() -> Option<bool> {
271    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
272    *ON.get_or_init(
273        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
274            Ok("1") => Some(true),
275            Ok("0") => Some(false),
276            _ => None,
277        },
278    )
279}
280/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
281/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
282/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
283/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
284/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
285/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
286/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
287/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
288/// 256-token run vs the serve session's thousands of rounds), and the two
289/// instruments must keep their own measured dispositions rather than share one flag.
290pub(crate) fn dspark_verify_graph_serve_on() -> bool {
291    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
292    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
293}
294/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
295/// pool's memory policy STATED instead of silently unbounded. The keyspace is
296/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
297/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
298/// on the q38 export — so the default (256) never engages there; the knob is the
299/// safety valve for a future export with a wider ladder. At the ceiling the pool
300/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
301/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
302/// cols-stashed layers inside one commit). No eviction by design: destroying a live
303/// exec graph re-opens the stale-address class the indirect tables exist to close,
304/// and the bounded keyspace makes reclaim worthless.
305pub(crate) fn dspark_vg_cap() -> usize {
306    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
307    *CAP.get_or_init(|| {
308        std::env::var("MEMRA_DSPARK_VG_MAX")
309            .ok()
310            .and_then(|v| v.parse().ok())
311            .unwrap_or(256)
312    })
313}
314
315/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
316/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
317/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
318/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
319/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
320/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
321///
322/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
323/// and proves nothing about another export): the debt is remaining capture slots x the
324/// MARGINAL bytes a capture adds to this device's graph mem pool.
325///
326/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
327/// version of this used the mean (`reserved / captures`) and the live serve log showed why
328/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
329/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
330/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
331/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
332/// boot can refuse admissions that would have fit, which is a worse defect than the
333/// under-charge this accounting exists to remove. The marginal reading prices what an
334/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
335/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
336/// tracks real growth on one that does.
337///
338/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
339/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
340/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
341/// the same direction as the old rule without the 255x extrapolation.
342///
343/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
344/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
345/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
346/// debt is 0 there too.
347pub fn dspark_vg_debt_projection(
348    captures: usize,
349    cap: usize,
350    reserved_bytes: usize,
351    prev: Option<(usize, usize)>,
352) -> usize {
353    if captures == 0 || cap == 0 {
354        return 0;
355    }
356    let remaining = cap.saturating_sub(captures);
357    if remaining == 0 {
358        return 0;
359    }
360    match prev {
361        // marginal growth between two observations of the same pool
362        Some((c0, r0)) if captures > c0 => {
363            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
364            remaining.saturating_mul(marginal)
365        }
366        // bootstrap: at most one more pool's worth
367        _ => remaining
368            .saturating_mul(reserved_bytes / captures)
369            .min(reserved_bytes),
370    }
371}
372/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
373/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
374/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
375/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
376/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
377/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
378/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
379/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
380/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
381/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
382/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
383/// empty partial the combine never reads, so the shared n_splits_max stride changes no
384/// bytes) and re-gated e2e by this lane's battery.
385pub(crate) fn dspark_fa_rows_on() -> bool {
386    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
387    *ON.get_or_init(|| {
388        std::env::var("MEMRA_DSPARK_FA_ROWS")
389            .map(|v| v != "0")
390            .unwrap_or(true)
391    })
392}
393
394/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
395///
396/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
397/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
398/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
399/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
400/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
401/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
402/// the flag crashed precisely the regime it exists to investigate.
403///
404/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
405/// indexing (an out-of-range pred there is a real bug and must still be loud).
406fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
407    if base == 0 {
408        return last_pred.to_string();
409    }
410    match preds.get(base - 1) {
411        Some(p) => p.to_string(),
412        // sampled: the greedy per-column argmax was never run for this round.
413        None => {
414            debug_assert!(
415                sampled,
416                "greedy spec: preds[{}] missing at base {base}",
417                base - 1
418            );
419            "n/a".to_string()
420        }
421    }
422}
423
424/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
425///
426/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
427/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
428/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
429/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
430/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
431/// not believe in — and `u * 0 < p` then accepts it unconditionally.
432///
433/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
434/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
435pub(crate) fn skey_probe() -> bool {
436    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
437    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
438}
439
440/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
441/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
442/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
443/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
444/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
445/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
446/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
447/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
448/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
449pub trait SpecConstraint {
450    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
451    /// masked argmax).
452    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
453    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
454    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
455    /// Is `tok` consumable in the CURRENT state?
456    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
457    /// Advance the state with an emitted token.
458    fn consume(&mut self, tok: u32) -> Result<(), String>;
459
460    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
461    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
462    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
463    // loose, research/constrained-full-20260803). These three methods let the engine mask the
464    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
465    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
466    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
467    // stays the correctness backstop and the emitted stream is unchanged by construction
468    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
469    // argmax; a cut slot is recomputed as the masked argmax either way).
470    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
471
472    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
473    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
474    fn draft_mask_enabled(&self) -> bool {
475        false
476    }
477    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
478    /// slot. Called once per spec round, before the first draft position.
479    fn draft_begin(&mut self) -> Result<(), String> {
480        Ok(())
481    }
482    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
483    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
484    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
485        Ok(None)
486    }
487    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
488    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
489    /// engine stops drafting; the token already pushed still goes through verify.
490    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
491        Ok(false)
492    }
493}
494
495/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
496/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
497/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
498/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
499/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
500/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
501/// verify emits the masked argmax as usual).
502fn upload_draft_mask(
503    e: &Engine,
504    c: &mut dyn SpecConstraint,
505    dst: &mut CudaSlice<u32>,
506    d2t: Option<&Vec<u32>>,
507    d_vocab: usize,
508    words: usize,
509) -> Result<bool, Box<dyn std::error::Error>> {
510    let Some(tw) = c
511        .draft_mask_words()
512        .map_err(|e2| format!("constraint: {e2}"))?
513    else {
514        return Ok(false);
515    };
516    let bit = |t: usize| -> bool {
517        let w = t >> 5;
518        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
519    };
520    let mut buf = vec![0u32; words];
521    match d2t {
522        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
523        Some(map) => {
524            for (i, &t) in map.iter().enumerate().take(d_vocab) {
525                if bit(t as usize) {
526                    buf[i >> 5] |= 1u32 << (i & 31);
527                }
528            }
529        }
530        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
531        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
532        None => {
533            let n = tw.len().min(words);
534            buf[..n].copy_from_slice(&tw[..n]);
535        }
536    }
537    if buf.iter().all(|w| *w == 0) {
538        return Ok(false);
539    }
540    e.htod_u32_into(dst, &buf)?;
541    Ok(true)
542}
543
544/// Keep the full token-embedding table in host memory and upload only the rows needed by each
545/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
546/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
547/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
548pub(crate) fn spec_host_embd() -> bool {
549    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
550    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
551}
552
553/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
554/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
555/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
556/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
557/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
558/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
559/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
560/// run-spec K=1..8 + acceptance identity arbitrate e2e).
561pub(crate) fn spec_fused_t() -> bool {
562    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
563    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
564    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
565    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
566    *F.get_or_init(|| {
567        std::env::var("MEMRA_SPEC_FUSED_T")
568            .map(|v| v != "0")
569            .unwrap_or(true)
570    })
571}
572
573/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
574/// Only call this on such buffers — the lean contract is "identical bytes by construction".
575fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
576    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
577}
578
579/// Scratch KV for the MTP block (one full-attn layer).
580///
581/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
582/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
583/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
584/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
585/// engine's "mtp_update" design). Entries come from two sources:
586///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
587///     hidden chain-approximate — the reference engine accepts the same);
588///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
589///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
590/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
591/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
592/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
593/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
594/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
595/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
596/// committed row across turns (the predecessor-pairing seed + fill anchor).
597/// Per-request sampling config for the sampled-spec serve path.
598#[derive(Clone, Copy, Debug)]
599pub struct SpecSampling {
600    pub temp: f32,
601    pub seed: u64,
602    pub top_k: i32,            // 0 = off
603    pub top_p: f32,            // 1.0 = off
604    pub min_p: f32,            // 0.0 = off
605    pub penalty_last_n: usize, // 0 = penalties off
606    pub penalty_repeat: f32,
607    pub penalty_freq: f32,
608    pub penalty_present: f32,
609}
610
611impl SpecSampling {
612    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
613    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
614    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
615    /// key their penalty arms off this.
616    pub fn pen_on(&self) -> bool {
617        self.penalty_last_n > 0
618            && (self.penalty_repeat != 1.0
619                || self.penalty_freq != 0.0
620                || self.penalty_present != 0.0)
621    }
622}
623
624/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
625/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
626/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
627/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
628/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
629/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
630/// is a distributional bug, not a style problem).
631pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
632    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
633    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
634    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
635    for _ in 0..10 {
636        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
637        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
638        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
639        c0 = n0;
640        c1 = n1;
641        c2 = n2;
642        c3 = n3;
643        k0 = k0.wrapping_add(0x9E3779B9);
644        k1 = k1.wrapping_add(0xBB67AE85);
645    }
646    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
647}
648
649/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
650/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
651pub const SPEC_TELEM_POS: usize = 8;
652
653/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
654/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
655/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
656/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
657/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
658/// in NEITHER drafted nor accepted.
659#[derive(Clone, Copy, Default, Debug)]
660pub struct SpecTelemetry {
661    /// verify rounds completed (a round-stream burst counts each of its M rounds).
662    pub rounds: u64,
663    /// tokens drafted / accepted across all rounds.
664    pub drafted: u64,
665    pub accepted: u64,
666    /// how often draft position j (0-based within a round's chain) was offered / accepted.
667    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
668    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
669    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
670    pub pos_drafted: [u64; SPEC_TELEM_POS],
671    pub pos_accepted: [u64; SPEC_TELEM_POS],
672}
673
674impl SpecTelemetry {
675    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
676    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
677    /// a wrapped counter.
678    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
679        let mut d = SpecTelemetry {
680            rounds: self.rounds.saturating_sub(prev.rounds),
681            drafted: self.drafted.saturating_sub(prev.drafted),
682            accepted: self.accepted.saturating_sub(prev.accepted),
683            ..Default::default()
684        };
685        for j in 0..SPEC_TELEM_POS {
686            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
687            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
688        }
689        d
690    }
691    /// Fieldwise `self += d` — the worker's per-model aggregation.
692    pub fn merge(&mut self, d: &SpecTelemetry) {
693        self.rounds += d.rounds;
694        self.drafted += d.drafted;
695        self.accepted += d.accepted;
696        for j in 0..SPEC_TELEM_POS {
697            self.pos_drafted[j] += d.pos_drafted[j];
698            self.pos_accepted[j] += d.pos_accepted[j];
699        }
700    }
701
702    /// Mean accepted draft-prefix length per verify round (tau).
703    pub fn tau(&self) -> f64 {
704        if self.rounds > 0 {
705            self.accepted as f64 / self.rounds as f64
706        } else {
707            0.0
708        }
709    }
710}
711
712/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
713/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
714/// launch, synchronization, allocation, or ordering dependency to the numeric path.
715struct SpecTelemetryCounters {
716    rounds: AtomicU64,
717    drafted: AtomicU64,
718    accepted: AtomicU64,
719    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
720    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
721}
722
723impl Default for SpecTelemetryCounters {
724    fn default() -> Self {
725        Self {
726            rounds: AtomicU64::new(0),
727            drafted: AtomicU64::new(0),
728            accepted: AtomicU64::new(0),
729            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
730            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
731        }
732    }
733}
734
735impl SpecTelemetryCounters {
736    fn record_round(&self, drafted: usize, accepted: usize) {
737        debug_assert!(accepted <= drafted);
738        self.rounds.fetch_add(1, Ordering::Relaxed);
739        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
740        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
741        for counter in self.pos_drafted.iter().take(drafted) {
742            counter.fetch_add(1, Ordering::Relaxed);
743        }
744        for counter in self.pos_accepted.iter().take(accepted) {
745            counter.fetch_add(1, Ordering::Relaxed);
746        }
747    }
748
749    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
750    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
751    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
752        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
753        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
754        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
755    }
756
757    fn snapshot(&self) -> SpecTelemetry {
758        SpecTelemetry {
759            rounds: self.rounds.load(Ordering::Relaxed),
760            drafted: self.drafted.load(Ordering::Relaxed),
761            accepted: self.accepted.load(Ordering::Relaxed),
762            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
763            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
764        }
765    }
766}
767
768pub struct SpecSession {
769    pub(crate) cache: Cache,
770    pub(crate) scratch: MtpScratch,
771    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
772    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
773    /// session must count them. Callers render output from this, not from their own echo.
774    pub committed: Vec<u32>,
775    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
776    pub(crate) last_h: Option<CudaSlice<f32>>,
777    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
778    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
779    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
780    pub next_pred: Option<u32>,
781    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
782    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
783    pub sctr: u32,
784    pub uctr: u32,
785    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
786    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
787    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
788    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
789    /// research/spec-serving-20260801). None before the first turn; error paths drop it
790    /// (next burst recaptures — serve retires errored sessions anyway).
791    pub(crate) draft_ctx: Option<DraftGraphCtx>,
792    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
793    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
794    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
795    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
796    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
797    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
798    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
799    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
800    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
801    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
802    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
803    pub pending_tok: Option<u32>,
804    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
805    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
806    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
807    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
808    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
809    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
810    /// accounting the loop already does — no syncs, no allocation. NOTE a
811    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
812    /// diff with [`SpecTelemetry::delta_since`] around each burst.
813    telem: SpecTelemetryCounters,
814    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
815    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
816    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
817    /// prime, result lands in `boundary_captures`.
818    pub capture_at: Option<usize>,
819    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
820    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
821    /// publication just isn't available for that request. Plural since
822    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
823    /// split (the shared-prefix class) and the stable pre-generation boundary (the
824    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
825    /// prefill tick publishes/checkpoints.
826    pub boundary_captures: Vec<SpecBoundaryCapture>,
827    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
828    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
829    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
830    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
831    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
832    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
833    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
834    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
835    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
836    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
837    /// prompt-end capture.
838    pub ckpt_at: Option<usize>,
839}
840impl SpecSession {
841    /// Context capacity of the session's caches (the server's ContextFull guard).
842    pub fn cache_max_ctx(&self) -> usize {
843        self.cache.max_ctx
844    }
845    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
846    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
847    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
848    /// the prime boundary), so no copy was taken at prime time.
849    pub fn cache_ref(&self) -> &Cache {
850        &self.cache
851    }
852    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
853    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
854    /// like the trunk KV — draft rows below the prompt end are append-only for the
855    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
856    /// committed length, never below the prime boundary, and the true-hidden refresh
857    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
858    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
859    /// prefix-addressable; the prefix cache already refuses that class end to end).
860    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
861        if self.scratch.kv.ring.is_some() {
862            return None;
863        }
864        Some((
865            &self.scratch.kv.k,
866            &self.scratch.kv.v,
867            self.scratch.kv.k_tok_bytes,
868            self.scratch.kv.v_tok_bytes,
869        ))
870    }
871    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
872    pub fn telemetry(&self) -> SpecTelemetry {
873        self.telem.snapshot()
874    }
875    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
876    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
877    /// `spec_rewind_to_checkpoint`.
878    pub fn rewind_pos(&self) -> Option<usize> {
879        self.turn_ckpt.as_ref().map(|c| c.pos)
880    }
881    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
882    pub fn rewind_is_resident(&self) -> bool {
883        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
884            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
885        })
886    }
887    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
888    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
889    /// session has never run a turn and has no prediction to hand over.
890    pub fn demote_ready(&self) -> bool {
891        self.pending_tok.is_none() && self.next_pred.is_some()
892    }
893    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
894    pub fn has_pending(&self) -> bool {
895        self.pending_tok.is_some()
896    }
897    /// Committed row count == cache rows (the session invariant), for the caller's own
898    /// `fed`-length cross-check at a handoff boundary.
899    pub fn committed_len(&self) -> usize {
900        self.committed.len()
901    }
902    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
903    /// cache + next-token prediction to the plain batched-decode path.
904    ///
905    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
906    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
907    /// tokenwise prime of the same `committed` sequence would have left it (that is the
908    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
909    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
910    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
911    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
912    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
913    /// a state indistinguishable from one the batched path produced itself: the batched tick
914    /// emits `next_pred`, feeds it into this same cache, and decodes on.
915    ///
916    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
917    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
918    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
919    /// path would silently skip a token.
920    ///
921    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
922    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
923    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
924    /// would mean an `mtp_kv_fill` over the whole committed history).
925    pub fn into_demoted(self) -> Option<(Cache, u32)> {
926        if self.pending_tok.is_some() {
927            return None;
928        }
929        let np = self.next_pred?;
930        debug_assert_eq!(
931            self.cache.pos,
932            self.committed.len(),
933            "demotion handoff: cache rows != committed tokens"
934        );
935        Some((self.cache, np))
936    }
937    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
938    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
939    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
940    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
941    pub fn reset_graph_fallback_on_resume(&mut self) {
942        if let Some(line) = self
943            .draft_ctx
944            .as_mut()
945            .and_then(|c| c.failed.reset_on_resume())
946        {
947            eprintln!("{line}");
948        }
949    }
950}
951
952/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
953///
954/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
955/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
956/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
957/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
958/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
959/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
960///
961/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
962/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
963/// position index, so it must be a real device COPY — that copy is the entire reason a spec
964/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
965/// below the boundary were written by this turn's fill and are never revisited (the per-round
966/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
967/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
968/// predecessor-pairing anchor the next prime's fill reads for its first row.
969///
970/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
971pub(crate) struct SpecCheckpoint {
972    snap: crate::cache::CacheSnapshot,
973    /// Committed length at the boundary (== cache.pos there, the session invariant).
974    pos: usize,
975    /// Pre-output_norm hidden of row `pos - 1`.
976    last_h: CudaSlice<f32>,
977}
978
979/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
980/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
981/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
982/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
983/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
984/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
985/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
986/// so the worker slices those from the live caches post-burst instead of copying at prime time.
987pub struct SpecBoundaryCapture {
988    pub snap: crate::cache::CacheSnapshot,
989    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
990    pub pos: usize,
991    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
992    pub logits: Vec<f32>,
993    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
994    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
995    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
996    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
997    pub last_h: Vec<f32>,
998}
999
1000/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1001/// spec boundary capture carries for later restored-session fills. Failure is silent
1002/// (`turn_ckpt` convention): the capture publishes without an anchor.
1003fn capture_boundary_hidden(
1004    e: &Engine,
1005    h_rows: &CudaSlice<f32>,
1006    pos: usize,
1007    n_embd: usize,
1008) -> Vec<f32> {
1009    if pos == 0 || h_rows.len() < pos * n_embd {
1010        return Vec::new();
1011    }
1012    let Ok(mut row) = e.uninit(n_embd) else {
1013        return Vec::new();
1014    };
1015    if e.copy_view_into(
1016        &mut row,
1017        0,
1018        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1019        n_embd,
1020    )
1021    .is_err()
1022    {
1023        return Vec::new();
1024    }
1025    e.dtoh(&row).unwrap_or_default()
1026}
1027
1028/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1029/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1030/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1031/// every boundary) without touching greedy, which is byte-unaffected either way.
1032pub fn spec_sampled_boundary_on() -> bool {
1033    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1034    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1035}
1036
1037/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1038/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1039/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1040/// restores the pre-lane posture (each burst restarts the window from its own prompt
1041/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1042/// must keep refusing penalized sampled prefix-cache restores, because the restored
1043/// session's continuation burst is handed no prompt slice at all.
1044pub fn spec_pen_session_on() -> bool {
1045    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1046    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1047}
1048
1049/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1050/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1051/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1052/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1053/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1054/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1055pub fn spec_restore_republish_on() -> bool {
1056    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1057    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1058}
1059
1060/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1061/// the argmax the pre-lane code would have emitted from the same row. This is how the
1062/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1063fn spec_boundary_trace() -> bool {
1064    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1065    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1066}
1067
1068/// llama-parity floor for the penalty window when the request does not ask for a bigger
1069/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = usize::MAX` for any
1070/// non-identity penalty, so this floor only matters to explicit small windows and to the
1071/// CLI env path.
1072const PEN_WINDOW_FLOOR: usize = 64;
1073
1074/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1075/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1076/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1077/// p column, the bonus column). The serve API arms `penalty_last_n = usize::MAX` for ANY
1078/// non-identity penalty — "the whole context", llama's `repeat_last_n = -1` — so an
1079/// uncapped session window would put a 128k-token history through that kernel: ~1.7e10
1080/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1081/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1082/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1083/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1084/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1085/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1086/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1087/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1088/// is a second thing to drift.
1089pub const PEN_WINDOW_MAX: usize = 8192;
1090
1091/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1092/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1093/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1094/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1095/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1096/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1097/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1098/// window through the SAME function (one definition of "the window" across both spec
1099/// routes and the gate binary's trunk-only reference arm).
1100pub fn pen_window_seed(
1101    session_committed: &[u32],
1102    burst_prompt: &[u32],
1103    penalty_last_n: usize,
1104) -> Vec<u32> {
1105    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1106    let take_prompt = burst_prompt.len().min(win);
1107    let take_sess = (win - take_prompt).min(session_committed.len());
1108    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1109    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1110    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1111    hist
1112}
1113
1114/// Draw a BOUNDARY token from the target distribution the request asked for
1115/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1116/// every burst boundary".
1117///
1118/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1119/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1120/// row after the last committed token on a continuation burst; the prefix-cache entry's
1121/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1122/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1123/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1124/// customer asked for a sampled token, so this draws one.
1125///
1126/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1127/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1128/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1129/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1130/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1131/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1132///
1133/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1134/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1135/// stream the accept walk uses — never a second, independently seeded stream (which would be
1136/// a new distributional bug: two streams from one seed correlate wherever their counters
1137/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1138/// to the cold session's own first draw from the same logits row, which is what preserves the
1139/// sampled-hit lane's per-seed hit==cold byte identity.
1140#[allow(clippy::too_many_arguments)]
1141pub fn sample_boundary_token_dev(
1142    e: &Engine,
1143    logits: &CudaSlice<f32>,
1144    n_vocab: usize,
1145    sp: &SpecSampling,
1146    pen_hist: &[u32],
1147    sctr: &mut u32,
1148    site: &str,
1149) -> Result<u32, Box<dyn std::error::Error>> {
1150    debug_assert!(
1151        sp.temp > 0.0,
1152        "boundary sampling is the sampled regime only"
1153    );
1154    // Own copy: penalize_logits mutates in place and the caller's row is live state
1155    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1156    let mut col = e.zeros(n_vocab)?;
1157    e.copy_into(&mut col, 0, logits, n_vocab)?;
1158    let pen_on = sp.penalty_last_n > 0
1159        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1160    if pen_on && !pen_hist.is_empty() {
1161        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1162        let w0 = pen_hist
1163            .len()
1164            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1165        let hist = &pen_hist[w0..];
1166        let hd = e.htod_u32_v(hist)?;
1167        e.penalize_logits(
1168            &mut col,
1169            &hd,
1170            hist.len(),
1171            sp.penalty_repeat,
1172            sp.penalty_freq,
1173            sp.penalty_present,
1174            n_vocab,
1175        )?;
1176    }
1177    let rows0 = e.htod_i32(&[0])?;
1178    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1179    e.filter_stats(
1180        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1181        sp.top_p, sp.min_p,
1182    )?;
1183    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1184    let mut perturb = e.zeros(n_vocab)?;
1185    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1186    *sctr = sctr.wrapping_add(1);
1187    let td = e.argmax_token_device(&perturb, n_vocab)?;
1188    let tok = e.dtoh_u32_one(&td)?;
1189    if spec_boundary_trace() {
1190        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1191        let raw = e.argmax_token_device(logits, n_vocab)?;
1192        let greedy = e.dtoh_u32_one(&raw)?;
1193        eprintln!(
1194            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1195             deviates={} temp={} sctr={}",
1196            (tok != greedy) as u8,
1197            sp.temp,
1198            sctr.wrapping_sub(1),
1199        );
1200    }
1201    Ok(tok)
1202}
1203
1204/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1205/// host `Vec<f32>`).
1206#[allow(clippy::too_many_arguments)]
1207pub fn sample_boundary_token(
1208    e: &Engine,
1209    logits: &[f32],
1210    sp: &SpecSampling,
1211    pen_hist: &[u32],
1212    sctr: &mut u32,
1213    site: &str,
1214) -> Result<u32, Box<dyn std::error::Error>> {
1215    let n_vocab = logits.len();
1216    let d = e.htod(logits)?;
1217    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1218}
1219
1220struct SpecPipeTraceClock {
1221    pair: usize,
1222    started: std::time::Instant,
1223}
1224
1225#[derive(Clone)]
1226struct SpecPipeTraceCtx {
1227    clock: std::sync::Arc<SpecPipeTraceClock>,
1228    round: usize,
1229    lane: usize,
1230}
1231
1232struct SpecPipeTraceMarker {
1233    trace: SpecPipeTraceCtx,
1234    phase: &'static str,
1235    edge: &'static str,
1236    slot: Option<usize>,
1237}
1238
1239unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1240    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1241    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1242    let slot = marker
1243        .slot
1244        .map(|v| v.to_string())
1245        .unwrap_or_else(|| "-".into());
1246    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1247    use std::io::Write as _;
1248    let stderr = std::io::stderr();
1249    let mut stderr = stderr.lock();
1250    let _ = writeln!(
1251        stderr,
1252        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1253         slot={slot} t_ms={t_ms:.3}",
1254        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1255    );
1256}
1257
1258fn enqueue_spec_pipe_trace_marker(
1259    stream: &cudarc::driver::CudaStream,
1260    trace: Option<&SpecPipeTraceCtx>,
1261    phase: &'static str,
1262    edge: &'static str,
1263    slot: Option<usize>,
1264) -> Result<(), Box<dyn std::error::Error>> {
1265    let Some(trace) = trace else {
1266        return Ok(());
1267    };
1268    let marker = Box::new(SpecPipeTraceMarker {
1269        trace: trace.clone(),
1270        phase,
1271        edge,
1272        slot,
1273    });
1274    let raw = Box::into_raw(marker);
1275    let result = unsafe {
1276        cudarc::driver::result::stream::launch_host_function(
1277            stream.cu_stream(),
1278            spec_pipe_trace_marker,
1279            raw.cast(),
1280        )
1281    };
1282    if let Err(err) = result {
1283        unsafe {
1284            drop(Box::from_raw(raw));
1285        }
1286        return Err(err.into());
1287    }
1288    Ok(())
1289}
1290
1291#[derive(Default)]
1292struct SpecPipeProgress {
1293    setup_done: [bool; 2],
1294    draft_done: [usize; 2],
1295    stage0_done: [usize; 2],
1296    verify_done: [usize; 2],
1297    accept_done: [usize; 2],
1298    finished: [bool; 2],
1299    aborted: bool,
1300}
1301
1302/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1303/// keeps its existing call stack and round locals; this object only orders phase entry. The
1304/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1305/// cannot be interleaved by the two host threads.
1306struct SpecPipeSync {
1307    progress: std::sync::Mutex<SpecPipeProgress>,
1308    changed: std::sync::Condvar,
1309    primary: std::sync::Mutex<()>,
1310    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1311}
1312
1313impl SpecPipeSync {
1314    fn new() -> Self {
1315        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1316        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1317            std::sync::Arc::new(SpecPipeTraceClock {
1318                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1319                started: std::time::Instant::now(),
1320            })
1321        });
1322        Self {
1323            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1324            changed: std::sync::Condvar::new(),
1325            primary: std::sync::Mutex::new(()),
1326            trace,
1327        }
1328    }
1329}
1330
1331#[derive(Clone)]
1332struct SpecPipeLane {
1333    sync: std::sync::Arc<SpecPipeSync>,
1334    lane: usize,
1335}
1336
1337impl SpecPipeLane {
1338    fn peer(&self) -> usize {
1339        1 - self.lane
1340    }
1341
1342    fn aborted() -> Box<dyn std::error::Error> {
1343        "paired speculative peer aborted".into()
1344    }
1345
1346    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1347        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1348            clock: clock.clone(),
1349            round,
1350            lane: self.lane,
1351        })
1352    }
1353
1354    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1355        let mut p = self.sync.progress.lock().unwrap();
1356        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1357            p = self.sync.changed.wait(p).unwrap();
1358        }
1359        if p.aborted {
1360            Err(Self::aborted())
1361        } else {
1362            Ok(())
1363        }
1364    }
1365
1366    fn setup_end(&self) {
1367        let mut p = self.sync.progress.lock().unwrap();
1368        p.setup_done[self.lane] = true;
1369        self.sync.changed.notify_all();
1370    }
1371
1372    fn draft_begin(
1373        &self,
1374        round: usize,
1375    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1376        let peer = self.peer();
1377        let mut p = self.sync.progress.lock().unwrap();
1378        loop {
1379            if p.aborted {
1380                return Err(Self::aborted());
1381            }
1382            let setup_ready =
1383                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1384            let prior_ready = p.accept_done[self.lane] >= round
1385                && (p.accept_done[peer] >= round || p.finished[peer]);
1386            let turn_ready = if self.lane == 0 {
1387                true
1388            } else {
1389                p.draft_done[0] > round || p.finished[0]
1390            };
1391            if setup_ready && prior_ready && turn_ready {
1392                break;
1393            }
1394            p = self.sync.changed.wait(p).unwrap();
1395        }
1396        drop(p);
1397        Ok(self.sync.primary.lock().unwrap())
1398    }
1399
1400    fn draft_end(&self, round: usize) {
1401        let mut p = self.sync.progress.lock().unwrap();
1402        p.draft_done[self.lane] = round + 1;
1403        self.sync.changed.notify_all();
1404    }
1405
1406    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1407    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1408    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1409        let peer = self.peer();
1410        let mut p = self.sync.progress.lock().unwrap();
1411        loop {
1412            if p.aborted {
1413                return Err(Self::aborted());
1414            }
1415            let ready = if self.lane == 0 {
1416                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1417            } else {
1418                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1419            };
1420            if ready {
1421                return Ok(self.lane == 0 || p.finished[peer]);
1422            }
1423            p = self.sync.changed.wait(p).unwrap();
1424        }
1425    }
1426
1427    fn stage0_end(&self, round: usize) {
1428        let mut p = self.sync.progress.lock().unwrap();
1429        p.stage0_done[self.lane] = round + 1;
1430        self.sync.changed.notify_all();
1431    }
1432
1433    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1434    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1435    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1436        let mut p = self.sync.progress.lock().unwrap();
1437        while !p.aborted
1438            && !(p.stage0_done[self.lane] > round
1439                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1440        {
1441            p = self.sync.changed.wait(p).unwrap();
1442        }
1443        if p.aborted {
1444            Err(Self::aborted())
1445        } else {
1446            Ok(())
1447        }
1448    }
1449
1450    fn verify_end(&self, round: usize) {
1451        let mut p = self.sync.progress.lock().unwrap();
1452        p.verify_done[self.lane] = round + 1;
1453        self.sync.changed.notify_all();
1454    }
1455
1456    fn accept_begin(
1457        &self,
1458        round: usize,
1459    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1460        let mut p = self.sync.progress.lock().unwrap();
1461        loop {
1462            if p.aborted {
1463                return Err(Self::aborted());
1464            }
1465            let ready = if self.lane == 0 {
1466                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1467            } else {
1468                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1469            };
1470            if ready {
1471                break;
1472            }
1473            p = self.sync.changed.wait(p).unwrap();
1474        }
1475        drop(p);
1476        Ok(self.sync.primary.lock().unwrap())
1477    }
1478
1479    fn accept_end(&self, round: usize) {
1480        let mut p = self.sync.progress.lock().unwrap();
1481        p.accept_done[self.lane] = round + 1;
1482        self.sync.changed.notify_all();
1483    }
1484
1485    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1486        self.sync.primary.lock().unwrap()
1487    }
1488
1489    fn finish(&self, failed: bool) {
1490        let mut p = self.sync.progress.lock().unwrap();
1491        p.finished[self.lane] = true;
1492        p.aborted |= failed;
1493        self.sync.changed.notify_all();
1494    }
1495}
1496
1497struct SpecPipeFinish<'a> {
1498    lane: &'a SpecPipeLane,
1499    closed: bool,
1500}
1501
1502impl<'a> SpecPipeFinish<'a> {
1503    fn new(lane: &'a SpecPipeLane) -> Self {
1504        Self {
1505            lane,
1506            closed: false,
1507        }
1508    }
1509
1510    fn close(&mut self, failed: bool) {
1511        self.lane.finish(failed);
1512        self.closed = true;
1513    }
1514}
1515
1516impl Drop for SpecPipeFinish<'_> {
1517    fn drop(&mut self) {
1518        if !self.closed {
1519            self.lane.finish(true);
1520        }
1521    }
1522}
1523
1524/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1525/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1526/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1527/// binds that context before touching the session, joins before returning, and never aliases the
1528/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1529/// session type Send.
1530struct SpecPipeSessionPtr(*mut SpecSession);
1531
1532unsafe impl Send for SpecPipeSessionPtr {}
1533
1534impl SpecPipeSessionPtr {
1535    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1536        unsafe { &mut *self.0 }
1537    }
1538}
1539
1540/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1541/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1542/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1543/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1544/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1545/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1546/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1547/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1548/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1549///
1550/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1551/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1552/// load-bearing:
1553///
1554/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1555///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1556///   This is all the key used to carry.
1557/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1558///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1559///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1560///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1561///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1562///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1563///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1564///
1565/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1566/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1567/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1568/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1569/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1570#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1571pub(crate) struct SampledGraphKey {
1572    seed: u64,
1573    temp_bits: u32,
1574    k: usize,
1575    top_k: i32,
1576    top_p_bits: u32,
1577    min_p_bits: u32,
1578    pen_on: bool,
1579}
1580
1581impl SampledGraphKey {
1582    pub(crate) fn new(
1583        seed: u64,
1584        temp: f32,
1585        k: usize,
1586        top_k: i32,
1587        top_p: f32,
1588        min_p: f32,
1589        pen_on: bool,
1590    ) -> Self {
1591        SampledGraphKey {
1592            seed,
1593            temp_bits: temp.to_bits(),
1594            k,
1595            top_k,
1596            top_p_bits: top_p.to_bits(),
1597            min_p_bits: min_p.to_bits(),
1598            pen_on,
1599        }
1600    }
1601
1602    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1603    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1604    /// the key can never drift apart (they were three separate expressions before this lane, and
1605    /// the launch site simply forgot to ask).
1606    pub(crate) fn pure_temp(&self) -> bool {
1607        self.top_k == 0
1608            && f32::from_bits(self.top_p_bits) >= 1.0
1609            && f32::from_bits(self.min_p_bits) <= 0.0
1610            && !self.pen_on
1611    }
1612}
1613
1614pub(crate) struct DraftGraphCtx {
1615    g_tok: CudaSlice<u32>,
1616    g_pos: CudaSlice<i32>,
1617    g_seed: CudaSlice<f32>,
1618    g_p: CudaSlice<f32>,
1619    g_ctr: CudaSlice<u32>,
1620    g_q: CudaSlice<f32>,
1621    g_perturb: CudaSlice<f32>,
1622    q_slots: Vec<CudaSlice<f32>>,
1623    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1624    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1625    /// per-position contents the host re-uploads before each replay (the graph-promote
1626    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1627    g_dmask: CudaSlice<u32>,
1628    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1629    graph_masked: bool,
1630    graph: Option<cudarc::driver::CudaGraph>,
1631    graph_s: Option<cudarc::driver::CudaGraph>,
1632    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1633    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1634    failed: DraftGraphFallback,
1635    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1636    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1637    s_key: Option<SampledGraphKey>,
1638    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1639    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1640    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1641    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1642    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1643    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1644    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1645    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1646    keeper: Vec<Box<dyn std::any::Any + Send>>,
1647    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1648}
1649
1650/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1651/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1652///
1653/// Three contracts:
1654/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1655///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1656///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1657///   an already-failed graph returns None (the per-burst memoization that keeps the eager
1658///   fallback from paying a doomed capture attempt every burst).
1659/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
1660///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
1661///   failure for the pool's whole lifetime. Returns the note line only when a flag was
1662///   actually set (quiet on the common clean-resume path).
1663/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
1664///   capture attempt whose own failure would re-flip loudly.
1665#[derive(Default)]
1666pub(crate) struct DraftGraphFallback {
1667    greedy: bool,
1668    sampled: bool,
1669}
1670impl DraftGraphFallback {
1671    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
1672        if self.greedy {
1673            return None;
1674        }
1675        self.greedy = true;
1676        Some(format!(
1677            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
1678        ))
1679    }
1680    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
1681        if self.sampled {
1682            return None;
1683        }
1684        self.sampled = true;
1685        Some(format!(
1686            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
1687        ))
1688    }
1689    fn greedy_failed(&self) -> bool {
1690        self.greedy
1691    }
1692    fn sampled_failed(&self) -> bool {
1693        self.sampled
1694    }
1695    fn clear_greedy(&mut self) {
1696        self.greedy = false;
1697    }
1698    fn clear_sampled(&mut self) {
1699        self.sampled = false;
1700    }
1701    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
1702    /// was set (so clean resumes stay quiet).
1703    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
1704        if !self.greedy && !self.sampled {
1705            return None;
1706        }
1707        let which = match (self.greedy, self.sampled) {
1708            (true, true) => "greedy+sampled",
1709            (true, false) => "greedy",
1710            _ => "sampled",
1711        };
1712        self.greedy = false;
1713        self.sampled = false;
1714        Some(format!(
1715            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
1716        ))
1717    }
1718}
1719
1720impl DraftGraphCtx {
1721    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
1722        Ok(DraftGraphCtx {
1723            g_tok: e.alloc_u32_zeroed(1)?,
1724            g_pos: e.htod_i32(&[0])?,
1725            g_seed: e.zeros(n_embd)?,
1726            g_p: e.zeros(1)?,
1727            g_ctr: e.alloc_u32_zeroed(1)?,
1728            g_q: e.zeros(qlen)?,
1729            g_perturb: e.zeros(qlen)?,
1730            q_slots: Vec::new(),
1731            g_dmask: e.alloc_u32_zeroed(1)?,
1732            graph_masked: false,
1733            graph: None,
1734            graph_s: None,
1735            failed: DraftGraphFallback::default(),
1736            s_key: None,
1737            keeper: Vec::new(),
1738            keeper_s: Vec::new(),
1739        })
1740    }
1741}
1742
1743pub(crate) struct MtpScratch {
1744    kv: KvLayer,
1745    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
1746    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
1747    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
1748    /// smaller host-indexed SWA ring instead.
1749    cap: usize,
1750    extra: Vec<MtpScratchPlane>,
1751}
1752
1753struct MtpScratchPlane {
1754    kv: KvLayer,
1755    cap: usize,
1756}
1757
1758fn mtp_scratch_layout(
1759    cfg: &memra_gguf::config::ModelConfig,
1760    geom: Option<&crate::hybrid::DraftGeom>,
1761) -> (usize, usize, usize, usize) {
1762    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
1763    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
1764    let head_dim_k = cfg.head_dim_k as usize;
1765    let head_dim_v = cfg.head_dim_v as usize;
1766    assert!(
1767        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
1768        "KVQUANT requires head_dim%32==0 (MTP scratch)"
1769    );
1770    let kv_dim_k = head_dim_k * n_head_kv;
1771    let kv_dim_v = head_dim_v * n_head_kv;
1772    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
1773    // policy shared with `MtpScratch::new` so admission scales the same allocation.
1774    let (kbb, vbb) = crate::kv_blk_bytes();
1775    let k_tok_bytes = (kv_dim_k / 32) * kbb;
1776    let v_tok_bytes = (kv_dim_v / 32) * vbb;
1777    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
1778}
1779
1780fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
1781    assert!(head_count > 0, "MTP chain requires at least one head");
1782    step % head_count
1783}
1784
1785impl MtpScratch {
1786    fn alloc_plane(
1787        e: &Engine,
1788        cfg: &memra_gguf::config::ModelConfig,
1789        plan: &memra_gguf::model_plan::ModelPlan,
1790        cap: usize,
1791        geom: Option<&crate::hybrid::DraftGeom>,
1792    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
1793        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
1794        let ring = if crate::cache::swa_ring_on()
1795            && crate::plan_backend::decode_batch_program(plan)
1796                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
1797        {
1798            let window = plan
1799                .layers
1800                .iter()
1801                .find_map(|layer| match layer.attention {
1802                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
1803                        Some(window as usize)
1804                    }
1805                    _ => None,
1806                })
1807                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
1808            Some(crate::cache::KvRing::new(
1809                crate::cache::swa_ring_rows(window, cap),
1810                window,
1811            ))
1812        } else {
1813            None
1814        };
1815        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
1816        Ok(MtpScratchPlane {
1817            kv: KvLayer {
1818                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
1819                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
1820                kv_dim_k,
1821                kv_dim_v,
1822                k_tok_bytes,
1823                v_tok_bytes,
1824                len: 0,
1825                ring,
1826                len_d: e.htod_i32(&[0])?,
1827            },
1828            cap,
1829        })
1830    }
1831
1832    fn new(
1833        e: &Engine,
1834        cfg: &memra_gguf::config::ModelConfig,
1835        plan: &memra_gguf::model_plan::ModelPlan,
1836        cap: usize,
1837        geom: Option<&crate::hybrid::DraftGeom>,
1838    ) -> Result<Self, Box<dyn std::error::Error>> {
1839        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
1840        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
1841        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
1842        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
1843        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
1844        Ok(MtpScratch {
1845            kv: primary.kv,
1846            cap: primary.cap,
1847            extra: Vec::new(),
1848        })
1849    }
1850
1851    fn push_plane(
1852        &mut self,
1853        e: &Engine,
1854        cfg: &memra_gguf::config::ModelConfig,
1855        plan: &memra_gguf::model_plan::ModelPlan,
1856        geom: Option<&crate::hybrid::DraftGeom>,
1857    ) -> Result<(), Box<dyn std::error::Error>> {
1858        self.extra
1859            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
1860        Ok(())
1861    }
1862
1863    fn plane_count(&self) -> usize {
1864        1 + self.extra.len()
1865    }
1866
1867    fn plane(&self, index: usize) -> (&KvLayer, usize) {
1868        if index == 0 {
1869            (&self.kv, self.cap)
1870        } else {
1871            let plane = &self.extra[index - 1];
1872            (&plane.kv, plane.cap)
1873        }
1874    }
1875
1876    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
1877        if index == 0 {
1878            (&mut self.kv, self.cap)
1879        } else {
1880            let plane = &mut self.extra[index - 1];
1881            (&mut plane.kv, plane.cap)
1882        }
1883    }
1884
1885    fn set_plane_len(
1886        &mut self,
1887        e: &Engine,
1888        index: usize,
1889        n: usize,
1890    ) -> Result<(), Box<dyn std::error::Error>> {
1891        let (kv, _) = self.plane_mut(index);
1892        if kv.ring.as_ref().is_some_and(|ring| !ring.can_rewind_to(n)) {
1893            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1894        }
1895        kv.len = n;
1896        e.set_i32_one(&mut kv.len_d, n as i32)
1897    }
1898
1899    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
1900    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
1901    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
1902    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
1903        if !self.can_rewind_to(n) {
1904            return Err("SWA ring MTP checkpoint has been lapped; full re-prime required".into());
1905        }
1906        for index in 0..self.plane_count() {
1907            self.set_plane_len(e, index, n)?;
1908        }
1909        Ok(())
1910    }
1911
1912    fn can_rewind_to(&self, n: usize) -> bool {
1913        (0..self.plane_count()).all(|index| {
1914            self.plane(index)
1915                .0
1916                .ring
1917                .as_ref()
1918                .is_none_or(|ring| ring.can_rewind_to(n))
1919        })
1920    }
1921}
1922
1923/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
1924/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
1925/// full weight reads per round — recomputing columns the verify had already produced
1926/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
1927/// to "after the first j verify columns" WITHOUT re-running the trunk:
1928/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
1929///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
1930///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
1931///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
1932///   pure-copy ring rebuild.
1933/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
1934///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
1935///   target: j <= t-1).
1936/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
1937/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
1938struct GdnStash {
1939    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
1940    q_l2: CudaSlice<f32>,
1941    k_l2: CudaSlice<f32>,
1942    v_g: CudaSlice<f32>, // [t, num_v, d_state]
1943    g_log: CudaSlice<f32>,
1944    beta: CudaSlice<f32>, // [t, num_v]
1945}
1946pub(crate) struct VerifyCkpt {
1947    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
1948    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
1949}
1950/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
1951pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
1952
1953/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
1954/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
1955/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
1956/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
1957/// layers between full-attention layers are shape-static given vt — no positions, no
1958/// t_kv, state addressed through pointer tables — so runs of them capture per
1959/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
1960/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
1961///
1962/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
1963/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
1964/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
1965/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
1966/// before and restored after — the graph's first real launch starts from the exact
1967/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
1968/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
1969/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
1970pub(crate) struct DsparkVerifyGraphs {
1971    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
1972    lin: Vec<usize>,
1973    lin_pos: std::collections::HashMap<usize, usize>,
1974    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
1975    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
1976    table_all: CudaSlice<u64>,
1977    host_table: Vec<u64>,
1978    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
1979    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
1980    stash_conv: Vec<CudaSlice<f32>>,
1981    stash_ssm: Vec<CudaSlice<f32>>,
1982    conv_words: usize,
1983    ssm_words: usize,
1984    /// Per-vt input/output staging (stable addresses the graphs bake).
1985    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
1986    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
1987    /// so the sink buffer must live (and persist) with the graphs, not with the round.
1988    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
1989    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
1990    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
1991    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
1992    save_conv: CudaSlice<f32>,
1993    save_ssm: CudaSlice<f32>,
1994    max_run: usize,
1995    n_embd: usize,
1996    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
1997    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
1998    pub(crate) round_slab: bool,
1999    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2000    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2001    fa: Vec<usize>,
2002    fa_pos: std::collections::HashMap<usize, usize>,
2003    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2004    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2005    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2006    fa_table: CudaSlice<u64>,
2007    fa_host_table: Vec<u64>,
2008    t_cap: usize,
2009    /// Per-vt position staging for the captured bodies — contents refreshed per round
2010    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2011    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2012    /// Full-verify graphs keyed (vt, rung_end, hi).
2013    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2014    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2015    covered: usize,
2016    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2017    /// full-verify capture walks all of them.
2018    walk_uniform: bool,
2019    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2020    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2021    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2022    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2023    debt_obs: Option<(usize, usize)>,
2024}
2025
2026struct DsparkSegGraph {
2027    graph: cudarc::driver::CudaGraph,
2028    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2029}
2030
2031/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2032/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2033/// modes without a second copy of the math.
2034pub(crate) struct FaLayerArgs<'a> {
2035    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2036    /// them per-z (append slot = pos, T_kv = pos + 1).
2037    pub pos_d: &'a CudaSlice<i32>,
2038    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2039    /// arm builds/uses them (graph mode refuses that arm).
2040    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2041    pub pos0: usize,
2042    pub seqs_append: bool,
2043    pub batch_fa_on: bool,
2044    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2045    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2046    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2047    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2048    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2049    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2050    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2051    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2052    /// for FA layers that never touch it.
2053    pub ckpt: Option<&'a mut VerifyCkpt>,
2054}
2055
2056// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2057// no automatic trait; CUDA driver graph handles are context-scoped rather than
2058// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2059// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2060// single decode-stream thread.
2061unsafe impl Send for DsparkVerifyGraphs {}
2062
2063impl DsparkVerifyGraphs {
2064    /// Live capture count (segment + full graphs) — the denominator of
2065    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2066    pub(crate) fn captures(&self) -> usize {
2067        self.graphs.len() + self.full.len()
2068    }
2069
2070    /// Take the marginal-growth debt reading and record this observation for the next one.
2071    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2072    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2073        let captures = self.captures();
2074        let debt =
2075            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2076        if captures > 0 {
2077            match self.debt_obs {
2078                Some((c0, _)) if captures <= c0 => {}
2079                _ => self.debt_obs = Some((captures, reserved_bytes)),
2080            }
2081        }
2082        debt
2083    }
2084
2085    /// Build for this cache's shape. None when there are no linear layers, sizes are
2086    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2087    pub(crate) fn new(
2088        e: &Engine,
2089        cache: &Cache,
2090        t_max: usize,
2091        n_embd: usize,
2092    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2093        let lin: Vec<usize> = (0..cache.recur.len())
2094            .filter(|&il| cache.recur[il].is_some())
2095            .collect();
2096        if lin.is_empty() || t_max < 2 {
2097            return Ok(None);
2098        }
2099        let first = cache.recur[lin[0]].as_ref().unwrap();
2100        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2101        for &il in &lin {
2102            let rl = cache.recur[il].as_ref().unwrap();
2103            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2104                return Ok(None);
2105            }
2106        }
2107        let n = lin.len();
2108        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2109        for (k, &il) in lin.iter().enumerate() {
2110            lin_pos.insert(il, k);
2111        }
2112        // longest run of consecutive linear layers (save-scratch sizing)
2113        let mut max_run = 1usize;
2114        let mut run = 1usize;
2115        for w in lin.windows(2) {
2116            if w[1] == w[0] + 1 {
2117                run += 1;
2118                max_run = max_run.max(run);
2119            } else {
2120                run = 1;
2121            }
2122        }
2123        let rows = t_max - 1;
2124        let mut stash_conv = Vec::with_capacity(n);
2125        let mut stash_ssm = Vec::with_capacity(n);
2126        for _ in 0..n {
2127            stash_conv.push(e.uninit(rows * conv_words)?);
2128            stash_ssm.push(e.uninit(rows * ssm_words)?);
2129        }
2130        let host_table = vec![0u64; n * 6];
2131        let table_all = e.htod_u64(&host_table)?;
2132        // slice 4c: full-attention census for the full-verify graphs.
2133        let fa: Vec<usize> = (0..cache.kv.len())
2134            .filter(|&il| cache.kv[il].is_some())
2135            .collect();
2136        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2137        for (k, &il) in fa.iter().enumerate() {
2138            fa_pos.insert(il, k);
2139        }
2140        let n_layers = cache.kv.len().max(cache.recur.len());
2141        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2142        let walk_uniform = (0..n_layers).all(|il| {
2143            cache.recur.get(il).is_some_and(|r| r.is_some())
2144                != cache.kv.get(il).is_some_and(|k| k.is_some())
2145        });
2146        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2147        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2148        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2149        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2150        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2151        let covered = (0..n_layers)
2152            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2153            .count();
2154        let t_cap = t_max;
2155        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2156        let fa_table = e.htod_u64(&fa_host_table)?;
2157        Ok(Some(Self {
2158            lin,
2159            lin_pos,
2160            table_all,
2161            host_table,
2162            stash_conv,
2163            stash_ssm,
2164            conv_words,
2165            ssm_words,
2166            stage: std::collections::HashMap::new(),
2167            tap_bufs: std::collections::HashMap::new(),
2168            graphs: std::collections::HashMap::new(),
2169            save_conv: e.uninit(n * conv_words)?,
2170            save_ssm: e.uninit(n * ssm_words)?,
2171            max_run,
2172            n_embd,
2173            round_slab: false,
2174            fa,
2175            fa_pos,
2176            fa_table,
2177            fa_host_table,
2178            t_cap,
2179            pos_stage: std::collections::HashMap::new(),
2180            full: std::collections::HashMap::new(),
2181            covered,
2182            walk_uniform,
2183            debt_obs: None,
2184        }))
2185    }
2186
2187    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2188    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2189    /// cache buffers land at new addresses; a stale table would read the wrong state).
2190    pub(crate) fn refresh_tables(
2191        &mut self,
2192        e: &Engine,
2193        cache: &Cache,
2194    ) -> Result<(), Box<dyn std::error::Error>> {
2195        use cudarc::driver::DevicePtr;
2196        {
2197            let s = &e.gpu.stream();
2198            for (k, &il) in self.lin.iter().enumerate() {
2199                let rl = cache.recur[il].as_ref().unwrap();
2200                let (pc, _g0) = rl.conv_state.device_ptr(s);
2201                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2202                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2203                let o = k * 6;
2204                self.host_table[o] = pc as u64;
2205                self.host_table[o + 1] = p0 as u64;
2206                self.host_table[o + 2] = p1 as u64;
2207                self.host_table[o + 3] = pc as u64;
2208                self.host_table[o + 4] = p1 as u64;
2209                self.host_table[o + 5] = p0 as u64;
2210            }
2211            for (k, &il) in self.fa.iter().enumerate() {
2212                let kvl = cache.kv[il].as_ref().unwrap();
2213                let (pk, _g0) = kvl.k.device_ptr(s);
2214                let (pv, _g1) = kvl.v.device_ptr(s);
2215                let o = k * 2 * self.t_cap;
2216                for z in 0..self.t_cap {
2217                    self.fa_host_table[o + 2 * z] = pk as u64;
2218                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2219                }
2220            }
2221        }
2222        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2223        if !self.fa_host_table.is_empty() {
2224            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2225        }
2226        Ok(())
2227    }
2228
2229    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2230    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2231    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2232    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2233    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2234    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2235    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2236    /// captured graph is bit-identical for every round the rung covers.
2237    #[allow(clippy::too_many_arguments)]
2238    pub(crate) fn full_rung(
2239        &self,
2240        model: &crate::hybrid::HybridModel,
2241        cache: &Cache,
2242        lo: usize,
2243        hi: usize,
2244        t: usize,
2245        seqs_arms_on: bool,
2246    ) -> Option<usize> {
2247        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2248            static ONCE: std::sync::Once = std::sync::Once::new();
2249            let len0 = self
2250                .fa
2251                .first()
2252                .and_then(|&il| cache.kv[il].as_ref())
2253                .map(|k| k.len);
2254            ONCE.call_once(|| {
2255                eprintln!(
2256                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2257                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2258                    self.lin.len(), self.fa.len(), self.t_cap, len0
2259                );
2260            });
2261        }
2262        if !self.walk_uniform
2263            || !seqs_arms_on
2264            || !dspark_fa_rows_on()
2265            || t < 2
2266            || lo != 0
2267            || hi > self.covered
2268            || t > self.t_cap
2269            || self.fa.is_empty()
2270        {
2271            return None;
2272        }
2273        let cfg = &model.cfg;
2274        let head_dim_global = cfg.head_dim_k as usize;
2275        let nkv = cfg.n_head_kv as usize;
2276        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2277        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2278        // projection stride (the body's guard, hoisted so ineligible models fall back
2279        // instead of refusing mid-capture).
2280        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2281        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2282        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2283            return None;
2284        }
2285        let len0 = kvl0.len;
2286        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2287        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2288            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2289            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2290        {
2291            return None;
2292        }
2293        let rung = t_kv_last.next_power_of_two().max(256);
2294        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2295            return None;
2296        }
2297        Some(rung)
2298    }
2299
2300    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2301    /// the residual + refresh the per-vt position staging, capture on first encounter
2302    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2303    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2304    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2305    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2306    #[allow(clippy::too_many_arguments)]
2307    pub(crate) fn run_full(
2308        &mut self,
2309        model: &crate::hybrid::HybridModel,
2310        e: &Engine,
2311        lo: usize,
2312        hi: usize,
2313        x: &CudaSlice<f32>,
2314        t: usize,
2315        pos0: usize,
2316        rung: usize,
2317        cache: &mut Cache,
2318    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2319        let n_embd = self.n_embd;
2320        if !self.stage.contains_key(&t) {
2321            let xin = e.uninit(t * n_embd)?;
2322            let xout = e.uninit(t * n_embd)?;
2323            self.stage.insert(t, (xin, xout));
2324        }
2325        if !self.pos_stage.contains_key(&t) {
2326            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2327        }
2328        // Per-round refresh: position contents + input staging (both addresses are baked
2329        // by the captured bodies; only their CONTENTS change round to round).
2330        {
2331            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2332            let pb = self.pos_stage.get_mut(&t).unwrap();
2333            e.htod_i32_into(pb, &pos_host)?;
2334            let (xin, _) = self.stage.get_mut(&t).unwrap();
2335            e.copy_into(xin, 0, x, t * n_embd)?;
2336        }
2337        let key = (t, rung, hi);
2338        if !self.full.contains_key(&key) {
2339            // The warmups EXECUTE the whole walk on live state — save every linear
2340            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2341            // graph mode never bumps host lens and the appends write this round's own
2342            // slots).
2343            for (k, &il) in self.lin.iter().enumerate() {
2344                let rl = cache.recur[il].as_ref().unwrap();
2345                e.copy_into(
2346                    &mut self.save_conv,
2347                    k * self.conv_words,
2348                    &rl.conv_state,
2349                    self.conv_words,
2350                )?;
2351                e.copy_into(
2352                    &mut self.save_ssm,
2353                    k * self.ssm_words,
2354                    &rl.ssm_state,
2355                    self.ssm_words,
2356                )?;
2357            }
2358            let (graph, keeper) = {
2359                let table_all = &self.table_all;
2360                let lin_pos = &self.lin_pos;
2361                let fa_pos = &self.fa_pos;
2362                let fa_table = &self.fa_table;
2363                let t_cap = self.t_cap;
2364                let stash_conv = &mut self.stash_conv;
2365                let stash_ssm = &mut self.stash_ssm;
2366                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2367                let (xin, xout) = self
2368                    .stage
2369                    .get_mut(&t)
2370                    .map(|(a, b)| (&*a, b))
2371                    .expect("stage bucket created above");
2372                let cache_ref: &mut Cache = cache;
2373                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2374                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2375                } else {
2376                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2377                };
2378                e.capture_graph_retained_flags(iflag, move |e| {
2379                    let mut xc: Option<CudaSlice<f32>> = None;
2380                    for il in lo..hi {
2381                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2382                        let nx = if let Some(&k) = lin_pos.get(&il) {
2383                            model.qwen35_tparallel_linear_layer(
2384                                e,
2385                                il,
2386                                xr,
2387                                t,
2388                                cache_ref,
2389                                None,
2390                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2391                                Some((table_all, k * 6)),
2392                            )?
2393                        } else if let Some(&kf) = fa_pos.get(&il) {
2394                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2395                            model.qwen35_tparallel_fa_layer(
2396                                e,
2397                                il,
2398                                xr,
2399                                t,
2400                                cache_ref,
2401                                FaLayerArgs {
2402                                    pos_d,
2403                                    pos_rows: &mut no_rows,
2404                                    pos0,
2405                                    seqs_append: true,
2406                                    batch_fa_on: true,
2407                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2408                                    stream: None,
2409                                    ckpt: None,
2410                                },
2411                            )?
2412                        } else {
2413                            return Err(format!(
2414                                "run_full: layer {il} is neither linear nor full-attention"
2415                            )
2416                            .into());
2417                        };
2418                        xc = Some(nx);
2419                    }
2420                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2421                    Ok(())
2422                })?
2423            };
2424            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2425            // is odd -> 3 runs = net one swap), then restore the device state the
2426            // warmups consumed (walk scope only — layers past hi never executed). The
2427            // launch below then behaves exactly like one run.
2428            if t % 2 == 1 {
2429                for &il in &self.lin {
2430                    if il < lo || il >= hi {
2431                        continue;
2432                    }
2433                    let rl = cache.recur[il].as_mut().unwrap();
2434                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2435                }
2436            }
2437            for (k, &il) in self.lin.iter().enumerate() {
2438                if il < lo || il >= hi {
2439                    continue;
2440                }
2441                let rl = cache.recur[il].as_mut().unwrap();
2442                let (cw, sw) = (self.conv_words, self.ssm_words);
2443                {
2444                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2445                    let win = sv.slice(k * cw..(k + 1) * cw);
2446                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2447                }
2448                {
2449                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2450                    let win = sv.slice(k * sw..(k + 1) * sw);
2451                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2452                }
2453            }
2454            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2455                if let Ok(c) = crate::graph_update::node_census(&graph) {
2456                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2457                }
2458            }
2459            self.full.insert(
2460                key,
2461                DsparkSegGraph {
2462                    graph,
2463                    _keeper: keeper,
2464                },
2465            );
2466        }
2467        self.full[&key].graph.launch()?;
2468        // Host bookkeeping for the replayed body (captured host code does not re-run):
2469        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2470        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2471        // head layer's kv) that the walk never touches.
2472        if t % 2 == 1 {
2473            for &il in &self.lin {
2474                if il < lo || il >= hi {
2475                    continue;
2476                }
2477                let rl = cache.recur[il].as_mut().unwrap();
2478                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2479            }
2480        }
2481        for &il in &self.fa {
2482            if il < lo || il >= hi {
2483                continue;
2484            }
2485            cache.kv[il].as_mut().unwrap().len += t;
2486        }
2487        let (_, xout) = self.stage.get(&t).unwrap();
2488        let mut out = e.uninit(t * n_embd)?;
2489        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2490        Ok(out)
2491    }
2492
2493    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2494    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2495    /// bracketed by a segment state save/restore), launch, then apply the host parity
2496    /// bookkeeping the captured body would have done. Returns the fresh residual.
2497    #[allow(clippy::too_many_arguments)]
2498    fn run_segment(
2499        &mut self,
2500        model: &crate::hybrid::HybridModel,
2501        e: &Engine,
2502        start: usize,
2503        end: usize,
2504        x: &CudaSlice<f32>,
2505        t: usize,
2506        cache: &mut Cache,
2507    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2508        let n_embd = self.n_embd;
2509        debug_assert!(end - start <= self.max_run);
2510        if !self.stage.contains_key(&t) {
2511            let xin = e.uninit(t * n_embd)?;
2512            let xout = e.uninit(t * n_embd)?;
2513            self.stage.insert(t, (xin, xout));
2514        }
2515        // Stage the residual at the bucket's baked input address.
2516        {
2517            let (xin, _) = self.stage.get_mut(&t).unwrap();
2518            e.copy_into(xin, 0, x, t * n_embd)?;
2519        }
2520        let key = (start, t);
2521        if !self.graphs.contains_key(&key) {
2522            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2523            // ssm of every segment layer first, restore after, so the graph's first real
2524            // launch starts from the exact pre-round state (bytes gated e2e).
2525            for (k, il) in (start..end).enumerate() {
2526                let rl = cache.recur[il].as_ref().unwrap();
2527                e.copy_into(
2528                    &mut self.save_conv,
2529                    k * self.conv_words,
2530                    &rl.conv_state,
2531                    self.conv_words,
2532                )?;
2533                e.copy_into(
2534                    &mut self.save_ssm,
2535                    k * self.ssm_words,
2536                    &rl.ssm_state,
2537                    self.ssm_words,
2538                )?;
2539            }
2540            let (graph, keeper) = {
2541                let table_all = &self.table_all;
2542                let lin_pos = &self.lin_pos;
2543                let stash_conv = &mut self.stash_conv;
2544                let stash_ssm = &mut self.stash_ssm;
2545                let (xin, xout) = self
2546                    .stage
2547                    .get_mut(&t)
2548                    .map(|(a, b)| (&*a, b))
2549                    .expect("stage bucket created above");
2550                let cache_ref: &mut Cache = cache;
2551                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2552                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2553                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2554                // = ~0.41 ms/round, most of the eager-launch savings. The captured
2555                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2556                // (every transient drops inside the capture region — the generic
2557                // capture path's census precedent, 1589/1589), so AUTO_FREE has
2558                // nothing to reclaim and the graph is legal to instantiate without
2559                // it; PRIORITY is the flag the gemma slotted door ships for exactly
2560                // this reason (both alternatives drop the scan; UPLOAD via
2561                // cuGraphInstantiateWithFlags is WithParams-only and refused).
2562                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2563                // the node census at capture (the ALLOC==FREE receipt).
2564                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2565                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2566                } else {
2567                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2568                };
2569                e.capture_graph_retained_flags(iflag, move |e| {
2570                    let mut xc: Option<CudaSlice<f32>> = None;
2571                    for il in start..end {
2572                        let k = lin_pos[&il];
2573                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2574                        let nx = model.qwen35_tparallel_linear_layer(
2575                            e,
2576                            il,
2577                            xr,
2578                            t,
2579                            cache_ref,
2580                            None,
2581                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
2582                            Some((table_all, k * 6)),
2583                        )?;
2584                        xc = Some(nx);
2585                    }
2586                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2587                    Ok(())
2588                })?
2589            };
2590            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2591            // is odd -> 3 runs = net one swap), then restore the device state the
2592            // warmups consumed. The launch below then behaves exactly like one run.
2593            if t % 2 == 1 {
2594                for il in start..end {
2595                    let rl = cache.recur[il].as_mut().unwrap();
2596                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2597                }
2598            }
2599            for (k, il) in (start..end).enumerate() {
2600                let rl = cache.recur[il].as_mut().unwrap();
2601                let (cw, sw) = (self.conv_words, self.ssm_words);
2602                {
2603                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2604                    let win = sv.slice(k * cw..(k + 1) * cw);
2605                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2606                }
2607                {
2608                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2609                    let win = sv.slice(k * sw..(k + 1) * sw);
2610                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2611                }
2612            }
2613            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2614                if let Ok(c) = crate::graph_update::node_census(&graph) {
2615                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
2616                }
2617            }
2618            self.graphs.insert(
2619                key,
2620                DsparkSegGraph {
2621                    graph,
2622                    _keeper: keeper,
2623                },
2624            );
2625        }
2626        self.graphs[&key].graph.launch()?;
2627        // Host parity bookkeeping for the replayed body (the captured host swaps do not
2628        // re-run at replay).
2629        if t % 2 == 1 {
2630            for il in start..end {
2631                let rl = cache.recur[il].as_mut().unwrap();
2632                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2633            }
2634        }
2635        let (_, xout) = self.stage.get(&t).unwrap();
2636        let mut out = e.uninit(t * n_embd)?;
2637        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2638        Ok(out)
2639    }
2640
2641    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
2642    fn can_capture(&self) -> bool {
2643        self.graphs.len() + self.full.len() < dspark_vg_cap()
2644    }
2645
2646    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
2647    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
2648    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
2649    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
2650    /// refusal would stash some layers in the ctx slabs and others in the round's cols
2651    /// while one commit reads only one of them.
2652    pub(crate) fn segments_ready(
2653        &self,
2654        model: &crate::hybrid::HybridModel,
2655        lo: usize,
2656        hi: usize,
2657        t: usize,
2658    ) -> bool {
2659        if self.can_capture() {
2660            return true;
2661        }
2662        let mut il = lo;
2663        while il < hi {
2664            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2665                let start = il;
2666                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
2667                    il += 1;
2668                }
2669                if !self.graphs.contains_key(&(start, t)) {
2670                    return false;
2671                }
2672            } else {
2673                il += 1;
2674            }
2675        }
2676        true
2677    }
2678
2679    /// Widest verify window this pool was built for. A caller whose round exceeds it must
2680    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
2681    /// past them is a panic rather than a refusal.
2682    pub(crate) fn t_capacity(&self) -> usize {
2683        self.t_cap
2684    }
2685
2686    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
2687    /// `row` (0-based) of layer `il`. None for non-linear layers.
2688    pub(crate) fn slab_row(
2689        &self,
2690        e: &Engine,
2691        il: usize,
2692        row: usize,
2693    ) -> Option<(u64, u64, usize, usize)> {
2694        use cudarc::driver::DevicePtr;
2695        let k = *self.lin_pos.get(&il)?;
2696        let s = &e.gpu.stream();
2697        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
2698        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
2699        Some((
2700            pc as u64 + (row * self.conv_words * 4) as u64,
2701            ps as u64 + (row * self.ssm_words * 4) as u64,
2702            self.conv_words,
2703            self.ssm_words,
2704        ))
2705    }
2706}
2707
2708impl VerifyCkpt {
2709    fn new(n_layer: usize) -> Self {
2710        VerifyCkpt {
2711            gdn: (0..n_layer).map(|_| None).collect(),
2712            cols: (0..n_layer).map(|_| None).collect(),
2713        }
2714    }
2715}
2716
2717/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
2718/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
2719/// a logical round number.
2720struct VerifyBoundaryTicket {
2721    rt: &'static crate::pp::PpNRt,
2722    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
2723    slot: usize,
2724    pos0: usize,
2725    t: usize,
2726    payload: usize,
2727    n_st: usize,
2728    pipelined: bool,
2729    pp_anatomy: bool,
2730    pp_started: std::time::Instant,
2731    reverse_ms: f64,
2732    stage0_ms: f64,
2733    tx_ms: f64,
2734    trace: Option<SpecPipeTraceCtx>,
2735}
2736
2737/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
2738/// increment-2 controller can also be armed by the server's fresh-process research door.
2739#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2740pub enum OptiForkGateMode {
2741    Disabled,
2742    Hit,
2743    Miss,
2744    Alternate,
2745    Abort,
2746    Controller,
2747}
2748
2749static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
2750static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
2751    std::sync::atomic::AtomicU32::new(0);
2752static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2753static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2754static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2755static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2756static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2757static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2758static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2759static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2760static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2761static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2762    std::sync::atomic::AtomicU64::new(0);
2763static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
2764    std::sync::atomic::AtomicU64::new(0);
2765static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2766
2767impl OptiForkGateMode {
2768    fn code(self) -> u8 {
2769        match self {
2770            Self::Disabled => 0,
2771            Self::Hit => 1,
2772            Self::Miss => 2,
2773            Self::Alternate => 3,
2774            Self::Abort => 4,
2775            Self::Controller => 5,
2776        }
2777    }
2778
2779    fn configured() -> Self {
2780        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
2781            1 => Self::Hit,
2782            2 => Self::Miss,
2783            3 => Self::Alternate,
2784            4 => Self::Abort,
2785            5 => Self::Controller,
2786            _ => Self::Disabled,
2787        }
2788    }
2789
2790    fn action(self, generation: u64) -> OptiForkAction {
2791        match self {
2792            Self::Hit => OptiForkAction::Hit,
2793            Self::Miss => OptiForkAction::Miss,
2794            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
2795            Self::Alternate => OptiForkAction::Miss,
2796            Self::Abort => OptiForkAction::Abort,
2797            Self::Disabled | Self::Controller => {
2798                unreachable!("non-forced mode cannot choose a forced fork action")
2799            }
2800        }
2801    }
2802
2803    fn is_forced(self) -> bool {
2804        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
2805    }
2806}
2807
2808/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
2809pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
2810    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
2811}
2812
2813/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
2814/// two-token draft-probability product. Serving can call this only through its explicit
2815/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
2816pub fn set_optipipe_controller_threshold(threshold: f32) {
2817    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
2818    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
2819    set_optipipe_gate_mode(OptiForkGateMode::Controller);
2820}
2821
2822#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2823pub struct OptiForkGateStats {
2824    pub attempts: u64,
2825    pub hits: u64,
2826    pub misses: u64,
2827    pub abort_drains: u64,
2828    pub refusals: u64,
2829    pub gate_checks: u64,
2830    pub gate_admits: u64,
2831    pub gate_rejects: u64,
2832    pub reconciles: u64,
2833    pub wasted_draft_tokens: u64,
2834    pub shadow_draft_tokens: u64,
2835    pub breaker_trips: u64,
2836}
2837
2838#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
2839pub struct OptiForkStateIdentity {
2840    pub trunk_kv_bytes: usize,
2841    pub recurrent_bytes: usize,
2842    pub scratch_kv_bytes: usize,
2843    pub hidden_bytes: usize,
2844}
2845
2846pub fn reset_optipipe_gate_stats() {
2847    for counter in [
2848        &OPTI_FORK_ATTEMPTS,
2849        &OPTI_FORK_HITS,
2850        &OPTI_FORK_MISSES,
2851        &OPTI_FORK_ABORT_DRAINS,
2852        &OPTI_FORK_REFUSALS,
2853        &OPTI_GATE_CHECKS,
2854        &OPTI_GATE_ADMITS,
2855        &OPTI_GATE_REJECTS,
2856        &OPTI_RECONCILES,
2857        &OPTI_WASTED_DRAFT_TOKENS,
2858        &OPTI_SHADOW_DRAFT_TOKENS,
2859        &OPTI_BREAKER_TRIPS,
2860    ] {
2861        counter.store(0, std::sync::atomic::Ordering::Relaxed);
2862    }
2863}
2864
2865pub fn optipipe_gate_stats() -> OptiForkGateStats {
2866    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
2867    OptiForkGateStats {
2868        attempts: load(&OPTI_FORK_ATTEMPTS),
2869        hits: load(&OPTI_FORK_HITS),
2870        misses: load(&OPTI_FORK_MISSES),
2871        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
2872        refusals: load(&OPTI_FORK_REFUSALS),
2873        gate_checks: load(&OPTI_GATE_CHECKS),
2874        gate_admits: load(&OPTI_GATE_ADMITS),
2875        gate_rejects: load(&OPTI_GATE_REJECTS),
2876        reconciles: load(&OPTI_RECONCILES),
2877        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
2878        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
2879        breaker_trips: load(&OPTI_BREAKER_TRIPS),
2880    }
2881}
2882
2883#[derive(Clone, Copy, Debug)]
2884struct OptiControllerPolicy {
2885    threshold: f32,
2886    consecutive_misses: u8,
2887    breaker_tripped: bool,
2888}
2889
2890impl OptiControllerPolicy {
2891    fn configured() -> Self {
2892        Self {
2893            threshold: f32::from_bits(
2894                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
2895            ),
2896            consecutive_misses: 0,
2897            breaker_tripped: false,
2898        }
2899    }
2900
2901    fn admit(&self, q_proxy: f32) -> bool {
2902        q_proxy.is_finite()
2903            && (0.0..=1.0).contains(&q_proxy)
2904            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
2905    }
2906
2907    /// Returns true exactly when this resolution newly trips the three-miss breaker.
2908    fn resolve(&mut self, hit: bool) -> bool {
2909        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
2910        // every optimistic opportunity, so the safety breaker is measured separately and must
2911        // not silently turn this arm into "three attempts then serial".
2912        if self.threshold == 0.0 {
2913            self.consecutive_misses = 0;
2914            return false;
2915        }
2916        if hit {
2917            self.consecutive_misses = 0;
2918            return false;
2919        }
2920        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
2921        if !self.breaker_tripped && self.consecutive_misses >= 3 {
2922            self.breaker_tripped = true;
2923            return true;
2924        }
2925        false
2926    }
2927}
2928
2929#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2930enum OptiForkAction {
2931    Hit,
2932    Miss,
2933    Abort,
2934}
2935
2936#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2937struct OptiForkGeneration {
2938    id: u64,
2939    slot: usize,
2940}
2941
2942#[derive(Default)]
2943struct OptiForkGenerationTracker {
2944    next: u64,
2945    live: [Option<u64>; 2],
2946}
2947
2948impl OptiForkGenerationTracker {
2949    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
2950        let generation = OptiForkGeneration {
2951            id: self.next,
2952            slot: (self.next & 1) as usize,
2953        };
2954        if let Some(live) = self.live[generation.slot] {
2955            return Err(format!(
2956                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
2957                generation.slot,
2958            )
2959            .into());
2960        }
2961        self.next += 1;
2962        self.live[generation.slot] = Some(generation.id);
2963        Ok(generation)
2964    }
2965
2966    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
2967        match self.live[generation.slot] {
2968            Some(id) if id == generation.id => {
2969                self.live[generation.slot] = None;
2970                Ok(())
2971            }
2972            other => Err(format!(
2973                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
2974                generation.id, generation.slot,
2975            )
2976            .into()),
2977        }
2978    }
2979}
2980
2981struct OptiForkSeedGeneration {
2982    h_seed: CudaSlice<f32>,
2983    fill_prev: CudaSlice<f32>,
2984    scratch_len: usize,
2985}
2986
2987/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
2988/// generic cache helper accepts one device and therefore cannot copy GDN state split across
2989/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
2990/// device ownership.
2991fn opti_snapshot_stage_owned(
2992    e: &Engine,
2993    cache: &Cache,
2994    rt: &'static crate::pp::PpNRt,
2995    fence: &[usize],
2996) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
2997    let n = cache.kv.len();
2998    let mut snapshot = crate::cache::CacheSnapshot {
2999        kv_len: vec![None; n],
3000        tp_kv_len: vec![None; n],
3001        conv: (0..n).map(|_| None).collect(),
3002        ssm: (0..n).map(|_| None).collect(),
3003        pos: cache.pos,
3004    };
3005    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3006    Ok(snapshot)
3007}
3008
3009fn opti_snapshot_stage_owned_into(
3010    e: &Engine,
3011    cache: &Cache,
3012    rt: &'static crate::pp::PpNRt,
3013    fence: &[usize],
3014    snapshot: &mut crate::cache::CacheSnapshot,
3015) -> Result<(), Box<dyn std::error::Error>> {
3016    if fence.len() != rt.n_stages() + 1
3017        || snapshot.kv_len.len() != cache.kv.len()
3018        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3019    {
3020        return Err("optipipe stage-owned snapshot shape mismatch".into());
3021    }
3022    for stage in 0..rt.n_stages() {
3023        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3024    }
3025    snapshot.pos = cache.pos;
3026    Ok(())
3027}
3028
3029/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3030/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3031/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3032/// either point would capture one side of the fork at the wrong generation.
3033fn opti_snapshot_one_stage_owned_into(
3034    e: &Engine,
3035    cache: &Cache,
3036    rt: &'static crate::pp::PpNRt,
3037    fence: &[usize],
3038    stage: usize,
3039    snapshot: &mut crate::cache::CacheSnapshot,
3040) -> Result<(), Box<dyn std::error::Error>> {
3041    if fence.len() != rt.n_stages() + 1
3042        || snapshot.kv_len.len() != cache.kv.len()
3043        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3044        || stage >= rt.n_stages()
3045    {
3046        return Err("optipipe single-stage snapshot shape mismatch".into());
3047    }
3048    let _scope = rt.enter(stage);
3049    let owner = rt.engine(stage, e);
3050    for il in fence[stage]..fence[stage + 1] {
3051        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3052        snapshot.tp_kv_len[il] = cache.tp_kv[il]
3053            .as_ref()
3054            .map(crate::tp::ResidentTpKvCache::committed_len);
3055        match &cache.recur[il] {
3056            Some(recur) => {
3057                match snapshot.conv[il].as_mut() {
3058                    Some(dst) => {
3059                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3060                    }
3061                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3062                }
3063                match snapshot.ssm[il].as_mut() {
3064                    Some(dst) => {
3065                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3066                    }
3067                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3068                }
3069            }
3070            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3071                return Err(
3072                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3073                );
3074            }
3075            None => {}
3076        }
3077    }
3078    snapshot.pos = cache.pos;
3079    Ok(())
3080}
3081
3082/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3083/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3084/// resolve, so the reconcile tables and conditional restores are stage-local.
3085struct OptiForkState {
3086    mode: OptiForkGateMode,
3087    controller: Option<OptiControllerPolicy>,
3088    generations: OptiForkGenerationTracker,
3089    active_snapshot_slot: usize,
3090    alternate_snapshot: crate::cache::CacheSnapshot,
3091    seeds: [OptiForkSeedGeneration; 2],
3092    rt: &'static crate::pp::PpNRt,
3093    fence: [usize; 3],
3094    split: usize,
3095    len_ptrs: CudaSlice<u64>,
3096    saved_lens: CudaSlice<i32>,
3097    forced_acc: CudaSlice<u32>,
3098    valid: CudaSlice<u32>,
3099    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3100    logical_payload_bytes: [usize; 2],
3101}
3102
3103struct OptiForkTicket {
3104    generation: OptiForkGeneration,
3105    boundary: Option<VerifyBoundaryTicket>,
3106    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3107    settled: bool,
3108}
3109
3110struct OptiControllerTicket {
3111    generation: OptiForkGeneration,
3112    boundary: Option<VerifyBoundaryTicket>,
3113    ckpt: Option<VerifyCkpt>,
3114    verify_tokens: [u32; 2],
3115    draft_prob: f32,
3116    eager_seed: Option<CudaSlice<f32>>,
3117    q_proxy: f32,
3118    scratch_len: usize,
3119    issued_at: std::time::Instant,
3120    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3121    settled: bool,
3122}
3123
3124struct OptiControllerPrepared {
3125    verify_tokens: [u32; 2],
3126    draft_prob: f32,
3127    eager_seed: Option<CudaSlice<f32>>,
3128    q_proxy: f32,
3129    scratch_len: usize,
3130}
3131
3132impl OptiControllerTicket {
3133    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3134        self.boundary
3135            .take()
3136            .expect("controller boundary ticket already consumed")
3137    }
3138
3139    fn take_ckpt(&mut self) -> VerifyCkpt {
3140        self.ckpt
3141            .take()
3142            .expect("controller verify checkpoint already consumed")
3143    }
3144
3145    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3146        self.eager_seed.take()
3147    }
3148
3149    fn settle(&mut self) {
3150        self.settled = true;
3151    }
3152}
3153
3154impl Drop for OptiControllerTicket {
3155    fn drop(&mut self) {
3156        if !self.settled {
3157            let _ = self.drain.synchronize();
3158            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3159        }
3160    }
3161}
3162
3163impl OptiForkTicket {
3164    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3165        self.boundary
3166            .take()
3167            .expect("fork ticket boundary already consumed")
3168    }
3169
3170    fn settle(&mut self) {
3171        self.settled = true;
3172    }
3173}
3174
3175impl Drop for OptiForkTicket {
3176    fn drop(&mut self) {
3177        if !self.settled {
3178            let _ = self.drain.synchronize();
3179            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3180        }
3181    }
3182}
3183
3184impl OptiForkState {
3185    #[allow(clippy::too_many_arguments)]
3186    fn new(
3187        e: &Engine,
3188        cache: &Cache,
3189        mode: OptiForkGateMode,
3190        alternate_snapshot: crate::cache::CacheSnapshot,
3191        h_seed: &CudaSlice<f32>,
3192        fill_prev: &CudaSlice<f32>,
3193        rt: &'static crate::pp::PpNRt,
3194        split: usize,
3195        n_layer: usize,
3196    ) -> Result<Self, Box<dyn std::error::Error>> {
3197        let fence = [0, split, n_layer];
3198        let mut logical_payload_bytes = [0usize; 2];
3199        for stage in 0..2 {
3200            for il in fence[stage]..fence[stage + 1] {
3201                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3202                    .as_ref()
3203                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3204                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3205                    .as_ref()
3206                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3207            }
3208        }
3209        let seeds = [
3210            OptiForkSeedGeneration {
3211                h_seed: e.clone_dtod(h_seed)?,
3212                fill_prev: e.clone_dtod(fill_prev)?,
3213                scratch_len: 0,
3214            },
3215            OptiForkSeedGeneration {
3216                h_seed: e.clone_dtod(h_seed)?,
3217                fill_prev: e.clone_dtod(fill_prev)?,
3218                scratch_len: 0,
3219            },
3220        ];
3221        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3222            let _stage = rt.enter(0);
3223            let e0 = rt.engine(0, e);
3224            (
3225                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3226                e0.htod_i32(&vec![0; split])?,
3227                e0.alloc_u32_zeroed(2)?,
3228                e0.alloc_u32_zeroed(1)?,
3229                e0.stream(),
3230            )
3231        };
3232        logical_payload_bytes[0] += seeds
3233            .iter()
3234            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3235            .sum::<usize>();
3236        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3237            + saved_lens.len() * std::mem::size_of::<i32>()
3238            + forced_acc.len() * std::mem::size_of::<u32>()
3239            + valid.len() * std::mem::size_of::<u32>();
3240        Ok(Self {
3241            mode,
3242            controller: (mode == OptiForkGateMode::Controller)
3243                .then(OptiControllerPolicy::configured),
3244            generations: OptiForkGenerationTracker::default(),
3245            active_snapshot_slot: 0,
3246            alternate_snapshot,
3247            seeds,
3248            rt,
3249            fence,
3250            split,
3251            len_ptrs,
3252            saved_lens,
3253            forced_acc,
3254            valid,
3255            stage0_stream,
3256            logical_payload_bytes,
3257        })
3258    }
3259
3260    fn reserve(
3261        &mut self,
3262        current_snapshot: &mut crate::cache::CacheSnapshot,
3263    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3264        let generation = self.generations.reserve()?;
3265        if generation.slot != self.active_snapshot_slot {
3266            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3267            self.active_snapshot_slot = generation.slot;
3268        }
3269        Ok(generation)
3270    }
3271
3272    fn capture_seed(
3273        &mut self,
3274        e: &Engine,
3275        generation: OptiForkGeneration,
3276        h_seed: &CudaSlice<f32>,
3277        fill_prev: &CudaSlice<f32>,
3278        scratch_len: usize,
3279    ) -> Result<(), Box<dyn std::error::Error>> {
3280        let seed = &mut self.seeds[generation.slot];
3281        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3282        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3283        seed.scratch_len = scratch_len;
3284        Ok(())
3285    }
3286
3287    fn ticket(
3288        &self,
3289        generation: OptiForkGeneration,
3290        boundary: VerifyBoundaryTicket,
3291    ) -> OptiForkTicket {
3292        OptiForkTicket {
3293            generation,
3294            boundary: Some(boundary),
3295            drain: self.stage0_stream.clone(),
3296            settled: false,
3297        }
3298    }
3299
3300    #[allow(clippy::too_many_arguments)]
3301    fn controller_ticket(
3302        &self,
3303        generation: OptiForkGeneration,
3304        boundary: VerifyBoundaryTicket,
3305        ckpt: VerifyCkpt,
3306        verify_tokens: [u32; 2],
3307        draft_prob: f32,
3308        eager_seed: Option<CudaSlice<f32>>,
3309        q_proxy: f32,
3310        scratch_len: usize,
3311    ) -> OptiControllerTicket {
3312        OptiControllerTicket {
3313            generation,
3314            boundary: Some(boundary),
3315            ckpt: Some(ckpt),
3316            verify_tokens,
3317            draft_prob,
3318            eager_seed,
3319            q_proxy,
3320            scratch_len,
3321            issued_at: std::time::Instant::now(),
3322            drain: self.stage0_stream.clone(),
3323            settled: false,
3324        }
3325    }
3326
3327    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3328        self.generations.reserve()
3329    }
3330
3331    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3332        &mut self.alternate_snapshot
3333    }
3334
3335    fn promote_successor_snapshot(
3336        &mut self,
3337        current_snapshot: &mut crate::cache::CacheSnapshot,
3338        generation: OptiForkGeneration,
3339    ) {
3340        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3341        self.active_snapshot_slot = generation.slot;
3342    }
3343
3344    fn queue_actual_reconcile(
3345        &mut self,
3346        e: &Engine,
3347        snapshot: &crate::cache::CacheSnapshot,
3348        acc: &CudaSlice<u32>,
3349        optimistic_pending: u32,
3350        base: usize,
3351    ) -> Result<(), Box<dyn std::error::Error>> {
3352        let saved: Vec<i32> = (0..self.split)
3353            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3354            .collect();
3355        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3356        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3357        // the validity/reconcile kernels must never peer-read acc before it is written. The
3358        // increment-1 harness uses primary stage 0, where stream order already provides this.
3359        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3360            self.rt.fence_stages_behind(&e.stream())?;
3361        }
3362        let _stage = self.rt.enter(0);
3363        let e0 = self.rt.engine(0, e);
3364        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3365        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3366        e0.spec_fork_reconcile_kv(
3367            &self.len_ptrs,
3368            &self.saved_lens,
3369            acc,
3370            &self.valid,
3371            base,
3372            self.split,
3373        )
3374    }
3375
3376    fn finish_actual_reconcile(
3377        &mut self,
3378        e: &Engine,
3379        cache: &mut Cache,
3380        snapshot: &crate::cache::CacheSnapshot,
3381        n_acc: usize,
3382        base: usize,
3383        hit: bool,
3384    ) -> Result<(), Box<dyn std::error::Error>> {
3385        if hit {
3386            return Ok(());
3387        }
3388        let len_delta = base + n_acc;
3389        for il in 0..self.split {
3390            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3391                kv.len = saved + len_delta;
3392            }
3393        }
3394        {
3395            let _stage = self.rt.enter(1);
3396            let e1 = self.rt.engine(1, e);
3397            for il in self.split..self.fence[2] {
3398                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3399                    kv.len = saved + len_delta;
3400                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3401                }
3402            }
3403        }
3404        self.rt.publish_to(0, &e.stream())?;
3405        Ok(())
3406    }
3407
3408    fn cancel_controller_ticket(
3409        &mut self,
3410        e: &Engine,
3411        cache: &mut Cache,
3412        scratch: &mut MtpScratch,
3413        snapshot: &crate::cache::CacheSnapshot,
3414        ticket: &mut OptiControllerTicket,
3415    ) -> Result<(), Box<dyn std::error::Error>> {
3416        {
3417            let _stage = self.rt.enter(0);
3418            let e0 = self.rt.engine(0, e);
3419            for il in 0..self.split {
3420                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3421                    kv.len = saved;
3422                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3423                }
3424            }
3425        }
3426        scratch.set_len(e, snapshot.pos)?;
3427        ticket.settle();
3428        self.generations.retire(ticket.generation)?;
3429        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3430        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3431        eprintln!(
3432            "[opti-controller] tail-drain generation={} slot={}",
3433            ticket.generation.id, ticket.generation.slot,
3434        );
3435        Ok(())
3436    }
3437
3438    #[allow(clippy::too_many_arguments)]
3439    fn reconcile(
3440        &mut self,
3441        e: &Engine,
3442        cache: &mut Cache,
3443        scratch: &mut MtpScratch,
3444        snapshot: &crate::cache::CacheSnapshot,
3445        h_seed: &mut CudaSlice<f32>,
3446        fill_prev: &mut CudaSlice<f32>,
3447        generation: OptiForkGeneration,
3448        action: OptiForkAction,
3449        optimistic_pending: u32,
3450    ) -> Result<(), Box<dyn std::error::Error>> {
3451        debug_assert!(action != OptiForkAction::Abort);
3452        let miss_started = std::time::Instant::now();
3453        let keep = action == OptiForkAction::Hit;
3454        let saved: Vec<i32> = (0..self.split)
3455            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3456            .collect();
3457        let seed = &self.seeds[generation.slot];
3458        {
3459            let _stage = self.rt.enter(0);
3460            let e0 = self.rt.engine(0, e);
3461            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3462            let forced = if keep {
3463                [1u32, optimistic_pending]
3464            } else {
3465                [0u32, optimistic_pending]
3466            };
3467            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3468            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3469            e0.spec_fork_reconcile_kv(
3470                &self.len_ptrs,
3471                &self.saved_lens,
3472                &self.forced_acc,
3473                &self.valid,
3474                0,
3475                self.split,
3476            )?;
3477            for il in 0..self.split {
3478                if let Some(recur) = cache.recur[il].as_mut() {
3479                    let conv = snapshot.conv[il]
3480                        .as_ref()
3481                        .ok_or("optipipe stage0 snapshot missing conv state")?;
3482                    let ssm = snapshot.ssm[il]
3483                        .as_ref()
3484                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
3485                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3486                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3487                }
3488            }
3489            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3490            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3491        }
3492
3493        if keep {
3494            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3495            return Ok(());
3496        }
3497
3498        for il in 0..self.split {
3499            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3500                kv.len = saved;
3501            }
3502        }
3503        scratch.set_len(e, seed.scratch_len)?;
3504        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3505        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3506        let caller = e.stream();
3507        self.rt.publish_to(0, &caller)?;
3508        caller.synchronize()?;
3509        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3510        eprintln!(
3511            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3512            generation.id, generation.slot,
3513        );
3514        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3515        Ok(())
3516    }
3517
3518    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3519        self.generations.retire(generation)
3520    }
3521}
3522
3523fn rewind_tp_kv_verified_prefix(
3524    tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3525    saved_lens: &[Option<usize>],
3526    accepted: usize,
3527) -> Result<(), Box<dyn std::error::Error>> {
3528    if tp_kv.len() != saved_lens.len() {
3529        return Err("spec TP KV snapshot shape mismatch".into());
3530    }
3531    for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3532        match (cache.as_mut(), *saved) {
3533            (Some(cache), Some(saved)) => {
3534                let target = saved
3535                    .checked_add(accepted)
3536                    .ok_or("spec TP KV committed length overflow")?;
3537                cache.rewind_to(target)?;
3538            }
3539            (None, None) => {}
3540            _ => {
3541                return Err(
3542                    format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3543                );
3544            }
3545        }
3546    }
3547    Ok(())
3548}
3549
3550impl HybridModel {
3551    fn mtp_head_count(&self) -> usize {
3552        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3553    }
3554
3555    fn mtp_head_at(&self, index: usize) -> &MtpHead {
3556        if index == 0 {
3557            self.mtp.as_ref().expect("MTP head 0 is unavailable")
3558        } else {
3559            &self.mtp_extra[index - 1]
3560        }
3561    }
3562
3563    fn new_mtp_scratch(
3564        &self,
3565        e: &Engine,
3566        cap: usize,
3567    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3568        let mut scratch = MtpScratch::new(
3569            e,
3570            &self.cfg,
3571            &self.plan,
3572            cap,
3573            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
3574        )?;
3575        for head in &self.mtp_extra {
3576            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
3577        }
3578        Ok(scratch)
3579    }
3580
3581    fn opti_graph_draft_step(
3582        &self,
3583        e: &Engine,
3584        mtp: &MtpHead,
3585        dctx: &mut DraftGraphCtx,
3586        scratch: &mut MtpScratch,
3587        d_vocab: usize,
3588    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3589        dctx.graph
3590            .as_ref()
3591            .ok_or("optipipe controller requires the greedy draft graph")?
3592            .launch()?;
3593        scratch.kv.len += 1;
3594        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
3595        if (idx as usize) >= d_vocab {
3596            return Err(
3597                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
3598            );
3599        }
3600        let probability = e.dtoh(&dctx.g_p)?[0];
3601        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3602            return Err(format!("optipipe draft probability is invalid: {probability}").into());
3603        }
3604        let token = match &mtp.d2t {
3605            Some(map) => map[idx as usize],
3606            None => idx,
3607        };
3608        if token != idx {
3609            e.set_u32_one(&mut dctx.g_tok, token)?;
3610        }
3611        Ok((token, probability))
3612    }
3613
3614    #[allow(clippy::too_many_arguments)]
3615    fn opti_controller_draft_step(
3616        &self,
3617        e: &Engine,
3618        mtp: &MtpHead,
3619        dctx: &mut DraftGraphCtx,
3620        scratch: &mut MtpScratch,
3621        d_vocab: usize,
3622        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
3623        eager_pos: usize,
3624        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3625    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
3626        if dctx.graph.is_some() {
3627            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
3628        }
3629        let (input_token, input_seed) = eager_state
3630            .take()
3631            .ok_or("optipipe eager continuation seed is unavailable")?;
3632        let (logits, next_seed) = self.mtp_head_forward_dev(
3633            e,
3634            mtp,
3635            input_token,
3636            &input_seed,
3637            scratch,
3638            eager_pos,
3639            embd_dev,
3640            None,
3641        )?;
3642        let token_d = e.argmax_token_device(&logits, d_vocab)?;
3643        let idx = e.dtoh_u32_one(&token_d)?;
3644        if (idx as usize) >= d_vocab {
3645            return Err(format!(
3646                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
3647            )
3648            .into());
3649        }
3650        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
3651        let probability = e.dtoh(&probability_d)?[0];
3652        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
3653            return Err(
3654                format!("optipipe eager draft probability is invalid: {probability}").into(),
3655            );
3656        }
3657        let token = match &mtp.d2t {
3658            Some(map) => map[idx as usize],
3659            None => idx,
3660        };
3661        *eager_state = Some((token, next_seed));
3662        Ok((token, probability))
3663    }
3664
3665    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
3666    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
3667    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
3668    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
3669    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
3670    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
3671    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
3672    /// transfer + host argmax per draft token from the K-token draft chain.
3673    #[allow(clippy::too_many_arguments)]
3674    fn mtp_head_forward_dev(
3675        &self,
3676        e: &Engine,
3677        mtp: &MtpHead,
3678        e_tok: u32,
3679        h_seed: &CudaSlice<f32>,
3680        scratch: &mut MtpScratch,
3681        mtp_pos: usize,
3682        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3683        mask: Option<(&CudaSlice<u32>, usize)>,
3684    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3685        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
3686    }
3687
3688    #[allow(clippy::too_many_arguments)]
3689    fn mtp_head_forward_dev_at(
3690        &self,
3691        e: &Engine,
3692        mtp: &MtpHead,
3693        e_tok: u32,
3694        h_seed: &CudaSlice<f32>,
3695        scratch: &mut MtpScratch,
3696        scratch_index: usize,
3697        mtp_pos: usize,
3698        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3699        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
3700        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
3701        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
3702        mask: Option<(&CudaSlice<u32>, usize)>,
3703    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3704        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
3705        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
3706        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
3707        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
3708        static ANAT_NS: [AtomicU64; 5] = [
3709            AtomicU64::new(0),
3710            AtomicU64::new(0),
3711            AtomicU64::new(0),
3712            AtomicU64::new(0),
3713            AtomicU64::new(0),
3714        ];
3715        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
3716        let anat = {
3717            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3718            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
3719        };
3720        if anat {
3721            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
3722        }
3723        let t_all = std::time::Instant::now();
3724        let mut t_ph = std::time::Instant::now();
3725        let mut anat_mark = |i: usize,
3726                             e: &Engine,
3727                             t: &mut std::time::Instant|
3728         -> Result<(), Box<dyn std::error::Error>> {
3729            if anat {
3730                e.stream().synchronize()?;
3731                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
3732                *t = std::time::Instant::now();
3733            }
3734            Ok(())
3735        };
3736        let cfg = &self.cfg;
3737        let n_embd = cfg.n_embd as usize;
3738        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
3739        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
3740        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
3741        let eps = cfg.rms_eps;
3742        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
3743
3744        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
3745        // expands this one row on CPU and transfers n_embd f32 values instead.
3746        let e_emb = match embd_dev {
3747            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
3748            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
3749        };
3750
3751        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
3752        let mut e_norm = e.zeros(n_embd)?;
3753        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
3754        let mut h_norm = e.zeros(n_embd)?;
3755        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
3756
3757        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
3758        let mut concat = e.zeros(2 * n_embd)?;
3759        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
3760        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
3761
3762        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
3763        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
3764
3765        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
3766        let mut a_norm = e.zeros(di)?;
3767        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
3768        anat_mark(0, e, &mut t_ph)?;
3769
3770        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
3771        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
3772        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
3773        // advances only the device counter).
3774        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
3775            // step35 MTP block: PER-LAYER geometry + a separate head-wise gate + an SWA window,
3776            // none of which the dc launcher can express (see `mtp_step35_attn`). Host-len arm.
3777            // Advances BOTH the host len and the device counter itself (unlike the dc arm,
3778            // whose host-side mirror the caller does).
3779            (Mixer::Full(fa), Some(g)) => {
3780                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
3781            }
3782            (Mixer::Full(fa), None) => {
3783                let out = self.mtp_full_attn_dc(
3784                    e,
3785                    fa,
3786                    &a_norm,
3787                    &pos_d,
3788                    scratch,
3789                    scratch_index,
3790                    mtp.geom.as_ref(),
3791                )?;
3792                scratch.plane_mut(scratch_index).0.len += 1;
3793                out
3794            }
3795            (Mixer::Linear(_), _) => {
3796                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
3797            }
3798            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
3799        };
3800        anat_mark(1, e, &mut t_ph)?;
3801
3802        // op 7: x1 = inpSA + attn_out
3803        let mut x1 = e.zeros(di)?;
3804        e.add(&inp_sa, &attn_out, &mut x1, di)?;
3805
3806        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
3807        let mut z = e.zeros(di)?;
3808        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
3809
3810        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
3811        let ffn_out = match &mtp.ffn {
3812            crate::hybrid::Ffn::Dense {
3813                ffn_gate,
3814                ffn_up,
3815                ffn_down,
3816            } => {
3817                let n_ff = ffn_gate.out_features();
3818                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
3819                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
3820                    (
3821                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
3822                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
3823                    )
3824                } else {
3825                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
3826                };
3827                let mut act = e.zeros(n_ff)?;
3828                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
3829                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
3830                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
3831                // passes None, which is `ffn_act`'s dispatch verbatim.
3832                Self::ffn_act_lim(
3833                    e,
3834                    &self.cfg,
3835                    &gate,
3836                    &up,
3837                    1.0,
3838                    1.0,
3839                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
3840                    &mut act,
3841                    n_ff,
3842                )?;
3843                e.matmul(ffn_down, &act, 1)?
3844            }
3845            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
3846            // so they never alias trunk layer 0's cache keys.
3847            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
3848        };
3849        anat_mark(2, e, &mut t_ph)?;
3850
3851        // op 10: h_nextn = x1 + ffn_out (at di)
3852        let mut h_inner = e.zeros(di)?;
3853        e.add(&x1, &ffn_out, &mut h_inner, di)?;
3854
3855        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
3856        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
3857        let h_nextn = match mtp.geom.as_ref() {
3858            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
3859            None => h_inner,
3860        };
3861
3862        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
3863        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
3864        let mut final_h = e.zeros(n_embd)?;
3865        e.rms_norm(
3866            &h_nextn,
3867            final_norm.float_data(),
3868            &mut final_h,
3869            n_embd,
3870            1,
3871            eps,
3872        )?;
3873
3874        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
3875        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
3876        let mut logits = e.matmul(head, &final_h, 1)?;
3877        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
3878        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
3879        if let Some((mask_d, mw)) = mask {
3880            let d_vocab = head.out_features();
3881            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
3882        }
3883        anat_mark(3, e, &mut t_ph)?;
3884        if anat {
3885            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
3886            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
3887            if n % 128 == 0 {
3888                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
3889                eprintln!(
3890                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
3891                    us(0),
3892                    us(1),
3893                    us(2),
3894                    us(3),
3895                    us(4)
3896                );
3897            }
3898        }
3899        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
3900        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
3901        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
3902    }
3903
3904    #[allow(clippy::too_many_arguments)]
3905    fn mtp_chain_forward_dev(
3906        &self,
3907        e: &Engine,
3908        tokens: &[u32],
3909        seeds: &[CudaSlice<f32>],
3910        scratch: &mut MtpScratch,
3911        committed_scratch_len: usize,
3912        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
3913        mask: Option<(&CudaSlice<u32>, usize)>,
3914    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
3915        if tokens.is_empty() || tokens.len() != seeds.len() {
3916            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
3917        }
3918        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
3919        let head = self.mtp_head_at(index);
3920        scratch.set_plane_len(e, index, committed_scratch_len)?;
3921
3922        let mut last = None;
3923        for row in 0..tokens.len() {
3924            let is_last = row + 1 == tokens.len();
3925            last = Some(self.mtp_head_forward_dev_at(
3926                e,
3927                head,
3928                tokens[row],
3929                &seeds[row],
3930                scratch,
3931                index,
3932                committed_scratch_len + row + 1,
3933                embd_dev,
3934                if is_last { mask } else { None },
3935            )?);
3936        }
3937        Ok(last.expect("non-empty MTP prefix produced no row"))
3938    }
3939
3940    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
3941    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
3942    /// the dc path, and all three are properties of this arch's MTP block:
3943    ///
3944    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
3945    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
3946    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
3947    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
3948    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
3949    ///    starts at row 0 — it cannot express a nonzero offset, so a windowed dc arm would need a
3950    ///    new kernel. That is deliberately not built here: see the CUDA-graph note below.
3951    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
3952    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
3953    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
3954    ///    resolved `Step35MtpGeom`, never from `cfg`.
3955    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
3956    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
3957    ///    fused-into-wq `q_gate_split` form the dc arm handles.
3958    ///
3959    /// WHY EAGER-ONLY IS NOT A GAP TODAY: `mtp_head_forward_cap` refuses step35 heads
3960    /// explicitly (the SWA-window refusal below), so the graph draft never engages for step35
3961    /// models regardless of eligibility — the eager chain IS the served path. `mtp_head_forward_cap` refuses step35 explicitly rather
3962    /// than silently capturing a window-less (wrong past `win` draft rows) graph.
3963    ///
3964    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
3965    /// caller must not mirror.
3966    fn mtp_step35_attn(
3967        &self,
3968        e: &Engine,
3969        fa: &FullAttnLayer,
3970        g: &crate::hybrid::Step35MtpGeom,
3971        h: &CudaSlice<f32>,
3972        pos_d: &CudaSlice<i32>,
3973        scratch: &mut MtpScratch,
3974        scratch_index: usize,
3975    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3976        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
3977        let eps = self.cfg.rms_eps;
3978        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
3979        let n_embd = self.cfg.n_embd as usize;
3980        let gw = fa
3981            .attn_gate
3982            .as_ref()
3983            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
3984
3985        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
3986            && e.uses_q8_1_fast(&fa.wk)
3987            && e.uses_q8_1_fast(&fa.wv)
3988            && e.uses_q8_1_fast(gw)
3989        {
3990            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
3991            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
3992                Some(t3) => t3,
3993                None => (
3994                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
3995                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
3996                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
3997                ),
3998            };
3999            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4000        } else {
4001            (
4002                e.matmul(&fa.wq, h, 1)?,
4003                e.matmul(&fa.wk, h, 1)?,
4004                e.matmul(&fa.wv, h, 1)?,
4005                e.matmul(gw, h, 1)?,
4006            )
4007        };
4008
4009        let mut q = e.uninit(nh * hd)?;
4010        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4011        let mut k = e.uninit(nkv * hd)?;
4012        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4013        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4014        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4015        // the resolved flag, not the constant, so an all-full sibling stays correct.
4016        let ff = if g.swa {
4017            None
4018        } else {
4019            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4020        };
4021        #[cfg(debug_assertions)]
4022        if let Some(ff) = ff {
4023            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4024        }
4025        e.rope_neox2(
4026            &mut q,
4027            &mut k,
4028            pos_d,
4029            hd,
4030            g.n_rot,
4031            nh,
4032            nkv,
4033            1,
4034            g.rope_base,
4035            1.0,
4036            ff,
4037        )?;
4038
4039        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4040        // length on the host anyway, and the windowed view below needs it there to compute the
4041        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4042        // dc-family consumer of this scratch still agree.
4043        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4044        assert!(
4045            kv.len < scratch_cap,
4046            "step35 MTP scratch overflow ({} >= {})",
4047            kv.len,
4048            scratch_cap
4049        );
4050        let next_len = kv.len + 1;
4051        let (off, t_kv) = if g.swa && next_len > g.window {
4052            (next_len - g.window, g.window)
4053        } else {
4054            (0, next_len)
4055        };
4056        let write_row = e.prepare_kv_append(kv, off & !31usize, 1)?;
4057        e.append_kv_quantized(
4058            &k,
4059            &v0,
4060            &mut kv.k,
4061            &mut kv.v,
4062            write_row,
4063            kv.kv_dim_k,
4064            kv.kv_dim_v,
4065            kv.k_tok_bytes,
4066            kv.v_tok_bytes,
4067            false,
4068        )?;
4069        kv.len = next_len;
4070        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4071        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4072        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4073        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4074        // therefore live, not theoretical.
4075        let physical = kv.physical_rows(off, off + t_kv)?;
4076        let k_view = e.view_u8_range(
4077            &kv.k,
4078            physical.start * kv.k_tok_bytes,
4079            physical.end * kv.k_tok_bytes,
4080        );
4081        let v_view = e.view_u8_range(
4082            &kv.v,
4083            physical.start * kv.v_tok_bytes,
4084            physical.end * kv.v_tok_bytes,
4085        );
4086        let mut attn = e.uninit(nh * hd)?;
4087        e.fa_decode_kvmod(
4088            &q,
4089            &k_view,
4090            &v_view,
4091            &mut attn,
4092            hd,
4093            nh,
4094            nkv,
4095            t_kv,
4096            scale,
4097            kv.k_tok_bytes,
4098            kv.v_tok_bytes,
4099            false,
4100        )?;
4101
4102        let mut ag = e.uninit(nh * hd)?;
4103        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4104        Ok(e.matmul(&fa.wo, &ag, 1)?)
4105    }
4106
4107    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4108    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4109    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4110    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4111    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4112    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4113    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4114    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4115    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4116    fn mtp_full_attn_dc(
4117        &self,
4118        e: &Engine,
4119        fa: &FullAttnLayer,
4120        h: &CudaSlice<f32>,
4121        pos_d: &CudaSlice<i32>,
4122        scratch: &mut MtpScratch,
4123        scratch_index: usize,
4124        geom: Option<&crate::hybrid::DraftGeom>,
4125    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4126        let cfg = &self.cfg;
4127        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4128        let geometry = cfg.full_attention_geometry_at(mtp_il);
4129        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4130        let n_head_kv = geom
4131            .map(|g| g.n_head_kv)
4132            .unwrap_or(geometry.n_head_kv as usize);
4133        let head_dim = geometry.head_dim_k as usize;
4134        let eps = cfg.rms_eps;
4135        let scale = geometry.attention_scale();
4136        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4137        let bucket_max = scratch.plane(scratch_index).1;
4138
4139        let (qf, mut k, v) =
4140            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4141                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4142                (
4143                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4144                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4145                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4146                )
4147            } else {
4148                (
4149                    e.matmul(&fa.wq, h, 1)?,
4150                    e.matmul(&fa.wk, h, 1)?,
4151                    e.matmul(&fa.wv, h, 1)?,
4152                )
4153            };
4154        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4155        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4156        let (mut q, gate) = if gated {
4157            let mut q = e.zeros(n_head * head_dim)?;
4158            let mut gate = e.zeros(n_head * head_dim)?;
4159            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4160            (q, Some(gate))
4161        } else {
4162            (qf, None)
4163        };
4164
4165        let mut qn = e.zeros(n_head * head_dim)?;
4166        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4167        q = qn;
4168        let mut kn = e.zeros(n_head_kv * head_dim)?;
4169        e.rms_norm(
4170            &k,
4171            fa.k_norm.float_data(),
4172            &mut kn,
4173            head_dim,
4174            n_head_kv,
4175            eps,
4176        )?;
4177        k = kn;
4178        let rope_dims = geometry.n_rot as usize;
4179        e.rope_neox(
4180            &mut q,
4181            pos_d,
4182            head_dim,
4183            rope_dims,
4184            n_head,
4185            1,
4186            geometry.rope_base,
4187            1.0,
4188        )?;
4189        e.rope_neox(
4190            &mut k,
4191            pos_d,
4192            head_dim,
4193            rope_dims,
4194            n_head_kv,
4195            1,
4196            geometry.rope_base,
4197            1.0,
4198        )?;
4199
4200        let kv = scratch.plane_mut(scratch_index).0;
4201        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4202        e.append_kv_quantized_dc(
4203            &k,
4204            &v,
4205            &mut kv.k,
4206            &mut kv.v,
4207            &kv.len_d,
4208            kv.kv_dim_k,
4209            kv.kv_dim_v,
4210            kv.k_tok_bytes,
4211            kv.v_tok_bytes,
4212            false,
4213        )?;
4214        e.inc_seqlen(&mut kv.len_d)?;
4215        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4216        // key range from the device counter.
4217        let k_view = e.view_u8(&kv.k, kv.k.len());
4218        let v_view = e.view_u8(&kv.v, kv.v.len());
4219        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4220        let mut attn = e.zeros(n_head * head_dim)?;
4221        e.fa_decode_dc(
4222            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4223            scale, ktb, vtb, false,
4224        )?;
4225
4226        let attn_g = match &gate {
4227            Some(gate) => {
4228                let mut gsig = e.zeros(n_head * head_dim)?;
4229                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4230                let mut ag = e.zeros(n_head * head_dim)?;
4231                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4232                ag
4233            }
4234            None => attn,
4235        };
4236        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4237    }
4238
4239    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4240    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4241    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4242    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4243    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4244    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4245    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4246    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4247    #[allow(clippy::too_many_arguments)]
4248    fn mtp_kv_fill_at(
4249        &self,
4250        e: &Engine,
4251        mtp: &MtpHead,
4252        tokens: &[u32],
4253        h: &CudaSlice<f32>,
4254        pos0: usize,
4255        scratch: &mut MtpScratch,
4256        scratch_index: usize,
4257        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4258    ) -> Result<(), Box<dyn std::error::Error>> {
4259        let cfg = &self.cfg;
4260        let n_embd = cfg.n_embd as usize;
4261        let eps = cfg.rms_eps;
4262        let t = tokens.len();
4263        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4264        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4265        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4266        let Mixer::Full(fa) = &mtp.mixer else {
4267            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4268        };
4269        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4270        let pos_d = e.htod_i32(&pos_vec)?;
4271
4272        // ops A/1/2: embed + the two input norms, T-wide.
4273        let e_emb = match embd_dev {
4274            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4275            None => e.htod(&self.embd.gather(n_embd, tokens))?,
4276        };
4277        let mut e_norm = e.zeros(t * n_embd)?;
4278        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4279        let mut h_norm = e.zeros(t * n_embd)?;
4280        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4281
4282        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4283        let mut concat = e.zeros(t * 2 * n_embd)?;
4284        for i in 0..t {
4285            e.copy_view_into(
4286                &mut concat,
4287                i * 2 * n_embd,
4288                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4289                n_embd,
4290            )?;
4291            e.copy_view_into(
4292                &mut concat,
4293                i * 2 * n_embd + n_embd,
4294                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4295                n_embd,
4296            )?;
4297        }
4298
4299        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4300        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4301        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4302        let mut a_norm = e.zeros(t * di)?;
4303        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4304
4305        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4306        // the fill only has to leave correct K/V rows behind for later chains to attend over.
4307        let n_head_kv = mtp
4308            .geom
4309            .as_ref()
4310            .map(|g| g.n_head_kv)
4311            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4312            .unwrap_or_else(|| {
4313                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4314                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4315            });
4316        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4317        let geometry = cfg.full_attention_geometry_at(mtp_il);
4318        let head_dim = geometry.head_dim_k as usize;
4319        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4320        let v = e.matmul(&fa.wv, &a_norm, t)?;
4321        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4322        e.rms_norm(
4323            &k,
4324            fa.k_norm.float_data(),
4325            &mut kn,
4326            head_dim,
4327            n_head_kv * t,
4328            eps,
4329        )?;
4330        k = kn;
4331        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4332        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4333        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4334        // writes K rows the attention arm then re-derives at a different theta: correct-looking
4335        // output with dead acceptance, invisible to the exactness gates.
4336        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4337            Some(s) => (
4338                s.n_rot,
4339                s.rope_base,
4340                if s.swa {
4341                    None
4342                } else {
4343                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4344                },
4345            ),
4346            None => (geometry.n_rot as usize, geometry.rope_base, None),
4347        };
4348        #[cfg(debug_assertions)]
4349        if let Some(ff) = ff {
4350            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
4351        }
4352        match ff {
4353            Some(f) => e.rope_neox_ff(
4354                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
4355            )?,
4356            None => e.rope_neox(
4357                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
4358            )?,
4359        }
4360
4361        let kv = scratch.plane_mut(scratch_index).0;
4362        // Match the trunk prime contract: a chunk may need the aligned window immediately before
4363        // its first row, so preserve that prefix when the physical tail rebases at wrap.
4364        let retain_from = kv
4365            .ring
4366            .as_ref()
4367            .map(|ring| pos0.saturating_sub(ring.window() - 1) & !31usize)
4368            .unwrap_or(0);
4369        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
4370        for i in 0..t {
4371            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
4372            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
4373            e.append_kv_quantized_view(
4374                &k_row,
4375                &v_row,
4376                &mut kv.k,
4377                &mut kv.v,
4378                write_row + i,
4379                kv.kv_dim_k,
4380                kv.kv_dim_v,
4381                kv.k_tok_bytes,
4382                kv.v_tok_bytes,
4383                false,
4384            )?;
4385        }
4386        kv.len = pos0 + t;
4387        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4388        Ok(())
4389    }
4390
4391    #[allow(clippy::too_many_arguments)]
4392    fn mtp_kv_fill_all(
4393        &self,
4394        e: &Engine,
4395        tokens: &[u32],
4396        h: &CudaSlice<f32>,
4397        pos0: usize,
4398        scratch: &mut MtpScratch,
4399        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4400    ) -> Result<(), Box<dyn std::error::Error>> {
4401        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
4402        for index in 0..self.mtp_head_count() {
4403            self.mtp_kv_fill_at(
4404                e,
4405                self.mtp_head_at(index),
4406                tokens,
4407                h,
4408                pos0,
4409                scratch,
4410                index,
4411                embd_dev,
4412            )?;
4413        }
4414        Ok(())
4415    }
4416
4417    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
4418    /// every varying input device-resident —
4419    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
4420    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
4421    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
4422    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
4423    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
4424    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
4425    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
4426    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
4427    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
4428    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
4429    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
4430    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
4431    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
4432    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
4433    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
4434    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
4435    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
4436    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
4437    #[allow(clippy::too_many_arguments)]
4438    fn mtp_head_forward_cap(
4439        &self,
4440        e: &Engine,
4441        mtp: &MtpHead,
4442        tok_d: &mut CudaSlice<u32>,
4443        pos_d: &mut CudaSlice<i32>,
4444        h_seed_d: &mut CudaSlice<f32>,
4445        p_d: &mut CudaSlice<f32>,
4446        scratch: &mut MtpScratch,
4447        with_prob: bool,
4448        with_head: bool,
4449        embd_gpu: &CudaSlice<u8>,
4450        embd_qt: i32,
4451        embd_rb: usize,
4452        d_vocab: usize,
4453        sampled_cap: Option<(
4454            &mut CudaSlice<u32>,
4455            &mut CudaSlice<f32>,
4456            &mut CudaSlice<f32>,
4457            u64,
4458            f32,
4459        )>,
4460        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
4461        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
4462        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
4463        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
4464        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
4465        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
4466        mask_cap: Option<(&CudaSlice<u32>, usize)>,
4467    ) -> Result<(), Box<dyn std::error::Error>> {
4468        let cfg = &self.cfg;
4469        let n_embd = cfg.n_embd as usize;
4470        // step35 REFUSAL (deliberate, named): the graph body's attention is `mtp_full_attn_dc`,
4471        // whose device-counter key bound always starts at row 0 — it cannot express this block's
4472        // SWA view offset, so a captured chain would silently attend OUTSIDE the window once the
4473        // persistent scratch passes 512 rows. Nothing is lost today: `mtp_head_forward_cap`
4474        // refuses step35 heads explicitly (SWA refusal), so the eager chain
4475        // (`mtp_head_forward_dev` -> `mtp_step35_attn`) is the served path. Returning Err (not a
4476        // panic) is what the two capture sites and the round-stream capture already handle by
4477        // degrading to eager / stream-off.
4478        if mtp.step35.is_some() {
4479            return Err(
4480                "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
4481                        block's SWA view offset; same root cause as the dc decode refusal) — the \
4482                        eager draft chain serves this arch"
4483                    .into(),
4484            );
4485        }
4486        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
4487        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4488        let eps = cfg.rms_eps;
4489        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
4490        let mut e_norm = e.zeros(n_embd)?;
4491        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4492        let mut h_norm = e.zeros(n_embd)?;
4493        e.rms_norm(
4494            &*h_seed_d,
4495            mtp.hnorm.float_data(),
4496            &mut h_norm,
4497            n_embd,
4498            1,
4499            eps,
4500        )?;
4501        let mut concat = e.zeros(2 * n_embd)?;
4502        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4503        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4504        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4505        let mut a_norm = e.zeros(di)?;
4506        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4507        let attn_out = match &mtp.mixer {
4508            Mixer::Full(fa) => {
4509                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
4510            }
4511            Mixer::Linear(_) => {
4512                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4513            }
4514            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
4515        };
4516        let mut x1 = e.zeros(di)?;
4517        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4518        let mut z = e.zeros(di)?;
4519        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4520        let ffn_out = match &mtp.ffn {
4521            crate::hybrid::Ffn::Dense {
4522                ffn_gate,
4523                ffn_up,
4524                ffn_down,
4525            } => {
4526                let n_ff = ffn_gate.out_features();
4527                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4528                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4529                    (
4530                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4531                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4532                    )
4533                } else {
4534                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4535                };
4536                let mut act = e.zeros(n_ff)?;
4537                Self::ffn_act(e, &self.cfg, &gate, &up, &mut act, n_ff)?;
4538                e.matmul(ffn_down, &act, 1)?
4539            }
4540            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
4541            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
4542            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
4543            // error arm degrades the caller to eager/stream-off.
4544            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
4545                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
4546            }
4547            crate::hybrid::Ffn::Moe(_) => {
4548                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
4549            }
4550        };
4551        let mut h_inner = e.zeros(di)?;
4552        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4553        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
4554        let h_nextn = match mtp.geom.as_ref() {
4555            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4556            None => h_inner,
4557        };
4558        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
4559        let final_h = if with_head || spec_hpost() {
4560            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4561            let mut fh = e.zeros(n_embd)?;
4562            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
4563            Some(fh)
4564        } else {
4565            None
4566        };
4567        if with_head {
4568            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4569            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
4570            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
4571            // before the argmax — proposals become legal by construction. Contents-only
4572            // per-replay upload keeps the capture valid.
4573            if let Some((mask_d, mw)) = mask_cap {
4574                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4575            }
4576            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
4577                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
4578                // own buffer is pool-recycled after the capture body returns, so it can't be the
4579                // retention target), bump the device event counter, gumbel-perturb reading it,
4580                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
4581                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
4582                e.sctr_inc(ctr_d)?;
4583                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
4584                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
4585                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
4586                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
4587                if with_prob {
4588                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4589                }
4590            } else {
4591                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
4592                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
4593                // p-min under a draft mask reads the MASKED row: confidence relative to the
4594                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
4595                // is the right semantics for "does the drafter know what comes next here" and
4596                // the same row the pick came from. Draft-quality only — verify arbitrates.
4597                if with_prob {
4598                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
4599                }
4600            }
4601        }
4602        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
4603        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
4604        if let Some((out, slot, d2t)) = stream_pack {
4605            e.pack_tok_p(tok_d, p_d, out, slot)?;
4606            if let Some(map) = d2t {
4607                e.tok_map_u32(tok_d, map)?;
4608            }
4609        }
4610        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
4611        if spec_hpost() {
4612            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
4613        } else {
4614            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
4615        }
4616        // advance the draft rope position in-graph.
4617        e.inc_seqlen(pos_d)?;
4618        Ok(())
4619    }
4620
4621    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
4622    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
4623    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
4624    /// Advances `cache.pos` by T.
4625    pub fn decode_step_t(
4626        &self,
4627        e: &Engine,
4628        tokens: &[u32],
4629        pos0: usize,
4630        cache: &mut Cache,
4631    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
4632        if self.is_gemma4_e4b() {
4633            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
4634        }
4635        if self.gemma_batch_program() {
4636            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
4637        }
4638        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
4639    }
4640
4641    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
4642    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
4643    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
4644    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
4645    pub fn decode_step_t_h(
4646        &self,
4647        e: &Engine,
4648        tokens: &[u32],
4649        pos0: usize,
4650        cache: &mut Cache,
4651    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4652        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
4653    }
4654
4655    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
4656    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
4657    pub fn decode_step_t_h_emb(
4658        &self,
4659        e: &Engine,
4660        tokens: &[u32],
4661        pos0: usize,
4662        cache: &mut Cache,
4663        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4664    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4665        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
4666        Ok((e.dtoh(&logits_d)?, h_seed))
4667    }
4668
4669    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
4670    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
4671    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
4672    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
4673    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
4674    pub fn decode_step_t_h_emb_dev(
4675        &self,
4676        e: &Engine,
4677        tokens: &[u32],
4678        pos0: usize,
4679        cache: &mut Cache,
4680        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4681    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4682        let n_embd = self.cfg.n_embd as usize;
4683        let t = tokens.len();
4684        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
4685        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
4686        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
4687        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
4688        Ok((logits, hs))
4689    }
4690
4691    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
4692    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
4693    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
4694    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
4695    /// retains/copies — they never change what any kernel computes).
4696    fn decode_step_t_core(
4697        &self,
4698        e: &Engine,
4699        tokens: &[u32],
4700        pos0: usize,
4701        cache: &mut Cache,
4702        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4703        mut ckpt: Option<&mut VerifyCkpt>,
4704    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4705        self.decode_step_t_core_stream(
4706            e,
4707            tokens,
4708            pos0,
4709            cache,
4710            embd_dev,
4711            ckpt.take(),
4712            None,
4713            None,
4714            None,
4715            None,
4716        )
4717    }
4718
4719    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
4720    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
4721    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
4722    fn decode_step_t_core_vg(
4723        &self,
4724        e: &Engine,
4725        tokens: &[u32],
4726        pos0: usize,
4727        cache: &mut Cache,
4728        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4729        mut ckpt: Option<&mut VerifyCkpt>,
4730        graphs: Option<&mut DsparkVerifyGraphs>,
4731    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4732        self.decode_step_t_core_stream(
4733            e,
4734            tokens,
4735            pos0,
4736            cache,
4737            embd_dev,
4738            ckpt.take(),
4739            None,
4740            None,
4741            None,
4742            graphs,
4743        )
4744    }
4745
4746    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
4747    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
4748    fn decode_step_t_core_pipelined(
4749        &self,
4750        e: &Engine,
4751        tokens: &[u32],
4752        pos0: usize,
4753        cache: &mut Cache,
4754        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4755        mut ckpt: Option<&mut VerifyCkpt>,
4756        pipe: &SpecPipeLane,
4757        round: usize,
4758    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4759        let fence = crate::pp::pp_cuts(self.layers.len())
4760            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
4761        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
4762            return Err("two-session speculative pipeline requires the PP verify split".into());
4763        }
4764        let interval_fence = pipe.stage0_begin(round)?;
4765        let ticket = self.verify_stage0_issue(
4766            e,
4767            tokens,
4768            pos0,
4769            cache,
4770            embd_dev,
4771            ckpt.as_deref_mut(),
4772            None,
4773            &fence,
4774            Some(interval_fence),
4775            pipe.trace(round),
4776        )?;
4777        pipe.stage0_end(round);
4778        pipe.stage1_begin(round)?;
4779        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
4780        pipe.verify_end(round);
4781        Ok(result)
4782    }
4783
4784    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
4785    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
4786    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
4787    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
4788    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
4789    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
4790    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
4791    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
4792    #[allow(clippy::too_many_arguments)]
4793    fn decode_step_t_core_stream(
4794        &self,
4795        e: &Engine,
4796        tokens: &[u32],
4797        pos0: usize,
4798        cache: &mut Cache,
4799        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4800        mut ckpt: Option<&mut VerifyCkpt>,
4801        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
4802        pp_pipe: Option<bool>,
4803        vtok_dev: Option<&CudaSlice<u32>>,
4804        graphs: Option<&mut DsparkVerifyGraphs>,
4805    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4806        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
4807        // exactly as the eager and batched steps do. This is the single funnel every verify
4808        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
4809        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
4810        // is untouched.
4811        //
4812        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
4813        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
4814        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
4815        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
4816        // or a placement whose PpNRt fails to build — so a config that would still walk the
4817        // whole trunk on one stream refuses instead of regressing 28x.
4818        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4819            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
4820                if vtok_dev.is_some() {
4821                    return Err(
4822                        "device-token dspark verify (slice-2 deferred readback) has no PP \
4823                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
4824                         route on one device"
4825                            .into(),
4826                    );
4827                }
4828                return self.decode_step_t_core_ppn(
4829                    e,
4830                    tokens,
4831                    pos0,
4832                    cache,
4833                    embd_dev,
4834                    ckpt.take(),
4835                    stream,
4836                    &fence,
4837                    pp_pipe,
4838                );
4839            }
4840        }
4841        crate::pp::refuse_unsplit_if_remote(
4842            "decode_step_t (spec verify)",
4843            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
4844             split (decode_step_t_core_ppn); or run spec on one device",
4845        )?;
4846        let cfg = &self.cfg;
4847        let n_embd = cfg.n_embd as usize;
4848        let eps = cfg.rms_eps;
4849        let t = tokens.len();
4850        let pos_d = match stream {
4851            Some((_, ctr)) => {
4852                let mut p = e.alloc_uninit::<i32>(t)?;
4853                e.pos_iota(ctr, &mut p, t)?;
4854                p
4855            }
4856            None => {
4857                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
4858                e.htod_i32(&pos_vec)?
4859            }
4860        };
4861
4862        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
4863        let x = match (stream, embd_dev) {
4864            (Some((vtok, _)), Some((g, qt, rb))) => {
4865                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
4866            }
4867            (None, Some((g, qt, rb))) => match vtok_dev {
4868                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
4869                // bit-identical rows to the host-token arm (same per-dtype deq).
4870                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
4871                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4872            },
4873            _ => {
4874                assert!(
4875                    vtok_dev.is_none(),
4876                    "device-token verify requires the resident embed table (embd_dev)"
4877                );
4878                e.htod(&self.embd.gather(n_embd, tokens))?
4879            }
4880        };
4881
4882        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
4883        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
4884        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
4885        let x = self.verify_layers(
4886            e,
4887            x,
4888            0,
4889            self.layers.len(),
4890            &pos_d,
4891            pos0,
4892            t,
4893            cache,
4894            ckpt.take(),
4895            stream,
4896            graphs,
4897        )?;
4898
4899        let mut hn = vbuf(e, t * n_embd)?;
4900        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
4901        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
4902        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
4903        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
4904        let eager_tail = self.sliding_gated_moe_batch_program()
4905            && std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1");
4906        if eager_tail {
4907            let n_vocab = self.cfg.n_vocab as usize;
4908            let mut logits = vbuf(e, t * n_vocab)?;
4909            for r in 0..t {
4910                let mut row = e.uninit(n_embd)?;
4911                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
4912                let mut hr = e.uninit(n_embd)?;
4913                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
4914                let lr = e.matmul(&self.output, &hr, 1)?;
4915                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
4916                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
4917            }
4918            if stream.is_none() {
4919                cache.pos += t;
4920            }
4921            return Ok((logits, if spec_hpost() { hn } else { x }));
4922        }
4923        let serving_head =
4924            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
4925        let logits = if serving_head {
4926            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
4927            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
4928            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
4929            // serve one batched numeric class at every live width, including B=1. Keep the
4930            // verify head in that same class; other generic families retain the decode-exact
4931            // head that their run-spec contract pins.
4932            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4933            e.matmul(&self.output, &hn, t)?
4934        } else {
4935            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
4936            e.matmul_decode_exact(&self.output, &hn, t)?
4937        };
4938        // stream: the device pos counter owns position; host mirror reconciles at drain.
4939        if stream.is_none() {
4940            cache.pos += t;
4941        }
4942        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
4943        Ok((logits, if spec_hpost() { hn } else { x }))
4944    }
4945
4946    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
4947    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
4948    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
4949    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
4950    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
4951    /// the payload).
4952    ///
4953    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
4954    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
4955    /// receipts):
4956    ///
4957    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
4958    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
4959    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
4960    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
4961    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
4962    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
4963    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
4964    ///
4965    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
4966    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
4967    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
4968    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
4969    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
4970    ///
4971    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
4972    ///    sharded loader leaves the table with stage 0 by construction).
4973    ///
4974    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
4975    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
4976    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
4977    ///    model, every round.
4978    ///
4979    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
4980    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
4981    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
4982    /// through the primary context by UVA — the same read the batched serving epilogue's
4983    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
4984    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
4985    ///
4986    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
4987    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
4988    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
4989    ///
4990    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
4991    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
4992    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
4993    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
4994    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
4995    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
4996    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
4997    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
4998    #[allow(clippy::too_many_arguments)]
4999    fn decode_step_t_core_ppn(
5000        &self,
5001        e: &Engine,
5002        tokens: &[u32],
5003        pos0: usize,
5004        cache: &mut Cache,
5005        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5006        mut ckpt: Option<&mut VerifyCkpt>,
5007        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5008        fence: &[usize],
5009        pp_pipe: Option<bool>,
5010    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5011        let ticket = self.verify_stage0_issue(
5012            e,
5013            tokens,
5014            pos0,
5015            cache,
5016            embd_dev,
5017            ckpt.as_deref_mut(),
5018            stream,
5019            fence,
5020            pp_pipe,
5021            None,
5022        )?;
5023        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5024    }
5025
5026    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5027    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5028    #[allow(clippy::too_many_arguments)]
5029    fn verify_stage0_issue(
5030        &self,
5031        e: &Engine,
5032        tokens: &[u32],
5033        pos0: usize,
5034        cache: &mut Cache,
5035        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5036        mut ckpt: Option<&mut VerifyCkpt>,
5037        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5038        fence: &[usize],
5039        pp_pipe: Option<bool>,
5040        trace: Option<SpecPipeTraceCtx>,
5041    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5042        assert!(
5043            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5044            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5045             (the gemma4 arms have their own decode_step_t twins)"
5046        );
5047        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5048            return Err(
5049                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5050                 boundary itself is host-staged, but device-resident verify still peer-reads \
5051                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5052                 serving on this host class; spec requires local per-stage inputs first."
5053                    .into(),
5054            );
5055        }
5056        let rt = crate::pp::PpNRt::get(e)?;
5057        let n_st = fence.len() - 1;
5058        assert_eq!(
5059            rt.n_stages(),
5060            n_st,
5061            "PpNRt stage count {} != fence stages {n_st}",
5062            rt.n_stages()
5063        );
5064        let n_embd = self.cfg.n_embd as usize;
5065        let t = tokens.len();
5066        let payload = t * n_embd;
5067        if pp_pipe.is_some() {
5068            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5069        }
5070        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5071        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5072        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5073        // the report below names exactly two stages and must never imply it measured middle ones.
5074        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5075        let pp_started = std::time::Instant::now();
5076        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5077        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5078        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5079        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5080        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5081        // stage stream and the wait would self-order into a no-op.
5082        let caller_stream = e.stream();
5083        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5084        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5085        // the primary stream still holds queued reads of them — with event tracking elided,
5086        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5087        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5088        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5089        // stage stream behind the caller before enqueueing new stage work.
5090        let reverse_started = std::time::Instant::now();
5091        if pp_pipe != Some(false) {
5092            rt.fence_stages_behind(&caller_stream)?;
5093        }
5094        if pp_pipe == Some(true) {
5095            // Both session verifies must alternate boundary slots even when the ordinary
5096            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5097            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5098            rt.prepare_overlap_slots(0, payload)?;
5099        }
5100        if pp_anatomy {
5101            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5102            // prices any primary-stream rollback/refresh tail inherited from the prior round.
5103            for s in 0..n_st {
5104                let _st = rt.enter(s);
5105                rt.engine(s, e).stream().synchronize()?;
5106            }
5107            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5108        }
5109
5110        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5111        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5112        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5113            match stream {
5114                Some((_, ctr)) => {
5115                    let mut p = es.alloc_uninit::<i32>(t)?;
5116                    es.pos_iota(ctr, &mut p, t)?;
5117                    Ok(p)
5118                }
5119                None => {
5120                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5121                    es.htod_i32(&pos_vec)
5122                }
5123            }
5124        };
5125
5126        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5127        let slot = {
5128            let _st0 = rt.enter(0);
5129            let e0 = rt.engine(0, e);
5130            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5131            let stage0_started = std::time::Instant::now();
5132            let pos_d = stage_pos(e0)?;
5133            let x = match (stream, embd_dev) {
5134                (Some((vtok, _)), Some((g, qt, rb))) => {
5135                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5136                }
5137                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5138                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5139            };
5140            let x = self.verify_layers(
5141                e0,
5142                x,
5143                fence[0],
5144                fence[1],
5145                &pos_d,
5146                pos0,
5147                t,
5148                cache,
5149                ckpt.as_deref_mut(),
5150                stream,
5151                None,
5152            )?;
5153            if pp_anatomy {
5154                e0.stream().synchronize()?;
5155                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5156            }
5157            let tx_started = std::time::Instant::now();
5158            let slot = if pp_pipe.is_some() {
5159                rt.tx_pipelined(0, &x, payload)?
5160            } else {
5161                rt.tx(0, &x, payload)?
5162            };
5163            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5164            if pp_anatomy {
5165                e0.stream().synchronize()?;
5166                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5167            }
5168            slot
5169            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5170        };
5171
5172        Ok(VerifyBoundaryTicket {
5173            rt,
5174            caller_stream,
5175            slot,
5176            pos0,
5177            t,
5178            payload,
5179            n_st,
5180            pipelined: pp_pipe.is_some(),
5181            pp_anatomy,
5182            pp_started,
5183            reverse_ms,
5184            stage0_ms,
5185            tx_ms,
5186            trace,
5187        })
5188    }
5189
5190    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5191    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5192    #[allow(clippy::too_many_arguments)]
5193    fn verify_stage1_finish(
5194        &self,
5195        e: &Engine,
5196        ticket: VerifyBoundaryTicket,
5197        cache: &mut Cache,
5198        mut ckpt: Option<&mut VerifyCkpt>,
5199        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5200        fence: &[usize],
5201        publish_to_caller: bool,
5202    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5203        let VerifyBoundaryTicket {
5204            rt,
5205            caller_stream,
5206            slot,
5207            pos0,
5208            t,
5209            payload,
5210            n_st,
5211            pipelined,
5212            pp_anatomy,
5213            pp_started,
5214            reverse_ms,
5215            stage0_ms,
5216            tx_ms,
5217            trace,
5218        } = ticket;
5219        let n_embd = self.cfg.n_embd as usize;
5220        let eps = self.cfg.rms_eps;
5221        let mut slot = slot;
5222        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5223        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5224            match stream {
5225                Some((_, ctr)) => {
5226                    let mut p = es.alloc_uninit::<i32>(t)?;
5227                    es.pos_iota(ctr, &mut p, t)?;
5228                    Ok(p)
5229                }
5230                None => {
5231                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5232                    es.htod_i32(&pos_vec)
5233                }
5234            }
5235        };
5236
5237        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5238        for s in 1..n_st - 1 {
5239            let _st = rt.enter(s);
5240            let es = rt.engine(s, e);
5241            let pos_d = stage_pos(es)?;
5242            let x = rt.rx(s - 1, slot, payload)?;
5243            let x = self.verify_layers(
5244                es,
5245                x,
5246                fence[s],
5247                fence[s + 1],
5248                &pos_d,
5249                pos0,
5250                t,
5251                cache,
5252                ckpt.as_deref_mut(),
5253                stream,
5254                None,
5255            )?;
5256            slot = if pipelined {
5257                rt.tx_pipelined(s, &x, payload)?
5258            } else {
5259                rt.tx(s, &x, payload)?
5260            };
5261        }
5262
5263        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5264        let _stl = rt.enter(n_st - 1);
5265        let el = rt.engine(n_st - 1, e);
5266        let pos_d = stage_pos(el)?;
5267        let rx_started = std::time::Instant::now();
5268        let x = rt.rx(n_st - 2, slot, payload)?;
5269        if pp_anatomy {
5270            el.stream().synchronize()?;
5271            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5272        }
5273        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5274        let stage1_started = std::time::Instant::now();
5275        let x = self.verify_layers(
5276            el,
5277            x,
5278            fence[n_st - 1],
5279            fence[n_st],
5280            &pos_d,
5281            pos0,
5282            t,
5283            cache,
5284            ckpt.as_deref_mut(),
5285            stream,
5286            None,
5287        )?;
5288
5289        let mut hn = vbuf(el, payload)?;
5290        let logits = if self.sliding_gated_moe_batch_program() {
5291            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
5292            // Verify must not switch numeric class merely because the same session speculates.
5293            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5294            el.matmul(&self.output, &hn, t)?
5295        } else {
5296            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5297            el.matmul_decode_exact(&self.output, &hn, t)?
5298        };
5299        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
5300        if pp_anatomy {
5301            el.stream().synchronize()?;
5302            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
5303        }
5304        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
5305        // stream. Order the caller's stream behind that work before the buffers escape this
5306        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
5307        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
5308        // the following arm's KV in the same process).
5309        if publish_to_caller {
5310            rt.publish_to(n_st - 1, &caller_stream)?;
5311        }
5312        if pp_anatomy {
5313            if publish_to_caller {
5314                caller_stream.synchronize()?;
5315            }
5316            eprintln!(
5317                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
5318                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
5319                pp_started.elapsed().as_secs_f64() * 1e3,
5320            );
5321        }
5322        // stream: the device pos counter owns position; host mirror reconciles at drain.
5323        if stream.is_none() {
5324            cache.pos += t;
5325        }
5326        Ok((logits, if spec_hpost() { hn } else { x }))
5327    }
5328
5329    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
5330    ///
5331    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
5332    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
5333    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
5334    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
5335    /// bytes when a request moves from batched plain serving into speculative verify. Run the
5336    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
5337    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
5338    /// every norm/projection/FFN uses exactly the live serving dispatch.
5339    #[allow(clippy::too_many_arguments)]
5340    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
5341    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
5342    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
5343    /// reference while replacing the host-canonical per-token prime. Requires the walk
5344    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
5345    #[allow(clippy::type_complexity)]
5346    pub(crate) fn step35_prime_trows(
5347        &self,
5348        e: &Engine,
5349        tokens: &[u32],
5350        cache: &mut Cache,
5351    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
5352    {
5353        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
5354        if std::env::var("MEMRA_PRIME_TROWS").as_deref() != Ok("1") {
5355            return Ok(None);
5356        }
5357        if !self.uses_sliding_gated_moe_program()
5358            || cache.pos != 0
5359            || cache.dflash_taps.is_some()
5360            || std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() != Ok("1")
5361            || std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() != Ok("1")
5362        {
5363            if dbg {
5364                eprintln!(
5365                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
5366                    self.uses_sliding_gated_moe_program(),
5367                    cache.pos,
5368                    cache.dflash_taps.is_some(),
5369                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
5370                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
5371                );
5372            }
5373            return Ok(None);
5374        }
5375        let n_embd = self.cfg.n_embd as usize;
5376        let n_layers = self.layers.len();
5377        let t_total = tokens.len();
5378        let Some(embd_gpu) = self.embd_gpu_try(e) else {
5379            if dbg {
5380                eprintln!("[prime-trows] refuse: no device embed table");
5381            }
5382            return Ok(None);
5383        };
5384        let embd_qtype = match self.embd.ggml_type {
5385            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
5386            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
5387            other => {
5388                if dbg {
5389                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
5390                }
5391                return Ok(None);
5392            }
5393        };
5394        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
5395        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
5396        // (the walk floor is t >= 2).
5397        let mut bounds = Vec::new();
5398        let mut start = 0usize;
5399        while start < t_total {
5400            let mut end = (start + 32).min(t_total);
5401            if t_total - end == 1 {
5402                end -= 1;
5403            }
5404            bounds.push((start, end));
5405            start = end;
5406        }
5407        if bounds.iter().any(|(a, b)| b - a < 2) {
5408            return Ok(None); // degenerate short prompt keeps the ordinary prime
5409        }
5410        let mut hiddens = e.uninit(t_total * n_embd)?;
5411        let mut last: Option<CudaSlice<f32>> = None;
5412        for &(a, b) in &bounds {
5413            let tc = b - a;
5414            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
5415            let x =
5416                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
5417            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
5418            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
5419            if b == t_total {
5420                let mut h = e.uninit(n_embd)?;
5421                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
5422                last = Some(h);
5423            }
5424        }
5425        let h_seed = last.expect("last chunk produced the seed row");
5426        let mut hn = e.uninit(n_embd)?;
5427        e.rms_norm_decode(
5428            &h_seed,
5429            self.output_norm.float_data(),
5430            &mut hn,
5431            n_embd,
5432            1,
5433            self.cfg.rms_eps,
5434        )?;
5435        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
5436        let logits = e.dtoh(&logits_d)?;
5437        cache.pos = t_total;
5438        Ok(Some((logits, h_seed, hiddens)))
5439    }
5440
5441    fn step35_verify_batch_layers(
5442        &self,
5443        e: &Engine,
5444        mut x: CudaSlice<f32>,
5445        lo: usize,
5446        hi: usize,
5447        pos0: usize,
5448        t: usize,
5449        cache: &mut Cache,
5450    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5451        let n_embd = self.cfg.n_embd as usize;
5452        if !self.uses_sliding_gated_moe_program() {
5453            return Err(
5454                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
5455            );
5456        }
5457        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
5458        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
5459        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
5460        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
5461        // and the tap path keep the batch-layer class.
5462        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5463        let eager_verify = *VE
5464            .get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_EAGER").as_deref() == Ok("1"))
5465            && lo == 0
5466            && hi == self.layers.len();
5467        if eager_verify {
5468            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
5469            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
5470            // column runs the UNMODIFIED t=1 attention program via the col-select door and
5471            // the ordinary residual/FFN body. Values per column are bit-equal to the
5472            // row-outer walk: rms over the materialized residual == the fused add+norm
5473            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
5474            // kernel, and every downstream op IS the t=1 program.
5475            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5476            let tcol =
5477                *TCOL.get_or_init(|| std::env::var("MEMRA_SPEC_VERIFY_TCOL").as_deref() == Ok("1"));
5478            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
5479            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
5480            // so a chunked call is value-identical to the row-outer loop it replaces.
5481            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5482            let trows_prefill = *TROWS_PREFILL
5483                .get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"));
5484            // MEMRA_PRIME_TROWS_T=<w>: chunk width (default 32, the walk's t cap). A
5485            // narrower width isolates slab-width faults from the chunking itself.
5486            static TROWS_W: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
5487            let trows_w = *TROWS_W.get_or_init(|| {
5488                std::env::var("MEMRA_PRIME_TROWS_T")
5489                    .ok()
5490                    .and_then(|v| v.parse::<usize>().ok())
5491                    .filter(|w| (2..=32).contains(w))
5492                    .unwrap_or(32)
5493            });
5494            if tcol && trows_prefill && t > trows_w {
5495                // One-time engagement receipt: without it a prefill gate cannot tell a
5496                // chunked walk from the row-outer fallback it is supposed to replace
5497                // (the first PRIME_TROWS gate passed vacuously on exactly that).
5498                static SEEN: std::sync::atomic::AtomicBool =
5499                    std::sync::atomic::AtomicBool::new(false);
5500                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
5501                    eprintln!(
5502                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
5503                        t.div_ceil(trows_w),
5504                        lo,
5505                        hi
5506                    );
5507                }
5508                let mut out = e.uninit(t * n_embd)?;
5509                let mut start = 0usize;
5510                while start < t {
5511                    let mut end = (start + trows_w).min(t);
5512                    if t - end == 1 {
5513                        end -= 1;
5514                    }
5515                    let tc = end - start;
5516                    let mut xc = e.uninit(tc * n_embd)?;
5517                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
5518                    let oc =
5519                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
5520                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
5521                    start = end;
5522                }
5523                return Ok(out);
5524            }
5525            if tcol && t >= 2 && t <= 32 {
5526                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
5527                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
5528                // syncs serialize the stream, so the split is for TARGETING amortization
5529                // work only — never a perf claim.
5530                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5531                let prof =
5532                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
5533                let mut prof_ms = [0f64; 3];
5534                let eps = self.cfg.rms_eps;
5535                let mut x_t = x;
5536                let mut h_t = e.uninit(t * n_embd)?;
5537                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
5538                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
5539                // pageable htod was an in-stream engine turnaround x t x 45).
5540                let mut pos_rows = Vec::with_capacity(t);
5541                for r in 0..t {
5542                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
5543                }
5544                let mut ok = true;
5545                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
5546                // stashes `gated` instead of joining per column; one b4_tcol per rank +
5547                // one slab join produce every column's `mixed` after the attention pass.
5548                // Bit-exact per column (t=1 b4 program per column; elementwise join).
5549                // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
5550                // MoE layer deferred, the residual norm runs as one t-grid launch
5551                // (per-row program == t=1) and the FFN as ONE two-column device-routed
5552                // sweep + per-column shexp — the two columns' expert weights dedup
5553                // through L2 instead of reading HBM twice.
5554                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5555                let ffn_batch =
5556                    *FFN2.get_or_init(|| std::env::var("MEMRA_TCOL_FFN").as_deref() == Ok("1"));
5557                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
5558                // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
5559                // the per-column pass norms/ropes/appends and stashes q+gate, then one
5560                // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
5561                // [2, o_out] mixed slab. The precheck runs before arming (stashing is
5562                // unrecoverable); ineligible/boundary layers run the ordinary program.
5563                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
5564                let mut mixed_row = e.uninit(n_embd)?;
5565                let mut pos_staged = false;
5566                for il in lo..hi {
5567                    let layer = &self.layers[il];
5568                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
5569                    let mut seg = std::time::Instant::now();
5570                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
5571                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
5572                        ok = false;
5573                        break;
5574                    }
5575                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
5576                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
5577                    // advance by t. Host cache bookkeeping mirrors the per-column tail.
5578                    if fa2_layer {
5579                        if let Some(mixed_t) =
5580                            self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
5581                        {
5582                            pos_staged = true;
5583                            {
5584                                let tp_kv = cache.tp_kv[il]
5585                                    .as_mut()
5586                                    .expect("precheck verified the distributed cache");
5587                                let transaction = tp_kv.begin_transaction()?;
5588                                let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
5589                                    return Err("verify rope pass expects full attention".into());
5590                                };
5591                                let tp = fa
5592                                    .step_tp_qkv
5593                                    .as_ref()
5594                                    .ok_or("verify rope pass lost its TP state")?;
5595                                let empty: [CudaSlice<f32>; 0] = [];
5596                                tp.runtime.append_tp_kv_transaction_inner(
5597                                    tp_kv,
5598                                    transaction,
5599                                    &empty,
5600                                    &empty,
5601                                    t,
5602                                    true,
5603                                )?;
5604                                tp.runtime.commit_tp_kv_transaction_external(
5605                                    tp_kv,
5606                                    transaction,
5607                                    t,
5608                                )?;
5609                                if let Some(local) = cache.kv[il].as_mut() {
5610                                    local.len = pos0 + t;
5611                                    if !crate::tp::len_mirror_lazy_on() {
5612                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
5613                                    }
5614                                }
5615                            }
5616                            if prof {
5617                                e.stream().synchronize()?;
5618                                prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5619                                seg = std::time::Instant::now();
5620                            }
5621                            let o_out = mixed_t.len() / t;
5622                            let mut next = e.uninit(t * n_embd)?;
5623                            let mut batched = false;
5624                            if ffn_batch && o_out == n_embd {
5625                                let mut x1_t = e.uninit(t * n_embd)?;
5626                                let mut z_t = e.uninit(t * n_embd)?;
5627                                e.add_rms_norm(
5628                                    &x_t,
5629                                    &mixed_t,
5630                                    layer.post_attn_norm.float_data(),
5631                                    &mut x1_t,
5632                                    &mut z_t,
5633                                    n_embd,
5634                                    t,
5635                                    eps,
5636                                )?;
5637                                if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
5638                                    let mut x2_t = e.uninit(t * n_embd)?;
5639                                    e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5640                                    next = x2_t;
5641                                    batched = true;
5642                                }
5643                            }
5644                            if !batched {
5645                                for r in 0..t {
5646                                    e.dtod_copy_view(
5647                                        &mixed_t.slice(r * o_out..(r + 1) * o_out),
5648                                        &mut mixed_row,
5649                                    )?;
5650                                    let mut x_row = e.uninit(n_embd)?;
5651                                    e.dtod_copy_view(
5652                                        &x_t.slice(r * n_embd..(r + 1) * n_embd),
5653                                        &mut x_row,
5654                                    )?;
5655                                    let (x1, ffn_out) = self.residual_norm_ffn(
5656                                        e, layer, &x_row, &mixed_row, n_embd, il, eps,
5657                                    )?;
5658                                    let mut x2 = e.uninit(n_embd)?;
5659                                    e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5660                                    e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
5661                                }
5662                            }
5663                            if prof {
5664                                e.stream().synchronize()?;
5665                                prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5666                            }
5667                            x_t = next;
5668                            continue;
5669                        }
5670                    }
5671                    if prof {
5672                        e.stream().synchronize()?;
5673                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
5674                        seg = std::time::Instant::now();
5675                    }
5676                    let mut next = e.uninit(t * n_embd)?;
5677                    // Columns whose o_proj was deferred (their FFN runs after the join).
5678                    // A NON-deferred column's FFN must run INSIDE the column loop: the
5679                    // oproj-tail handoff is a single cell that the same column's
5680                    // residual_norm_ffn consumes before the next column's finish.
5681                    let mut deferred: Vec<usize> = Vec::new();
5682                    let mut fa2_deferred: Vec<usize> = Vec::new();
5683                    let mut ffn_col =
5684                        |r: usize,
5685                         mixed: &CudaSlice<f32>,
5686                         next: &mut CudaSlice<f32>|
5687                         -> Result<(), Box<dyn std::error::Error>> {
5688                            let mut x_row = e.uninit(n_embd)?;
5689                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
5690                            let (x1, ffn_out) =
5691                                self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
5692                            let mut x2 = e.uninit(n_embd)?;
5693                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
5694                            e.dtod_copy_into(&x2, next, r * n_embd)?;
5695                            Ok(())
5696                        };
5697                    for r in 0..t {
5698                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
5699                        let row_pos = &pos_rows[r];
5700                        crate::tp::set_verify_tcol(Some(r));
5701                        if fa2_layer {
5702                            crate::tp::set_spec_fa2_defer(Some(r));
5703                        } else if oproj_batch {
5704                            crate::tp::set_tcol_oproj_defer(Some(r));
5705                        }
5706                        let mixed = match &layer.mixer {
5707                            crate::hybrid::Mixer::Full(fa) => {
5708                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
5709                            }
5710                            _ => Err("step35 verify expects full attention".into()),
5711                        };
5712                        crate::tp::set_verify_tcol(None);
5713                        crate::tp::set_spec_fa2_defer(None);
5714                        crate::tp::set_tcol_oproj_defer(None);
5715                        let mixed = mixed?;
5716                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
5717                            fa2_deferred.push(r);
5718                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
5719                            deferred.push(r);
5720                        } else {
5721                            ffn_col(r, &mixed, &mut next)?;
5722                        }
5723                    }
5724                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
5725                        // The precheck guarantees both columns stash or neither; a strict
5726                        // subset means a column's output was never produced anywhere.
5727                        return Err("spec fa2 stash engaged for a subset of columns".into());
5728                    }
5729                    if prof {
5730                        e.stream().synchronize()?;
5731                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
5732                        seg = std::time::Instant::now();
5733                    }
5734                    if !fa2_deferred.is_empty() {
5735                        deferred = fa2_deferred;
5736                    }
5737                    if !deferred.is_empty() {
5738                        let mixed_t = if fa2_layer {
5739                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
5740                        } else {
5741                            self.step35_verify_oproj_tcol(e, il, t)?
5742                        };
5743                        let o_out = mixed_t.len() / t;
5744                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
5745                        // program == t=1; bit-identical to the oproj-tail join per the
5746                        // M2 verbatim-program contract) feeding the two-column routed
5747                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
5748                        // to the per-column body.
5749                        let mut batched = false;
5750                        if ffn_batch && deferred.len() == t && o_out == n_embd {
5751                            let mut x1_t = e.uninit(t * n_embd)?;
5752                            let mut z_t = e.uninit(t * n_embd)?;
5753                            e.add_rms_norm(
5754                                &x_t,
5755                                &mixed_t,
5756                                layer.post_attn_norm.float_data(),
5757                                &mut x1_t,
5758                                &mut z_t,
5759                                n_embd,
5760                                t,
5761                                eps,
5762                            )?;
5763                            if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
5764                                let mut x2_t = e.uninit(t * n_embd)?;
5765                                e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
5766                                next = x2_t;
5767                                batched = true;
5768                            }
5769                        }
5770                        if !batched {
5771                            for &r in &deferred {
5772                                e.dtod_copy_view(
5773                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
5774                                    &mut mixed_row,
5775                                )?;
5776                                ffn_col(r, &mixed_row, &mut next)?;
5777                            }
5778                        }
5779                    }
5780                    if prof {
5781                        e.stream().synchronize()?;
5782                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
5783                    }
5784                    drop(ffn_col);
5785                    x_t = next;
5786                }
5787                if prof {
5788                    eprintln!(
5789                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
5790                        prof_ms[0], prof_ms[1], prof_ms[2]
5791                    );
5792                }
5793                if ok {
5794                    return Ok(x_t);
5795                }
5796                // fall through to the row-outer walk on ineligible layers
5797                x = x_t;
5798            }
5799            let mut next = e.uninit(t * n_embd)?;
5800            for r in 0..t {
5801                let mut row = e.uninit(n_embd)?;
5802                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5803                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5804                let out = self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?;
5805                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5806            }
5807            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
5808            // row-outer walk does not materialize); the door is a step37 MTP bring-up
5809            // surface where taps are unused.
5810            return Ok(next);
5811        }
5812        let mut ph_last = std::time::Instant::now();
5813        for il in lo..hi {
5814            let mut next = e.uninit(t * n_embd)?;
5815            for r in 0..t {
5816                let mut row = e.uninit(n_embd)?;
5817                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5818                // The caller owns this verify's position. During controller overlap, cache.pos
5819                // still describes generation N while this stage-0 walk belongs to N+1.
5820                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
5821                let mut one = [&mut *cache];
5822                let out = self.step35_decode_batch_layers(
5823                    e,
5824                    row,
5825                    &mut one,
5826                    &[(pos0 + r) as i32],
5827                    &row_pos,
5828                    il,
5829                    il + 1,
5830                    &mut ph_last,
5831                )?;
5832                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
5833            }
5834            self.dflash_tap(e, cache, il, &next, t)?;
5835            x = next;
5836        }
5837        Ok(x)
5838    }
5839
5840    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
5841    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
5842    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
5843    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
5844    /// prefix-keep, not all-or-nothing).
5845    pub(crate) fn dspark_verify_t_am(
5846        &self,
5847        e: &Engine,
5848        tokens: &[u32],
5849        pos0: usize,
5850        cache: &mut Cache,
5851    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
5852        let (logits, _hn) = self.decode_step_t_core_stream(
5853            e, tokens, pos0, cache, None, None, None, None, None, None,
5854        )?;
5855        let t = tokens.len();
5856        let v = self.output.out_features();
5857        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5858        for r in 0..t {
5859            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5860        }
5861        Ok(e.dtoh_u32(&am_d)?)
5862    }
5863
5864    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
5865    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
5866    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
5867    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
5868    pub(crate) fn dspark_verify_t_logits(
5869        &self,
5870        e: &Engine,
5871        tokens: &[u32],
5872        pos0: usize,
5873        cache: &mut Cache,
5874    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
5875        let (logits, _hn) = self.decode_step_t_core_stream(
5876            e, tokens, pos0, cache, None, None, None, None, None, None,
5877        )?;
5878        Ok(logits)
5879    }
5880
5881    /// DSpark verify with the MTP column-stash armed: identical forward to
5882    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
5883    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
5884    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
5885    pub(crate) fn dspark_verify_t_am_ckpt(
5886        &self,
5887        e: &Engine,
5888        tokens: &[u32],
5889        pos0: usize,
5890        cache: &mut Cache,
5891    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5892        let mut ck = VerifyCkpt::new(self.layers.len());
5893        let (logits, _hn) = self.decode_step_t_core_stream(
5894            e,
5895            tokens,
5896            pos0,
5897            cache,
5898            None,
5899            Some(&mut ck),
5900            None,
5901            None,
5902            None,
5903            None,
5904        )?;
5905        let t = tokens.len();
5906        let v = self.output.out_features();
5907        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5908        for r in 0..t {
5909            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5910        }
5911        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
5912    }
5913
5914    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
5915    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
5916    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
5917    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
5918    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
5919    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
5920    pub(crate) fn dspark_verify_t_am_ckpt_dev(
5921        &self,
5922        e: &Engine,
5923        vtok: &CudaSlice<u32>,
5924        t: usize,
5925        pos0: usize,
5926        cache: &mut Cache,
5927        embd_dev: (&CudaSlice<u8>, i32, usize),
5928        graphs: Option<&mut DsparkVerifyGraphs>,
5929    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5930        debug_assert!(
5931            vtok.len() >= t,
5932            "verify window exceeds the device token buffer"
5933        );
5934        // The slab flag is a per-round statement: clear it here so a verify that never
5935        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
5936        // stale `true` steering the commit at slabs the round never wrote.
5937        let mut graphs = graphs;
5938        if let Some(g) = graphs.as_deref_mut() {
5939            g.round_slab = false;
5940        }
5941        let mut ck = VerifyCkpt::new(self.layers.len());
5942        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
5943        // arm's established pattern — spec.rs stream-mode verify does the same).
5944        let dummy = vec![0u32; t];
5945        let (logits, _hn) = self.decode_step_t_core_stream(
5946            e,
5947            &dummy,
5948            pos0,
5949            cache,
5950            Some(embd_dev),
5951            Some(&mut ck),
5952            None,
5953            None,
5954            Some(vtok),
5955            graphs,
5956        )?;
5957        let v = self.output.out_features();
5958        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
5959        for r in 0..t {
5960            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
5961        }
5962        Ok((am_d, DsparkVerifyCkpt(ck)))
5963    }
5964
5965    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
5966    pub(crate) fn dspark_verify_t_logits_ckpt(
5967        &self,
5968        e: &Engine,
5969        tokens: &[u32],
5970        pos0: usize,
5971        cache: &mut Cache,
5972    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
5973        let mut ck = VerifyCkpt::new(self.layers.len());
5974        let (logits, _hn) = self.decode_step_t_core_stream(
5975            e,
5976            tokens,
5977            pos0,
5978            cache,
5979            None,
5980            Some(&mut ck),
5981            None,
5982            None,
5983            None,
5984            None,
5985        )?;
5986        Ok((logits, DsparkVerifyCkpt(ck)))
5987    }
5988
5989    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
5990    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
5991    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
5992    pub(crate) fn dspark_commit_prefix(
5993        &self,
5994        e: &Engine,
5995        cache: &mut Cache,
5996        snap: &crate::cache::CacheSnapshot,
5997        ckpt: &DsparkVerifyCkpt,
5998        keep: usize,
5999    ) -> Result<(), Box<dyn std::error::Error>> {
6000        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6001    }
6002
6003    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6004    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6005    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6006    /// from the stash of column keep-1), slab-addressed and batched into two copy
6007    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6008    pub(crate) fn dspark_commit_prefix_slab(
6009        &self,
6010        e: &Engine,
6011        cache: &mut Cache,
6012        snap: &crate::cache::CacheSnapshot,
6013        ctx: &DsparkVerifyGraphs,
6014        keep: usize,
6015    ) -> Result<(), Box<dyn std::error::Error>> {
6016        use cudarc::driver::DevicePtr;
6017        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6018        let mut conv_src: Vec<u64> = Vec::new();
6019        let mut ssm_src: Vec<u64> = Vec::new();
6020        let mut conv_dst: Vec<u64> = Vec::new();
6021        let mut ssm_dst: Vec<u64> = Vec::new();
6022        for il in 0..self.layers.len() {
6023            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6024                kvl.len = saved + keep;
6025                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6026            }
6027            if let Some(rl) = cache.recur[il].as_ref() {
6028                let (pc, ps, _cw, _sw) = ctx
6029                    .slab_row(e, il, keep - 1)
6030                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6031                conv_src.push(pc);
6032                ssm_src.push(ps);
6033                let st = &e.gpu.stream();
6034                let (dc, _g0) = rl.conv_state.device_ptr(st);
6035                let (ds, _g1) = rl.ssm_state.device_ptr(st);
6036                conv_dst.push(dc as u64);
6037                ssm_dst.push(ds as u64);
6038            }
6039        }
6040        let n = conv_src.len();
6041        if n > 0 {
6042            if state_copy_batch_on() {
6043                let mut tt = vec![0u64; 2 * n];
6044                tt[..n].copy_from_slice(&conv_src);
6045                tt[n..].copy_from_slice(&conv_dst);
6046                let ct = e.htod_u64(&tt)?;
6047                tt[..n].copy_from_slice(&ssm_src);
6048                tt[n..].copy_from_slice(&ssm_dst);
6049                let st = e.htod_u64(&tt)?;
6050                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6051                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6052            } else {
6053                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6054                let row = keep - 1;
6055                for il in 0..self.layers.len() {
6056                    let Some(rl) = cache.recur[il].as_mut() else {
6057                        continue;
6058                    };
6059                    let k = ctx.lin_pos[&il];
6060                    {
6061                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
6062                        let win = sv.slice(row * cw..(row + 1) * cw);
6063                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
6064                    }
6065                    {
6066                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
6067                        let win = sv.slice(row * sw..(row + 1) * sw);
6068                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
6069                    }
6070                }
6071            }
6072        }
6073        cache.pos = snap.pos + keep;
6074        Ok(())
6075    }
6076
6077    /// Qwen35-family verify trunk in the live serving numeric class.
6078    ///
6079    /// Serving intentionally keeps this architecture in the generic batched program even at
6080    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6081    ///
6082    /// Two arms, one numeric class:
6083    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6084    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6085    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6086    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6087    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
6088    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
6089    ///   program its isolated serving step would). One weight read per layer per round
6090    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
6091    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
6092    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
6093    ///   serving layer body, preserving single-session autoregressive cache order (the
6094    ///   correctness reference; also the rollback seam for the t-parallel arm).
6095    ///
6096    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
6097    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
6098    #[allow(clippy::too_many_arguments)]
6099    fn qwen35_verify_batch_layers(
6100        &self,
6101        e: &Engine,
6102        x: CudaSlice<f32>,
6103        lo: usize,
6104        hi: usize,
6105        pos0: usize,
6106        t: usize,
6107        cache: &mut Cache,
6108        ckpt: Option<&mut VerifyCkpt>,
6109        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6110        graphs: Option<&mut DsparkVerifyGraphs>,
6111    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6112        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
6113        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
6114        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
6115        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
6116        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
6117        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
6118        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
6119            || !self.batched_serving_numeric_class()
6120            || t > 16;
6121        if rowwise {
6122            if stream.is_some() {
6123                // rowwise replays per row with host cache.pos — irreconcilable with a
6124                // device position counter. Burst callers must keep t <= 16 and the
6125                // ROWWISE env unset; refusing beats silently mispositioned rows.
6126                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
6127                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
6128                    .into());
6129            }
6130            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
6131        } else {
6132            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
6133        }
6134    }
6135
6136    /// The per-row correctness reference: replay each verify row through the authoritative
6137    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
6138    #[allow(clippy::too_many_arguments)]
6139    fn qwen35_verify_rowwise(
6140        &self,
6141        e: &Engine,
6142        mut x: CudaSlice<f32>,
6143        lo: usize,
6144        hi: usize,
6145        pos0: usize,
6146        t: usize,
6147        cache: &mut Cache,
6148        mut ckpt: Option<&mut VerifyCkpt>,
6149    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6150        let n_embd = self.cfg.n_embd as usize;
6151        let saved_pos = cache.pos;
6152        let mut ph_last = std::time::Instant::now();
6153        for il in lo..hi {
6154            let mut next = e.uninit(t * n_embd)?;
6155            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6156                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6157                    Some(Vec::with_capacity(t - 1))
6158                } else {
6159                    None
6160                };
6161            for r in 0..t {
6162                cache.pos = pos0 + r;
6163                let mut row = e.uninit(n_embd)?;
6164                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6165                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6166                let mut one = [&mut *cache];
6167                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
6168                let out = match self.decode_batch_layers(
6169                    e,
6170                    row,
6171                    &mut one,
6172                    &ctx,
6173                    &row_pos,
6174                    &mut ph_last,
6175                ) {
6176                    Ok(out) => out,
6177                    Err(error) => {
6178                        cache.pos = saved_pos;
6179                        return Err(error);
6180                    }
6181                };
6182                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6183                if r + 1 < t {
6184                    if let Some(states) = col_states.as_mut() {
6185                        let recur = cache.recur[il]
6186                            .as_ref()
6187                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
6188                        states.push((
6189                            e.clone_dtod(&recur.conv_state)?,
6190                            e.clone_dtod(&recur.ssm_state)?,
6191                        ));
6192                    }
6193                }
6194            }
6195            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
6196                checkpoint.cols[il] = Some(states);
6197            }
6198            x = next;
6199        }
6200        cache.pos = saved_pos;
6201        Ok(x)
6202    }
6203
6204    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
6205    ///
6206    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
6207    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
6208    /// pins the serving batch tier already carries:
6209    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
6210    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
6211    ///     alone;
6212    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
6213    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
6214    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
6215    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
6216    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
6217    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
6218    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
6219    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
6220    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
6221    /// program its isolated B=1 serving step would.
6222    ///
6223    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
6224    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
6225    #[allow(clippy::too_many_arguments)]
6226    fn qwen35_verify_tparallel(
6227        &self,
6228        e: &Engine,
6229        mut x: CudaSlice<f32>,
6230        lo: usize,
6231        hi: usize,
6232        pos0: usize,
6233        t: usize,
6234        cache: &mut Cache,
6235        mut ckpt: Option<&mut VerifyCkpt>,
6236        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6237        mut graphs: Option<&mut DsparkVerifyGraphs>,
6238    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6239        let seqs_append =
6240            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
6241        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
6242
6243        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
6244        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
6245        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
6246        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
6247        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
6248        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
6249        // full-verify bodies).
6250        if stream.is_some() && graphs.is_some() {
6251            return Err(
6252                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
6253                        cannot arm together"
6254                    .into(),
6255            );
6256        }
6257        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
6258        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
6259        // moves the kv caches). Then:
6260        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
6261        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
6262        //    full-verify graph per (vt, rung) — linear layers through the shared
6263        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
6264        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
6265        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
6266        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
6267        //    the full-attention layers run eager (batched rows when eligible).
6268        if let Some(g) = graphs.as_deref_mut() {
6269            g.refresh_tables(e, cache)?;
6270            g.round_slab = false;
6271            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
6272                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
6273                // full capture past the ceiling falls through to the segment/eager arms.
6274                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
6275                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
6276                    g.round_slab = true;
6277                    return Ok(out);
6278                }
6279            }
6280            // Round-atomic ceiling check for the segment door: if any linear run in this
6281            // walk would need a NEW capture past the ceiling, the whole round runs the
6282            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
6283            // would corrupt the commit).
6284            if !g.segments_ready(self, lo, hi, t) {
6285                graphs = None;
6286            }
6287        }
6288        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
6289        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
6290        let pos_d = match stream {
6291            Some((_, ctr)) => {
6292                let mut p = e.alloc_uninit::<i32>(t)?;
6293                e.pos_iota(ctr, &mut p, t)?;
6294                p
6295            }
6296            None => {
6297                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
6298                e.htod_i32(&pos_host)?
6299            }
6300        };
6301        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
6302        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
6303        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
6304        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
6305        // rides the dc rows kernels and never reaches the fallback).
6306        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
6307        let mut il = lo;
6308        while il < hi {
6309            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
6310                let mut end = il;
6311                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
6312                    end += 1;
6313                }
6314                let g = graphs.as_deref_mut().expect("checked above");
6315                x = g.run_segment(self, e, il, end, &x, t, cache)?;
6316                g.round_slab = true;
6317                il = end;
6318                continue;
6319            }
6320            let layer = &self.layers[il];
6321            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
6322                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
6323                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
6324                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
6325                x = self.qwen35_tparallel_linear_layer(
6326                    e,
6327                    il,
6328                    &x,
6329                    t,
6330                    cache,
6331                    ckpt.as_deref_mut(),
6332                    None,
6333                    None,
6334                )?;
6335                il += 1;
6336                continue;
6337            }
6338            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
6339            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
6340            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
6341            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
6342            // run (lane/draftcost-moe).
6343            x = self.qwen35_tparallel_fa_layer(
6344                e,
6345                il,
6346                &x,
6347                t,
6348                cache,
6349                FaLayerArgs {
6350                    pos_d: &pos_d,
6351                    pos_rows: &mut pos_rows,
6352                    pos0,
6353                    seqs_append,
6354                    batch_fa_on,
6355                    graph_cap: None,
6356                    stream,
6357                    ckpt: ckpt.as_deref_mut(),
6358                },
6359            )?;
6360            il += 1;
6361        }
6362        Ok(x)
6363    }
6364
6365    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
6366    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
6367    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
6368    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
6369    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
6370    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
6371    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
6372    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
6373    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
6374    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
6375    /// original singles chain, byte-for-byte.
6376    #[allow(clippy::too_many_arguments)]
6377    fn qwen35_tparallel_dense_ffn(
6378        &self,
6379        e: &Engine,
6380        ffn_gate: &crate::model::GpuTensor,
6381        ffn_up: &crate::model::GpuTensor,
6382        ffn_down: &crate::model::GpuTensor,
6383        zn: &CudaSlice<f32>,
6384        t: usize,
6385        n_embd: usize,
6386    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6387        let n_ff = ffn_gate.out_features();
6388        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
6389        if Engine::tk_ffn_dual_on() {
6390            if let Some(((g, gs), (u, us))) =
6391                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
6392            {
6393                if e.uses_q8_1_fast(ffn_down) {
6394                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
6395                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
6396                }
6397                let mut act = e.uninit(t * n_ff)?;
6398                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
6399                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6400                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
6401            }
6402        }
6403        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
6404        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
6405        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
6406        let mut act = e.uninit(t * n_ff)?;
6407        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
6408        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
6409        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
6410    }
6411
6412    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
6413    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
6414    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
6415    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
6416    ///
6417    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
6418    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
6419    ///   generation's cache lands at new addresses that only the per-verify table refresh
6420    ///   knows — the slice-3 baked-address lesson);
6421    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
6422    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
6423    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
6424    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
6425    ///   round whose rows all sit inside the rung;
6426    /// - the host len bump moves to the replay caller (captured host code does not
6427    ///   re-run at replay).
6428    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
6429    /// host-branches on t_kv and must never be captured.
6430    #[allow(clippy::too_many_arguments)]
6431    fn qwen35_tparallel_fa_layer(
6432        &self,
6433        e: &Engine,
6434        il: usize,
6435        x: &CudaSlice<f32>,
6436        t: usize,
6437        cache: &mut Cache,
6438        args: FaLayerArgs<'_>,
6439    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6440        use cudarc::driver::DevicePtr;
6441        let cfg = &self.cfg;
6442        let n_embd = cfg.n_embd as usize;
6443        let eps = cfg.rms_eps;
6444        let head_dim_global = cfg.head_dim_k as usize;
6445        let layer = &self.layers[il];
6446        let FaLayerArgs {
6447            pos_d,
6448            pos_rows,
6449            pos0,
6450            seqs_append,
6451            batch_fa_on,
6452            graph_cap,
6453            stream,
6454            mut ckpt,
6455        } = args;
6456
6457        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6458        let anorm = layer.attn_norm.float_data();
6459        let mut xn = e.uninit(t * n_embd)?;
6460        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6461        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6462
6463        let mixed: CudaSlice<f32> = match &layer.mixer {
6464            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
6465            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
6466            // per-row serving-kernel chain cannot run (host state swaps keyed on host
6467            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
6468            // rebuild — the per-row chain only produces per-column clones). GDN rides
6469            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
6470            // and its one-scan recurrence is pinned bit-identical to T chained T=1
6471            // steps (its header + kernel-check). Position-independent, so no counter
6472            // plumbing is needed. Guards mirror the generic call site exactly.
6473            Mixer::Linear(la) if stream.is_some() => {
6474                if !(t >= 3 || (t == 2 && spec_m2()))
6475                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
6476                    || !e.uses_q8_1_fast(&la.ssm_out)
6477                {
6478                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
6479                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
6480                        .into());
6481                }
6482                let want = ckpt.is_some();
6483                let (out, stash) =
6484                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
6485                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
6486                    ck.gdn[il] = Some(st);
6487                }
6488                out
6489            }
6490            Mixer::Linear(_) => {
6491                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
6492            }
6493            Mixer::Full(fa) => {
6494                let geometry = cfg.full_attention_geometry_at(il as u32);
6495                let n_head = geometry.n_head as usize;
6496                let n_head_kv = geometry.n_head_kv as usize;
6497                let head_dim = geometry.head_dim_k as usize;
6498                let rope_dims = geometry.n_rot as usize;
6499                let rope_base = geometry.rope_base;
6500                let scale = geometry.attention_scale();
6501                // Batched projections: one weight read serves all T rows.
6502                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
6503                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
6504                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
6505                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
6506                    [&fa.wq, &fa.wk, &fa.wv],
6507                    &hq,
6508                    &hd,
6509                    t,
6510                )? {
6511                    Some(mut g3) => {
6512                        let v = g3.pop().unwrap();
6513                        let k = g3.pop().unwrap();
6514                        let qf = g3.pop().unwrap();
6515                        (qf, k, v)
6516                    }
6517                    None => (
6518                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
6519                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
6520                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
6521                    ),
6522                };
6523                let gated =
6524                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
6525                let (mut q, gate) = if gated {
6526                    let mut qs = e.uninit(t * n_head * head_dim)?;
6527                    let mut gs = e.uninit(t * n_head * head_dim)?;
6528                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
6529                    (qs, Some(gs))
6530                } else {
6531                    (qf, None)
6532                };
6533                let mut qn = e.uninit(t * n_head * head_dim)?;
6534                e.rms_norm(
6535                    &q,
6536                    fa.q_norm.float_data(),
6537                    &mut qn,
6538                    head_dim,
6539                    t * n_head,
6540                    eps,
6541                )?;
6542                q = qn;
6543                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
6544                e.rms_norm(
6545                    &k,
6546                    fa.k_norm.float_data(),
6547                    &mut kn,
6548                    head_dim,
6549                    t * n_head_kv,
6550                    eps,
6551                )?;
6552                k = kn;
6553                e.rope_neox(
6554                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
6555                )?;
6556                e.rope_neox(
6557                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
6558                )?;
6559
6560                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
6561                // draft), each through the b_n=1 serving kernels at its own t_kv.
6562                let q_dim = n_head * head_dim;
6563                let kv_dim = n_head_kv * head_dim;
6564                let mut attn = e.uninit(t * q_dim)?;
6565                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
6566                    let kvl = cache.kv[il].as_ref().unwrap();
6567                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
6568                    // the batched twins; the per-row fallback reads pair 0 (same cache
6569                    // for every row of one layer). Graph mode reads the ctx table.
6570                    let local: Option<CudaSlice<u64>> = match graph_cap {
6571                        Some(_) => None,
6572                        None => {
6573                            let s = &e.gpu.stream();
6574                            let (pk, _g) = kvl.k.device_ptr(s);
6575                            let (pv, _g2) = kvl.v.device_ptr(s);
6576                            let mut tbl = Vec::with_capacity(2 * t);
6577                            for _ in 0..t {
6578                                tbl.push(pk as u64);
6579                                tbl.push(pv as u64);
6580                            }
6581                            Some(e.htod_u64(&tbl)?)
6582                        }
6583                    };
6584                    (
6585                        kvl.kv_dim_k,
6586                        kvl.kv_dim_v,
6587                        kvl.k_tok_bytes,
6588                        kvl.v_tok_bytes,
6589                        kvl.len,
6590                        local,
6591                    )
6592                };
6593                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
6594                    Some((tb, off, _)) => (tb, off),
6595                    None => (kv_local.as_ref().expect("built above"), 0),
6596                };
6597                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
6598                // section batches into the z-batched serving twins when every row of
6599                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
6600                // guards are evaluated at the round's FIRST and LAST t_kv — the
6601                // eligibility window (vec floor .. v4 max) and each split-ladder rung
6602                // are intervals in t_kv, so ends-inside means all-inside (the straddle
6603                // law). Appending all T rows before any attend is read-equivalent to
6604                // the interleaved order: row r's walk reads keys 0..len0+r only, and
6605                // rows > r land at slots it never touches; every written cache row is
6606                // the per-token appender's exact warp program (kernel-check pinned).
6607                let t_kv_first = len0 + 1;
6608                let t_kv_last = len0 + t;
6609                let rows_batched = t >= 2
6610                    && seqs_append
6611                    && batch_fa_on
6612                    && dspark_fa_rows_on()
6613                    // the z-batched twins read stacked rows at the CACHE's kv dims;
6614                    // the projection stack is [T, n_head_kv*head_dim] — they must be
6615                    // the same stride or row z misaligns (true for this family; the
6616                    // guard keeps any asymmetric-kv model on the per-row loop).
6617                    && kdk == kv_dim
6618                    && kdv == kv_dim
6619                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
6620                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
6621                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
6622                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
6623                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
6624                // grid only — bytes proven equal above). Capture-time invariants refuse
6625                // loudly rather than bake a divergent body.
6626                let (size_kv_max, sp) = match graph_cap {
6627                    Some((_, _, rung)) => {
6628                        if !rows_batched {
6629                            return Err(format!(
6630                                "fa graph capture: layer {il} round is not batchable \
6631                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
6632                                 must never be captured"
6633                            )
6634                            .into());
6635                        }
6636                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
6637                        if t_kv_last > rung
6638                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
6639                        {
6640                            return Err(format!(
6641                                "fa graph capture: rung {rung} does not cover round \
6642                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
6643                            )
6644                            .into());
6645                        }
6646                        (rung, sp_r)
6647                    }
6648                    None => (
6649                        t_kv_last,
6650                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
6651                    ),
6652                };
6653                if let Some((_, ctr)) = stream {
6654                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
6655                    // — the generic stream arm's exact shape (rows kernels are pinned
6656                    // byte-identical to the per-row programs by kernel-check). Host len
6657                    // stays a stale lower bound; the burst drain reconciles it.
6658                    let kvl = cache.kv[il].as_mut().unwrap();
6659                    e.append_kv_quantized_rows_dc(
6660                        &k,
6661                        &v,
6662                        &mut kvl.k,
6663                        &mut kvl.v,
6664                        ctr,
6665                        t,
6666                        kdk,
6667                        kdv,
6668                        ktb,
6669                        vtb,
6670                        Engine::kv_fp8_on(),
6671                    )?;
6672                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
6673                    let k_view = e.view_u8(&kvl.k, upper * ktb);
6674                    let v_view = e.view_u8(&kvl.v, upper * vtb);
6675                    e.fa_decode_rows_dc(
6676                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
6677                        t, scale, ktb, vtb, 0, false,
6678                    )?;
6679                } else if rows_batched {
6680                    e.append_kv_quantized_seqs(
6681                        &k,
6682                        &v,
6683                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6684                        pos_d,
6685                        t,
6686                        kdk,
6687                        kdv,
6688                        ktb,
6689                        vtb,
6690                    )?;
6691                    if graph_cap.is_none() {
6692                        cache.kv[il].as_mut().unwrap().len += t;
6693                    }
6694                    e.fa_decode_batch_seqs_v4(
6695                        &q,
6696                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
6697                        pos_d,
6698                        &mut attn,
6699                        head_dim,
6700                        n_head,
6701                        n_head_kv,
6702                        t,
6703                        size_kv_max,
6704                        scale,
6705                        sp,
6706                        ktb,
6707                        vtb,
6708                    )?;
6709                } else {
6710                    if pos_rows.is_none() {
6711                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
6712                        // the dc rows kernels above and never reaches this fallback).
6713                        *pos_rows = Some(match stream {
6714                            Some((_, ctr)) => (0..t)
6715                                .map(|r| {
6716                                    let mut b = e.alloc_uninit::<i32>(1)?;
6717                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
6718                                    Ok(b)
6719                                })
6720                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
6721                            None => (0..t)
6722                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
6723                                .collect::<Result<_, _>>()?,
6724                        });
6725                    }
6726                    let pos_rows = pos_rows.as_ref().unwrap();
6727                    for r in 0..t {
6728                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
6729                        // whose row 0 is this row (arithmetic-free materialization copies,
6730                        // same as decode's per-seq fallback arm).
6731                        let mut k_row = e.uninit(kv_dim)?;
6732                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
6733                        let mut v_row = e.uninit(kv_dim)?;
6734                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
6735                        let pos_row = &pos_rows[r];
6736                        let kvl = cache.kv[il].as_mut().unwrap();
6737                        if seqs_append {
6738                            e.append_kv_quantized_seqs(
6739                                &k_row,
6740                                &v_row,
6741                                &kv_tbl.slice(kv_off..kv_off + 2),
6742                                pos_row,
6743                                1,
6744                                kdk,
6745                                kdv,
6746                                ktb,
6747                                vtb,
6748                            )?;
6749                            kvl.len += 1;
6750                        } else {
6751                            e.append_kv_quantized_view(
6752                                &k_row.slice(0..kv_dim),
6753                                &v_row.slice(0..kv_dim),
6754                                &mut kvl.k,
6755                                &mut kvl.v,
6756                                kvl.len,
6757                                kvl.kv_dim_k,
6758                                kvl.kv_dim_v,
6759                                kvl.k_tok_bytes,
6760                                kvl.v_tok_bytes,
6761                                Engine::kv_fp8_on(),
6762                            )?;
6763                            kvl.len += 1;
6764                        }
6765                        let t_kv = kvl.len;
6766                        let mut q_row = e.uninit(q_dim)?;
6767                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
6768                        let mut a_row = e.uninit(q_dim)?;
6769                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
6770                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
6771                            e.fa_decode_batch_seqs_v4(
6772                                &q_row,
6773                                &kv_tbl.slice(kv_off..kv_off + 2),
6774                                pos_row,
6775                                &mut a_row,
6776                                head_dim,
6777                                n_head,
6778                                n_head_kv,
6779                                1,
6780                                t_kv,
6781                                scale,
6782                                sp0_r,
6783                                ktb,
6784                                vtb,
6785                            )?;
6786                        } else {
6787                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
6788                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
6789                            let mut a_view = a_row.slice_mut(0..q_dim);
6790                            e.fa_decode_kvmod_view(
6791                                &q_row.slice(0..q_dim),
6792                                &k_view,
6793                                &v_view,
6794                                &mut a_view,
6795                                head_dim,
6796                                n_head,
6797                                n_head_kv,
6798                                t_kv,
6799                                scale,
6800                                kvl.k_tok_bytes,
6801                                kvl.v_tok_bytes,
6802                                Engine::kv_fp8_on(),
6803                            )?;
6804                        }
6805                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
6806                    }
6807                }
6808
6809                // Output gate (element-wise) + o-proj at m=T.
6810                let attn_g = match &gate {
6811                    Some(g) => {
6812                        let n = t * q_dim;
6813                        let mut gsig = e.uninit(n)?;
6814                        e.sigmoid(g, &mut gsig, n)?;
6815                        let mut ag = e.uninit(n)?;
6816                        e.mul(&attn, &gsig, &mut ag, n)?;
6817                        ag
6818                    }
6819                    None => attn,
6820                };
6821                e.matmul(&fa.wo, &attn_g, t)?
6822            }
6823        };
6824
6825        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
6826        let pnorm = layer.post_attn_norm.float_data();
6827        let mut x1 = e.uninit(t * n_embd)?;
6828        let mut zn = e.uninit(t * n_embd)?;
6829        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
6830        let ffn_out = match &layer.ffn {
6831            crate::hybrid::Ffn::Dense {
6832                ffn_gate,
6833                ffn_up,
6834                ffn_down,
6835            } => {
6836                assert!(
6837                    self.cfg.m3.is_none(),
6838                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
6839                );
6840                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
6841            }
6842            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
6843        };
6844        let mut x2 = e.uninit(t * n_embd)?;
6845        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
6846        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
6847        self.dflash_tap(e, cache, il, &x2, t)?;
6848        Ok(x2)
6849    }
6850
6851    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
6852    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
6853    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
6854    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
6855    /// bit-identical by construction:
6856    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
6857    ///   the device sequence is driven entirely by the 6-entry pointer table, which
6858    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
6859    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
6860    ///   legacy post-swap clone read.
6861    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
6862    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
6863    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
6864    /// None builds the per-verify table exactly as before.
6865    #[allow(clippy::too_many_arguments)]
6866    fn qwen35_tparallel_linear_layer(
6867        &self,
6868        e: &Engine,
6869        il: usize,
6870        x: &CudaSlice<f32>,
6871        t: usize,
6872        cache: &mut Cache,
6873        mut ckpt: Option<&mut VerifyCkpt>,
6874        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
6875        table_src: Option<(&CudaSlice<u64>, usize)>,
6876    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6877        use cudarc::driver::DevicePtr;
6878        let cfg = &self.cfg;
6879        let n_embd = cfg.n_embd as usize;
6880        let eps = cfg.rms_eps;
6881        let layer = &self.layers[il];
6882        let Mixer::Linear(la) = &layer.mixer else {
6883            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
6884        };
6885        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
6886        let anorm = layer.attn_norm.float_data();
6887        let mut xn = e.uninit(t * n_embd)?;
6888        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
6889        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
6890
6891        let geometry = la.geometry;
6892        let d_state = geometry.key_head_dim as usize;
6893        let num_k = geometry.key_heads as usize;
6894        let num_v = geometry.value_heads as usize;
6895        let d_conv = geometry.conv_kernel as usize;
6896        let key_dim = d_state * num_k;
6897        let value_dim = geometry.value_head_dim as usize * num_v;
6898        let conv_dim = key_dim * 2 + value_dim;
6899        let gdn_scale = 1.0 / (d_state as f32).sqrt();
6900
6901        // ---- batched projections: one weight read for all T rows ----
6902        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
6903        // per (tensor, token, row) to the four singles; refused (layout/tier) or
6904        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
6905        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
6906            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
6907            &hq,
6908            &hd,
6909            t,
6910        )? {
6911            Some(mut g4) => {
6912                let alpha = g4.pop().unwrap();
6913                let beta_raw = g4.pop().unwrap();
6914                let z = g4.pop().unwrap();
6915                let qkv_mixed = g4.pop().unwrap();
6916                (qkv_mixed, z, beta_raw, alpha)
6917            }
6918            None => (
6919                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
6920                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
6921                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
6922                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
6923            ),
6924        };
6925        let beta_w = la.ssm_beta.out_features();
6926        let alpha_w = la.ssm_alpha.out_features();
6927        let qkv_w = la.wqkv.out_features();
6928
6929        // ---- per-row state chain through the b_n=1 serving kernels ----
6930        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
6931        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
6932        let table_local: Option<CudaSlice<u64>> = match table_src {
6933            Some(_) => None,
6934            None => {
6935                let rl = cache.recur[il].as_ref().unwrap();
6936                let s = &e.gpu.stream();
6937                let (pc, _g0) = rl.conv_state.device_ptr(s);
6938                let (p0, _g1) = rl.ssm_state.device_ptr(s);
6939                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
6940                Some(e.htod_u64(&[
6941                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
6942                ])?)
6943            }
6944        };
6945        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
6946            Some((tb, off)) => (tb, off),
6947            None => (table_local.as_ref().unwrap(), 0),
6948        };
6949        let mut o_all = e.uninit(t * value_dim)?;
6950        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
6951            if ckpt.is_some() && stash.is_none() && t >= 2 {
6952                Some(Vec::with_capacity(t - 1))
6953            } else {
6954                None
6955            };
6956        let mut stash = stash;
6957        // Per-row scratch reused across rows (uninit is cheap but not free at
6958        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
6959        // [T, ...] buffers — zero arithmetic-free copies in this loop.
6960        let mut conv_out = e.uninit(conv_dim)?;
6961        let mut q_l2 = e.uninit(value_dim)?;
6962        let mut k_l2 = e.uninit(value_dim)?;
6963        let mut v_gd = e.uninit(value_dim)?;
6964        let mut beta_b = e.uninit(num_v)?;
6965        let mut g_log = e.uninit(num_v)?;
6966        for r in 0..t {
6967            let base = toff + if r % 2 == 0 { 0 } else { 3 };
6968            let conv_view = table.slice(base..base + 1);
6969            let in_view = table.slice(base + 1..base + 2);
6970            let out_view = table.slice(base + 2..base + 3);
6971            e.ssm_conv1d_fused_decode_b_view(
6972                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
6973                &conv_view,
6974                la.ssm_conv1d.float_data(),
6975                &mut conv_out,
6976                conv_dim,
6977                d_conv,
6978                1,
6979            )?;
6980            e.gdn_prep_decode_b_view(
6981                &conv_out,
6982                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
6983                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
6984                la.ssm_dt.float_data(),
6985                la.ssm_a.float_data(),
6986                &mut q_l2,
6987                &mut k_l2,
6988                &mut v_gd,
6989                &mut beta_b,
6990                &mut g_log,
6991                d_state,
6992                num_v,
6993                num_k,
6994                key_dim,
6995                eps,
6996                conv_dim,
6997                1,
6998            )?;
6999            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7000            e.gdn_scan_s128_batched_view(
7001                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7002                gdn_scale,
7003            )?;
7004            if r + 1 < t {
7005                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7006                // odd rows write s0 — the same physical state the legacy post-swap
7007                // canonical clone read.
7008                let rl = cache.recur[il]
7009                    .as_ref()
7010                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
7011                let ssm_src = if r % 2 == 0 {
7012                    &rl.ssm_state_alt
7013                } else {
7014                    &rl.ssm_state
7015                };
7016                match stash.as_mut() {
7017                    Some((conv_slab, ssm_slab)) => {
7018                        // BOTH stash reads go through the pointer table at run time: the
7019                        // ssm handles ping-pong between rounds, and the ctx (with its
7020                        // captured graphs) outlives the Cache — a fresh generation's
7021                        // conv/ssm buffers land at new addresses that only the per-round
7022                        // table refresh knows. A baked direct copy would read freed
7023                        // memory (parity was the slice-3 smoke divergence; cache
7024                        // lifetime is the cross-generation twin).
7025                        e.copy_indirect_src_f32(
7026                            &conv_view,
7027                            conv_slab,
7028                            r * conv_dim * (d_conv - 1),
7029                            conv_dim * (d_conv - 1),
7030                        )?;
7031                        // The ssm handles PING-PONG between rounds: a captured direct
7032                        // copy would bake the capture-time physical buffer and read the
7033                        // wrong parity after any odd-vt round (the slice-3 smoke
7034                        // divergence). Read the src address from row r's OUT table
7035                        // entry at run time — the same entry the scan just wrote.
7036                        e.copy_indirect_src_f32(
7037                            &out_view,
7038                            ssm_slab,
7039                            r * d_state * d_state * num_v,
7040                            d_state * d_state * num_v,
7041                        )?;
7042                    }
7043                    None => {
7044                        if let Some(states) = col_states.as_mut() {
7045                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7046                        }
7047                    }
7048                }
7049            }
7050        }
7051        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7052        // handle motion is identical and the device sequence never read the handles.
7053        if t % 2 == 1 {
7054            let rl = cache.recur[il].as_mut().unwrap();
7055            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7056        }
7057        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7058            checkpoint.cols[il] = Some(states);
7059        }
7060
7061        // ---- batched gated norm + out-projection at m=T ----
7062        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7063            let (gq, gd) = e.gated_rmsnorm_q8_1(
7064                &o_all,
7065                la.ssm_norm.float_data(),
7066                &z,
7067                d_state,
7068                t * num_v,
7069                eps,
7070            )?;
7071            let g0 = e.zeros(0)?;
7072            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7073        } else {
7074            let mut gn = e.uninit(t * value_dim)?;
7075            e.gated_rmsnorm(
7076                &o_all,
7077                la.ssm_norm.float_data(),
7078                &z,
7079                &mut gn,
7080                d_state,
7081                t * num_v,
7082                eps,
7083            )?;
7084            e.matmul(&la.ssm_out, &gn, t)?
7085        };
7086
7087        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7088        let pnorm = layer.post_attn_norm.float_data();
7089        let mut x1 = e.uninit(t * n_embd)?;
7090        let mut zn = e.uninit(t * n_embd)?;
7091        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7092        let ffn_out = match &layer.ffn {
7093            crate::hybrid::Ffn::Dense {
7094                ffn_gate,
7095                ffn_up,
7096                ffn_down,
7097            } => {
7098                assert!(
7099                    self.cfg.m3.is_none(),
7100                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7101                );
7102                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7103            }
7104            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7105        };
7106        let mut x2 = e.uninit(t * n_embd)?;
7107        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7108        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7109        self.dflash_tap(e, cache, il, &x2, t)?;
7110        Ok(x2)
7111    }
7112
7113    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
7114    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
7115    /// carried in from outside the range) and exits with the range's final residual materialized
7116    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
7117    /// instead of one.
7118    ///
7119    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
7120    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
7121    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
7122    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
7123    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
7124    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
7125    /// code — there is no "split version" of the verify math.
7126    ///
7127    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
7128    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
7129    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
7130    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
7131    #[allow(clippy::too_many_arguments)]
7132    fn verify_layers(
7133        &self,
7134        e: &Engine,
7135        mut x: CudaSlice<f32>,
7136        lo: usize,
7137        hi: usize,
7138        pos_d: &CudaSlice<i32>,
7139        pos0: usize,
7140        t: usize,
7141        cache: &mut Cache,
7142        mut ckpt: Option<&mut VerifyCkpt>,
7143        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7144        graphs: Option<&mut DsparkVerifyGraphs>,
7145    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7146        if self.sliding_gated_moe_batch_program() {
7147            if stream.is_some() {
7148                return Err(
7149                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
7150                            cannot express the SWA offset KV view)"
7151                        .into(),
7152                );
7153            }
7154            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
7155        }
7156        if self.batched_serving_numeric_class() {
7157            return self.qwen35_verify_batch_layers(
7158                e,
7159                x,
7160                lo,
7161                hi,
7162                pos0,
7163                t,
7164                cache,
7165                ckpt.take(),
7166                stream,
7167                graphs,
7168            );
7169        }
7170        let n_embd = self.cfg.n_embd as usize;
7171        let eps = self.cfg.rms_eps;
7172        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
7173        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
7174        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
7175        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
7176        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
7177        // residual the next layer needs) as its `res` output. Falls back to the separate add
7178        // when the next layer is off the fused-q8 path.
7179        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
7180        for il in lo..hi {
7181            let layer = &self.layers[il];
7182            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
7183            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
7184            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
7185            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
7186            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
7187            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
7188            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
7189            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
7190            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
7191            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
7192            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
7193            // projections only; Linear mixer: the batched arm — the per-column fallback needs
7194            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
7195            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
7196            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
7197            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
7198            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
7199            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
7200            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
7201            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
7202            let lin_q8_only = match &layer.mixer {
7203                Mixer::Linear(la) => {
7204                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
7205                }
7206                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
7207                _ => true,
7208            };
7209            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
7210            // a non-fused layer still performs the residual add.
7211            let taken = pending.take();
7212            let (h, h_q8) = if norm_fused && lin_q8_only {
7213                let pair = match taken {
7214                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
7215                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
7216                    Some((x1p, f1p)) => {
7217                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
7218                        let p = e.add_rms_norm_q8_1(
7219                            &x1p,
7220                            &f1p,
7221                            layer.attn_norm.float_data(),
7222                            &mut x2,
7223                            n_embd,
7224                            t,
7225                            eps,
7226                        )?;
7227                        x = x2;
7228                        p
7229                    }
7230                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
7231                };
7232                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
7233            } else {
7234                if let Some((x1p, f1p)) = taken {
7235                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7236                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7237                    x = x2;
7238                }
7239                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
7240                if norm_fused {
7241                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7242                } else {
7243                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
7244                }
7245                (h, None)
7246            };
7247            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
7248
7249            let mixed = match &layer.mixer {
7250                Mixer::Full(fa) => self.full_attn_verify(
7251                    e,
7252                    fa,
7253                    &h,
7254                    h_q8_ref,
7255                    pos_d,
7256                    t,
7257                    cache,
7258                    il,
7259                    stream.map(|(_, c)| c),
7260                )?,
7261                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7262                Mixer::Linear(la) => {
7263                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
7264                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
7265                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
7266                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
7267                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
7268                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
7269                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
7270                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
7271                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
7272                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
7273                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
7274                    if (t >= 3 || (t == 2 && spec_m2()))
7275                        && mixer_fast
7276                        && e.uses_q8_1_fast(&la.ssm_out)
7277                    {
7278                        let want = ckpt.is_some();
7279                        let (out, stash) =
7280                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
7281                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7282                            ck.gdn[il] = Some(st);
7283                        }
7284                        out
7285                    } else {
7286                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
7287                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7288                            if ckpt.is_some() && t >= 2 {
7289                                Some(Vec::with_capacity(t - 1))
7290                            } else {
7291                                None
7292                            };
7293                        for col in 0..t {
7294                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
7295                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
7296                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
7297                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
7298                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
7299                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
7300                            // (pure dtod — cannot change any computed value). Last column skipped:
7301                            // rebuild targets are j <= t-1 columns.
7302                            if let Some(cs) = col_states.as_mut() {
7303                                if col + 1 < t {
7304                                    let rl = cache.recur[il].as_ref().unwrap();
7305                                    cs.push((
7306                                        e.clone_dtod(&rl.conv_state)?,
7307                                        e.clone_dtod(&rl.ssm_state)?,
7308                                    ));
7309                                }
7310                            }
7311                        }
7312                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
7313                            // ReplaySSM-assessment instrumentation (2026-07-30): the
7314                            // per-column clones are the only true state snapshots left in
7315                            // the verify (the batched path stashes INPUTS and replays).
7316                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
7317                                static ONCE: std::sync::Once = std::sync::Once::new();
7318                                let bytes: usize =
7319                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
7320                                ONCE.call_once(|| eprintln!(
7321                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
7322                                    cs.len(), bytes as f64 / 1e6));
7323                            }
7324                            ck.cols[il] = Some(cs);
7325                        }
7326                        out
7327                    }
7328                }
7329            };
7330
7331            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
7332            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
7333            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
7334            let ffn_fuse = match &layer.ffn {
7335                crate::hybrid::Ffn::Dense {
7336                    ffn_gate, ffn_up, ..
7337                } => {
7338                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
7339                        && e.uses_q8_1_fast(ffn_gate)
7340                        && e.uses_q8_1_fast(ffn_up)
7341                }
7342                crate::hybrid::Ffn::Moe(_) => false,
7343            };
7344            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
7345            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
7346            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
7347            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
7348            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
7349            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
7350            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
7351            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
7352            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
7353            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
7354            // mirror decode's dispatch or spec self-consistency fails.
7355            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
7356            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
7357            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
7358            let mut z = e.zeros(0)?; // replaced below on the unfused arms
7359            let z_q8 = if fuse_q8 {
7360                Some(e.add_rms_norm_q8_1(
7361                    &x,
7362                    &mixed,
7363                    layer.post_attn_norm.float_data(),
7364                    &mut x1,
7365                    n_embd,
7366                    t,
7367                    eps,
7368                )?)
7369            } else {
7370                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
7371                if ffn_fuse {
7372                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
7373                    e.rms_norm_decode(
7374                        &x1,
7375                        layer.post_attn_norm.float_data(),
7376                        &mut zf,
7377                        n_embd,
7378                        t,
7379                        eps,
7380                    )?;
7381                } else {
7382                    e.add_rms_norm(
7383                        &x,
7384                        &mixed,
7385                        layer.post_attn_norm.float_data(),
7386                        &mut x1,
7387                        &mut zf,
7388                        n_embd,
7389                        t,
7390                        eps,
7391                    )?;
7392                }
7393                z = zf;
7394                None
7395            };
7396            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
7397            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
7398            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
7399            let ffn_out = match &layer.ffn {
7400                crate::hybrid::Ffn::Dense {
7401                    ffn_gate,
7402                    ffn_up,
7403                    ffn_down,
7404                } => {
7405                    let n_ff = ffn_gate.out_features();
7406                    if let Some((zq, zd)) = z_q8.as_ref() {
7407                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
7408                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
7409                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
7410                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
7411                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
7412                        // structure at nrows=t.
7413                        let pair =
7414                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
7415                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
7416                                None => None,
7417                            };
7418                        let (gate, gs, up, us) = match pair {
7419                            Some(x4) => x4,
7420                            None => (
7421                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
7422                                1.0, // scale already applied inside _pre
7423                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
7424                                1.0,
7425                            ),
7426                        };
7427                        if e.uses_q8_1_fast(ffn_down) {
7428                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
7429                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
7430                        } else {
7431                            let mut act = vbuf(e, t * n_ff)?;
7432                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
7433                            e.matmul_decode_exact(ffn_down, &act, t)?
7434                        }
7435                    } else {
7436                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
7437                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
7438                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
7439                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
7440                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
7441                        let (gate, up) =
7442                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
7443                                Some(pair) => pair,
7444                                None => (
7445                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
7446                                    e.matmul_decode_exact(ffn_up, &z, t)?,
7447                                ),
7448                            };
7449                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
7450                        Self::ffn_act_lim(
7451                            e,
7452                            &self.cfg,
7453                            &gate,
7454                            &up,
7455                            1.0,
7456                            1.0,
7457                            dense_lim,
7458                            &mut act,
7459                            t * n_ff,
7460                        )?;
7461                        e.matmul_decode_exact(ffn_down, &act, t)?
7462                    }
7463                }
7464                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
7465            };
7466            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
7467            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
7468            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
7469            pending = Some((x1, ffn_out));
7470        }
7471        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
7472        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
7473        if let Some((x1p, f1p)) = pending.take() {
7474            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
7475            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
7476            x = x2;
7477        }
7478        Ok(x)
7479    }
7480    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
7481    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
7482    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
7483    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
7484    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
7485    /// ssm state exactly like T sequential decode steps.
7486    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
7487    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
7488    #[allow(clippy::too_many_arguments)]
7489    fn linear_attn_verify_t(
7490        &self,
7491        e: &Engine,
7492        la: &LinearAttnLayer,
7493        h: &CudaSlice<f32>,
7494        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
7495        t: usize,
7496        cache: &mut Cache,
7497        il: usize,
7498        want_stash: bool,
7499    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
7500        let cfg = &self.cfg;
7501        let geometry = la.geometry;
7502        let d_state = geometry.key_head_dim as usize;
7503        let num_k = geometry.key_heads as usize;
7504        let num_v = geometry.value_heads as usize;
7505        let d_conv = geometry.conv_kernel as usize;
7506        let key_dim = d_state * num_k;
7507        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
7508        let eps = cfg.rms_eps;
7509        let scale = 1.0 / (d_state as f32).sqrt();
7510
7511        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
7512        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
7513        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
7514        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
7515        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
7516        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
7517        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
7518        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
7519        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
7520        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
7521        // Bit-identical per (tensor,token,row) — see spec_fused_t().
7522        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
7523        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
7524        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
7525        // and feeds every projection; the caller guaranteed all four input projections are
7526        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
7527        let h_q8_t = if h_q8.is_none()
7528            && spec_fused_t()
7529            && (2..=4).contains(&t)
7530            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
7531                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
7532        {
7533            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
7534        } else {
7535            None
7536        };
7537        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
7538        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
7539            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
7540        let (qkv_mixed, z) = {
7541            let mut fused = None;
7542            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
7543                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7544                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
7545            } else if let Some((hq, hd)) = hq8_any {
7546                if spec_fused_t() && (2..=4).contains(&t) {
7547                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
7548                }
7549            }
7550            match (fused, hq8_any) {
7551                (Some(pair), _) => pair,
7552                (None, Some((hq, hd))) if h_q8.is_some() => (
7553                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
7554                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
7555                ),
7556                (None, _) => (
7557                    e.matmul_decode_exact(&la.wqkv, h, t)?,
7558                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
7559                ),
7560            }
7561        };
7562        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
7563        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
7564        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
7565        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
7566        let (beta_raw, alpha) = if t == 1 {
7567            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
7568            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
7569                Some(((mut b, bs), (mut a, as_))) => {
7570                    if bs != 1.0 {
7571                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
7572                    }
7573                    if as_ != 1.0 {
7574                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
7575                    }
7576                    (b, a)
7577                }
7578                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
7579                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
7580                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
7581                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
7582                    Some((b, a)) => (b, a),
7583                    None => (
7584                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
7585                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
7586                    ),
7587                },
7588            }
7589        } else {
7590            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
7591            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
7592            let mut nvfp4_fused = None;
7593            let mut q8_fused = None;
7594            if let Some((hq, hd)) = hq8_any {
7595                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
7596                    nvfp4_fused =
7597                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7598                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
7599                        static ONCE: std::sync::Once = std::sync::Once::new();
7600                        ONCE.call_once(|| {
7601                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
7602                        });
7603                    }
7604                }
7605                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
7606                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
7607                }
7608            }
7609            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
7610                if bs != 1.0 {
7611                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
7612                }
7613                if as_ != 1.0 {
7614                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
7615                }
7616                (b, a)
7617            } else if let Some(pair) = q8_fused {
7618                pair
7619            } else {
7620                match hq8_any {
7621                    Some((hq, hd)) if h_q8.is_some() => (
7622                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
7623                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
7624                    ),
7625                    _ => (
7626                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
7627                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
7628                    ),
7629                }
7630            }
7631        };
7632
7633        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
7634        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
7635        let rl = cache.recur[il].as_mut().unwrap();
7636        let mut conv_out = e.uninit(conv_dim * t)?;
7637        e.ssm_conv1d_tm_state(
7638            &qkv_mixed,
7639            &mut rl.conv_state,
7640            la.ssm_conv1d.float_data(),
7641            &mut conv_out,
7642            conv_dim,
7643            t,
7644            d_conv,
7645        )?;
7646
7647        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
7648        let mut q_g = e.uninit(d_state * num_v * t)?;
7649        let mut k_g = e.uninit(d_state * num_v * t)?;
7650        let mut v_g = e.uninit(d_state * num_v * t)?;
7651        e.qkv_to_gdn_repack(
7652            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
7653        )?;
7654        let mut q_l2 = e.uninit(d_state * num_v * t)?;
7655        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
7656        let mut k_l2 = e.uninit(d_state * num_v * t)?;
7657        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
7658        let mut beta = e.uninit(t * num_v)?;
7659        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
7660        let mut g_log = e.uninit(t * num_v)?;
7661        e.gdn_glog(
7662            &alpha,
7663            la.ssm_dt.float_data(),
7664            la.ssm_a.float_data(),
7665            &mut g_log,
7666            num_v,
7667            t,
7668        )?;
7669
7670        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
7671        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
7672        let mut o = e.uninit(d_state * num_v * t)?;
7673        {
7674            let crate::cache::RecurLayer {
7675                ssm_state,
7676                ssm_state_alt,
7677                ..
7678            } = rl;
7679            e.gdn_scan_s128(
7680                &q_l2,
7681                &k_l2,
7682                &v_g,
7683                &g_log,
7684                &beta,
7685                ssm_state,
7686                ssm_state_alt,
7687                &mut o,
7688                num_v,
7689                t,
7690                scale,
7691            )?;
7692        }
7693        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7694
7695        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
7696        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
7697        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
7698        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
7699        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
7700        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
7701        let out = if e.uses_q8_1_fast(&la.ssm_out) {
7702            let (gq, gd) =
7703                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
7704            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
7705        } else {
7706            let mut gn = e.uninit(d_state * num_v * t)?;
7707            e.gated_rmsnorm(
7708                &o,
7709                la.ssm_norm.float_data(),
7710                &z,
7711                &mut gn,
7712                d_state,
7713                num_v * t,
7714                eps,
7715            )?;
7716            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
7717            // would fall to dp4a with a different FP reduction order — same class of bug as
7718            // the input projs).
7719            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
7720        };
7721        let stash = if want_stash {
7722            Some(GdnStash {
7723                qkv_mixed,
7724                q_l2,
7725                k_l2,
7726                v_g,
7727                g_log,
7728                beta,
7729            })
7730        } else {
7731            None
7732        };
7733        Ok((out, stash))
7734    }
7735
7736    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
7737    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
7738    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
7739    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
7740    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
7741    ///   replaying them.
7742    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
7743    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
7744    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
7745    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
7746    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
7747    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
7748    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
7749    fn commit_verified_prefix(
7750        &self,
7751        e: &Engine,
7752        cache: &mut Cache,
7753        snap: &crate::cache::CacheSnapshot,
7754        ckpt: &VerifyCkpt,
7755        j: usize,
7756        kv_lens_done: bool,
7757        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
7758    ) -> Result<(), Box<dyn std::error::Error>> {
7759        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
7760        // recurrent state and must never be forced through a synthetic SSM geometry.
7761        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
7762        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
7763        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
7764        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
7765        // buffers and stream order are identical to the per-layer memcpy sequence; the
7766        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
7767        let mut batched_cols = false;
7768        if state_copy_batch_on() && dev_j.is_none() {
7769            use cudarc::driver::DevicePtr;
7770            let s = &e.gpu.stream();
7771            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
7772            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
7773            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
7774            let mut uniform = true;
7775            for il in 0..self.layers.len() {
7776                let Some(rl) = cache.recur[il].as_ref() else {
7777                    continue;
7778                };
7779                if ckpt.gdn[il].is_some() {
7780                    continue; // kernel-rebuild arm restores below, per layer
7781                }
7782                let Some(cols) = &ckpt.cols[il] else {
7783                    continue; // missing-ckpt error surfaces in the main loop
7784                };
7785                let (c, st) = &cols[j - 1];
7786                if conv_pairs.is_empty() {
7787                    conv_words = c.len();
7788                    ssm_words = st.len();
7789                } else if c.len() != conv_words || st.len() != ssm_words {
7790                    uniform = false;
7791                    break;
7792                }
7793                let (pc, _g0) = c.device_ptr(s);
7794                let (dc, _g1) = rl.conv_state.device_ptr(s);
7795                let (ps, _g2) = st.device_ptr(s);
7796                let (ds, _g3) = rl.ssm_state.device_ptr(s);
7797                conv_pairs.push((pc as u64, dc as u64));
7798                ssm_pairs.push((ps as u64, ds as u64));
7799            }
7800            if uniform && !conv_pairs.is_empty() {
7801                let n = conv_pairs.len();
7802                let mut t = vec![0u64; 2 * n];
7803                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
7804                    t[k] = src;
7805                    t[n + k] = dst;
7806                }
7807                let conv_t = e.htod_u64(&t)?;
7808                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
7809                    t[k] = src;
7810                    t[n + k] = dst;
7811                }
7812                let ssm_t = e.htod_u64(&t)?;
7813                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
7814                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
7815                batched_cols = true;
7816            }
7817        }
7818        rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
7819        for il in 0..self.layers.len() {
7820            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
7821                kvl.len = saved + j;
7822                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
7823                if !kv_lens_done {
7824                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
7825                }
7826            }
7827            if let Some(rl) = cache.recur[il].as_mut() {
7828                let Mixer::Linear(linear) = &self.layers[il].mixer else {
7829                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7830                };
7831                let geometry = linear.geometry;
7832                let d_state = geometry.key_head_dim as usize;
7833                let num_k = geometry.key_heads as usize;
7834                let num_v = geometry.value_heads as usize;
7835                let d_conv = geometry.conv_kernel as usize;
7836                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7837                let scale = 1.0 / (d_state as f32).sqrt();
7838                if let Some(st) = &ckpt.gdn[il] {
7839                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7840                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7841                    if let Some((acc, base, t_v)) = dev_j {
7842                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
7843                        e.ssm_conv_ring_rebuild_dc(
7844                            &st.qkv_mixed,
7845                            ring_old,
7846                            &mut rl.conv_state,
7847                            conv_dim,
7848                            acc,
7849                            base,
7850                            t_v,
7851                            d_conv,
7852                        )?;
7853                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
7854                        e.gdn_scan_s128_dc(
7855                            &st.q_l2,
7856                            &st.k_l2,
7857                            &st.v_g,
7858                            &st.g_log,
7859                            &st.beta,
7860                            state_in,
7861                            &mut rl.ssm_state,
7862                            &mut o,
7863                            num_v,
7864                            acc,
7865                            base,
7866                            t_v,
7867                            scale,
7868                        )?;
7869                    } else {
7870                        e.ssm_conv_ring_rebuild(
7871                            &st.qkv_mixed,
7872                            ring_old,
7873                            &mut rl.conv_state,
7874                            conv_dim,
7875                            j,
7876                            d_conv,
7877                        )?;
7878                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
7879                        e.gdn_scan_s128(
7880                            &st.q_l2,
7881                            &st.k_l2,
7882                            &st.v_g,
7883                            &st.g_log,
7884                            &st.beta,
7885                            state_in,
7886                            &mut rl.ssm_state,
7887                            &mut o,
7888                            num_v,
7889                            j,
7890                            scale,
7891                        )?;
7892                    }
7893                } else if let Some(cols) = &ckpt.cols[il] {
7894                    if !batched_cols {
7895                        let (c, s) = &cols[j - 1];
7896                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
7897                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
7898                    }
7899                } else {
7900                    return Err(
7901                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
7902                    );
7903                }
7904            }
7905        }
7906        cache.pos = snap.pos + j;
7907        Ok(())
7908    }
7909
7910    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
7911    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
7912    fn commit_verified_prefix_stream(
7913        &self,
7914        e: &Engine,
7915        cache: &mut Cache,
7916        snap: &crate::cache::CacheSnapshot,
7917        ckpt: &VerifyCkpt,
7918        acc: &CudaSlice<u32>,
7919        base: usize,
7920        t_v: usize,
7921    ) -> Result<(), Box<dyn std::error::Error>> {
7922        for il in 0..self.layers.len() {
7923            if let Some(rl) = cache.recur[il].as_mut() {
7924                let Mixer::Linear(linear) = &self.layers[il].mixer else {
7925                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
7926                };
7927                let geometry = linear.geometry;
7928                let d_state = geometry.key_head_dim as usize;
7929                let num_k = geometry.key_heads as usize;
7930                let num_v = geometry.value_heads as usize;
7931                let d_conv = geometry.conv_kernel as usize;
7932                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
7933                let scale = 1.0 / (d_state as f32).sqrt();
7934                let st = ckpt.gdn[il]
7935                    .as_ref()
7936                    .ok_or("stream restore: batched-linear stash missing")?;
7937                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
7938                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
7939                e.ssm_conv_ring_rebuild_dc(
7940                    &st.qkv_mixed,
7941                    ring_old,
7942                    &mut rl.conv_state,
7943                    conv_dim,
7944                    acc,
7945                    base,
7946                    t_v,
7947                    d_conv,
7948                )?;
7949                let mut o = e.uninit(d_state * num_v * t_v)?;
7950                e.gdn_scan_s128_dc(
7951                    &st.q_l2,
7952                    &st.k_l2,
7953                    &st.v_g,
7954                    &st.g_log,
7955                    &st.beta,
7956                    state_in,
7957                    &mut rl.ssm_state,
7958                    &mut o,
7959                    num_v,
7960                    acc,
7961                    base,
7962                    t_v,
7963                    scale,
7964                )?;
7965            }
7966        }
7967        Ok(())
7968    }
7969
7970    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
7971    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
7972    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
7973    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
7974    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
7975    pub fn decode_step_t_aux2(
7976        &self,
7977        e: &Engine,
7978        tokens: &[u32],
7979        pos0: usize,
7980        cache: &mut Cache,
7981        aux_layers: &[usize],
7982        pred_col: Option<usize>,
7983    ) -> Result<
7984        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
7985        Box<dyn std::error::Error>,
7986    > {
7987        let cfg = &self.cfg;
7988        let n_embd = cfg.n_embd as usize;
7989        let eps = cfg.rms_eps;
7990        let t = tokens.len();
7991        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
7992        let pos_d = e.htod_i32(&pos_vec)?;
7993        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
7994        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
7995        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
7996        let want_pred = pred_col.is_some();
7997
7998        for (il, layer) in self.layers.iter().enumerate() {
7999            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8000            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8001            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8002            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8003            if norm_fused {
8004                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8005            } else {
8006                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8007            }
8008            let mixed = match &layer.mixer {
8009                Mixer::Full(fa) => {
8010                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8011                }
8012                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8013                Mixer::Linear(la) => {
8014                    let mut out = e.zeros(t * n_embd)?;
8015                    for col in 0..t {
8016                        let mut h_col = e.zeros(n_embd)?;
8017                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
8018                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8019                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8020                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8021                    }
8022                    out
8023                }
8024            };
8025            let ffn_fuse = match &layer.ffn {
8026                crate::hybrid::Ffn::Dense {
8027                    ffn_gate, ffn_up, ..
8028                } => {
8029                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8030                        && e.uses_q8_1_fast(ffn_gate)
8031                        && e.uses_q8_1_fast(ffn_up)
8032                }
8033                crate::hybrid::Ffn::Moe(_) => false,
8034            };
8035            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8036            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8037            if ffn_fuse {
8038                e.add(&x, &mixed, &mut x1, t * n_embd)?;
8039                e.rms_norm_decode(
8040                    &x1,
8041                    layer.post_attn_norm.float_data(),
8042                    &mut z,
8043                    n_embd,
8044                    t,
8045                    eps,
8046                )?;
8047            } else {
8048                e.add_rms_norm(
8049                    &x,
8050                    &mixed,
8051                    layer.post_attn_norm.float_data(),
8052                    &mut x1,
8053                    &mut z,
8054                    n_embd,
8055                    t,
8056                    eps,
8057                )?;
8058            }
8059            let ffn_out = match &layer.ffn {
8060                crate::hybrid::Ffn::Dense {
8061                    ffn_gate,
8062                    ffn_up,
8063                    ffn_down,
8064                } => {
8065                    let n_ff = ffn_gate.out_features();
8066                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8067                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8068                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8069                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8070                    Self::ffn_act_lim(
8071                        e,
8072                        &self.cfg,
8073                        &gate,
8074                        &up,
8075                        1.0,
8076                        1.0,
8077                        self.cfg.clamp_shexp_at(il as u32),
8078                        &mut act,
8079                        t * n_ff,
8080                    )?;
8081                    e.matmul_decode_exact(ffn_down, &act, t)?
8082                }
8083                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8084            };
8085            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8086            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8087            if aux_layers.contains(&il) {
8088                let mut a = e.zeros(n_embd)?;
8089                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
8090                aux_last.push(a);
8091                if let Some(pc) = pred_col {
8092                    let mut ap = e.zeros(n_embd)?;
8093                    e.copy_view_into(
8094                        &mut ap,
8095                        0,
8096                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
8097                        n_embd,
8098                    )?;
8099                    aux_pred.push(ap);
8100                }
8101            }
8102            x = x2;
8103        }
8104        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
8105        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
8106        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
8107        let host = e.dtoh(&logits)?;
8108        cache.pos += t;
8109        Ok((
8110            host,
8111            aux_last,
8112            if want_pred { Some(aux_pred) } else { None },
8113        ))
8114    }
8115
8116    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
8117    /// `step35_decode_attn`.
8118    ///
8119    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
8120    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
8121    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
8122    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
8123    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
8124    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
8125    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
8126    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
8127    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
8128    /// position of each query row. A batched twin would have to reproduce all of that AND the
8129    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
8130    /// take one `base_len`, not a per-row offset).
8131    ///
8132    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
8133    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
8134    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
8135    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
8136    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
8137    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
8138    /// step35 twin is a perf lane's job and must be gated against this arm.
8139    ///
8140    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
8141    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
8142    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
8143    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
8144    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
8145    #[allow(clippy::too_many_arguments)]
8146    fn step35_verify(
8147        &self,
8148        e: &Engine,
8149        fa: &FullAttnLayer,
8150        h: &CudaSlice<f32>,
8151        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8152        t: usize,
8153        cache: &mut Cache,
8154        il: usize,
8155    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8156        let n_embd = self.cfg.n_embd as usize;
8157        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
8158        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
8159        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
8160        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
8161        // cannot regress it into silently reading an empty buffer.
8162        assert_eq!(
8163            h.len(),
8164            t * n_embd,
8165            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
8166             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
8167            h_q8.is_some()
8168        );
8169        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
8170        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
8171        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
8172        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
8173        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
8174        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
8175        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
8176        for r in 0..t {
8177            // Absolute position of this query row. `cache.pos` is the committed length at round
8178            // start and every row before r has already been appended by this loop, so the r-th
8179            // verify token sits at cache.pos + r — the same position eager decode would give it.
8180            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
8181            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
8182            e.copy_view_into(
8183                &mut h_row,
8184                0,
8185                &h.slice(r * n_embd..(r + 1) * n_embd),
8186                n_embd,
8187            )?;
8188            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
8189            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
8190            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
8191            debug_assert_eq!(
8192                o.len(),
8193                n_embd,
8194                "step35_decode_attn returns post-wo [n_embd]"
8195            );
8196            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
8197        }
8198        Ok(out)
8199    }
8200
8201    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
8202    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
8203    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
8204    #[allow(clippy::too_many_arguments)]
8205    fn full_attn_verify(
8206        &self,
8207        e: &Engine,
8208        fa: &FullAttnLayer,
8209        h: &CudaSlice<f32>,
8210        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8211        pos_d: &CudaSlice<i32>,
8212        t: usize,
8213        cache: &mut Cache,
8214        il: usize,
8215        stream_ctr: Option<&CudaSlice<i32>>,
8216    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8217        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
8218        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
8219        // its own arm. A verify that silently computes different attention than decode defeats the
8220        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
8221        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
8222        // shape and not laziness.
8223        if self.sliding_gated_moe_batch_program() {
8224            if stream_ctr.is_some() {
8225                return Err(
8226                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8227                            cannot express the SWA offset KV view; same root cause as the dc \
8228                            decode refusal) — run spec without the stream arm"
8229                        .into(),
8230                );
8231            }
8232            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
8233        }
8234        let cfg = &self.cfg;
8235        let geometry = cfg.full_attention_geometry_at(il as u32);
8236        let n_head = geometry.n_head as usize;
8237        let n_head_kv = geometry.n_head_kv as usize;
8238        let head_dim = geometry.head_dim_k as usize;
8239        let eps = cfg.rms_eps;
8240        let scale = geometry.attention_scale();
8241        let n_embd = cfg.n_embd as usize;
8242
8243        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
8244        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
8245        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
8246        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
8247        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
8248        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
8249        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
8250        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
8251        let (qf, mut k, v) = {
8252            let mut fused = None;
8253            let qkv_fast =
8254                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
8255            if t == 1 && qkv_fast {
8256                let (hq_o, hd_o);
8257                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8258                    Some(p) => p,
8259                    None => {
8260                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
8261                        (&hq_o, &hd_o)
8262                    }
8263                };
8264                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
8265            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
8266                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
8267                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
8268                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
8269                let (hq_o, hd_o);
8270                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
8271                    Some(p) => p,
8272                    None => {
8273                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
8274                        (&hq_o, &hd_o)
8275                    }
8276                };
8277                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
8278            }
8279            match (fused, h_q8) {
8280                (Some(triple), _) => triple,
8281                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
8282                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
8283                (None, Some((hq, hd))) if qkv_fast => (
8284                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
8285                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
8286                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
8287                ),
8288                (None, _) => (
8289                    e.matmul_decode_exact(&fa.wq, h, t)?,
8290                    e.matmul_decode_exact(&fa.wk, h, t)?,
8291                    e.matmul_decode_exact(&fa.wv, h, t)?,
8292                ),
8293            }
8294        };
8295        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
8296        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
8297        let (mut q, gate) = if gated {
8298            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8299            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
8300            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
8301            (q, Some(gate))
8302        } else {
8303            (qf, None)
8304        };
8305
8306        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
8307        e.rms_norm(
8308            &q,
8309            fa.q_norm.float_data(),
8310            &mut qn,
8311            head_dim,
8312            n_head * t,
8313            eps,
8314        )?;
8315        q = qn;
8316        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
8317        e.rms_norm(
8318            &k,
8319            fa.k_norm.float_data(),
8320            &mut kn,
8321            head_dim,
8322            n_head_kv * t,
8323            eps,
8324        )?;
8325        k = kn;
8326        let rope_dims = geometry.n_rot as usize;
8327        e.rope_neox(
8328            &mut q,
8329            pos_d,
8330            head_dim,
8331            rope_dims,
8332            n_head,
8333            t,
8334            geometry.rope_base,
8335            1.0,
8336        )?;
8337        e.rope_neox(
8338            &mut k,
8339            pos_d,
8340            head_dim,
8341            rope_dims,
8342            n_head_kv,
8343            t,
8344            geometry.rope_base,
8345            1.0,
8346        )?;
8347
8348        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
8349        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
8350        let kvl = cache.kv[il].as_mut().unwrap();
8351        let (kv_dim_k, kv_dim_v, ktb, vtb) =
8352            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
8353        if let Some(ctr) = stream_ctr {
8354            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
8355            // math on a (block, token) grid, documented byte-identical); host len is a stale
8356            // LOWER BOUND under pre-issue (drain reconciles it).
8357            e.append_kv_quantized_rows_dc(
8358                &k,
8359                &v,
8360                &mut kvl.k,
8361                &mut kvl.v,
8362                ctr,
8363                t,
8364                kv_dim_k,
8365                kv_dim_v,
8366                ktb,
8367                vtb,
8368                crate::Engine::kv_fp8_on(),
8369            )?;
8370        } else {
8371            for i in 0..t {
8372                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
8373                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
8374                e.append_kv_quantized_view(
8375                    &k_row,
8376                    &v_row,
8377                    &mut kvl.k,
8378                    &mut kvl.v,
8379                    kvl.len + i,
8380                    kv_dim_k,
8381                    kv_dim_v,
8382                    ktb,
8383                    vtb,
8384                    crate::Engine::kv_fp8_on(),
8385                )?;
8386            }
8387            kvl.len += t;
8388        }
8389
8390        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
8391        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
8392        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
8393        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
8394        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
8395        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
8396        // keys. The verify appends all T tokens first but bounds the key range per row.
8397        //
8398        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
8399        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
8400        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
8401        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
8402        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
8403        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
8404        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
8405        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
8406        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
8407        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
8408        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
8409        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
8410        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
8411        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
8412        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
8413        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
8414        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
8415        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
8416        if let Some(ctr) = stream_ctr {
8417            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
8418            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
8419            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
8420            let upper = kvl.len + t + 64;
8421            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
8422            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
8423            e.fa_decode_rows_dc(
8424                &q,
8425                &k_view,
8426                &v_view,
8427                &mut attn,
8428                head_dim,
8429                n_head,
8430                n_head_kv,
8431                ctr,
8432                upper.min(cache.max_ctx),
8433                t,
8434                scale,
8435                ktb,
8436                vtb,
8437                0,
8438                false,
8439            )?;
8440        } else if spec_lean() && t == 1 {
8441            let t_kv = base_len + 1;
8442            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
8443            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
8444            e.fa_decode_kvmod(
8445                &q,
8446                &k_view,
8447                &v_view,
8448                &mut attn,
8449                head_dim,
8450                n_head,
8451                n_head_kv,
8452                t_kv,
8453                scale,
8454                ktb,
8455                vtb,
8456                crate::Engine::kv_fp8_on(),
8457            )?;
8458        } else if e.fa_rows_eligible(base_len, head_dim) {
8459            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
8460            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
8461            e.fa_decode_rows(
8462                &q,
8463                &k_view,
8464                &v_view,
8465                &mut attn,
8466                head_dim,
8467                n_head,
8468                n_head_kv,
8469                base_len,
8470                t,
8471                scale,
8472                ktb,
8473                vtb,
8474                None,
8475                false,
8476                crate::Engine::kv_fp8_on(),
8477                None,
8478            )?;
8479        } else {
8480            for r in 0..t {
8481                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
8482                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
8483                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
8484                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
8485                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
8486                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
8487                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
8488                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
8489                e.fa_decode_kvmod(
8490                    &q_row,
8491                    &k_view_r,
8492                    &v_view_r,
8493                    &mut attn_row,
8494                    head_dim,
8495                    n_head,
8496                    n_head_kv,
8497                    t_kv_r,
8498                    scale,
8499                    ktb,
8500                    vtb,
8501                    crate::Engine::kv_fp8_on(),
8502                )?;
8503                e.copy_into(
8504                    &mut attn,
8505                    r * n_head * head_dim,
8506                    &attn_row,
8507                    n_head * head_dim,
8508                )?;
8509            }
8510        }
8511
8512        let attn_g = match &gate {
8513            Some(gate) => {
8514                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
8515                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
8516                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
8517                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
8518                ag
8519            }
8520            None => attn,
8521        };
8522        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
8523        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
8524        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
8525    }
8526
8527    /// Context-linear bytes for a plain serving session's trunk cache.
8528    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
8529        crate::cache::cache_bytes_per_token_for_plan(
8530            &self.cfg,
8531            &self.plan,
8532            0,
8533            self.plan.layers.len(),
8534        )
8535    }
8536
8537    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
8538    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
8539        (
8540            self.plain_session_kv_bytes_per_token(),
8541            crate::cache::cache_ring_bytes_per_token_for_plan(
8542                &self.cfg,
8543                &self.plan,
8544                0,
8545                self.plan.layers.len(),
8546            ),
8547            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
8548        )
8549    }
8550
8551    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
8552    /// scratch. With no MTP head this equals the plain coefficient.
8553    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
8554        let scratch = self
8555            .mtp
8556            .iter()
8557            .chain(self.mtp_extra.iter())
8558            .map(|mtp| {
8559                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8560                k + v
8561            })
8562            .sum::<usize>();
8563        self.plain_session_kv_bytes_per_token()
8564            .saturating_add(scratch)
8565    }
8566
8567    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
8568    /// capped by the same SWA ring rows as the trunk.
8569    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
8570        let total = self.spec_session_kv_bytes_per_token();
8571        let (_, mut ring, rows) = self.plain_session_kv_shape();
8572        if rows > 0 {
8573            ring = ring.saturating_add(
8574                self.mtp
8575                    .iter()
8576                    .chain(self.mtp_extra.iter())
8577                    .map(|mtp| {
8578                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
8579                        k + v
8580                    })
8581                    .sum::<usize>(),
8582            );
8583        }
8584        (total, ring, rows)
8585    }
8586
8587    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
8588    /// the NextN head to draft K tokens then verifies them in one batched target forward.
8589    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
8590    /// acceptance rate. `k` = draft length per round.
8591    ///
8592    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
8593    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
8594    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
8595    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
8596    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
8597    /// captured graph references is event-free; the spec loop is strictly single-stream.
8598    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
8599    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
8600    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
8601    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
8602    /// generate_spec_inner2.
8603    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
8604    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
8605    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
8606    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
8607    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
8608    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
8609    pub fn new_session(
8610        &self,
8611        e: &Engine,
8612        max_ctx: usize,
8613    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
8614        Ok(SpecSession {
8615            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
8616            // is the SERVING spec-session path, and with the ppN door open across two cards a
8617            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
8618            // round — the wrong-card class already fixed on the two batched serving paths
8619            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
8620            // branch, same allocations), so single-device behavior is byte-unchanged.
8621            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
8622            scratch: self.new_mtp_scratch(e, max_ctx)?,
8623            committed: Vec::new(),
8624            last_h: None,
8625            next_pred: None,
8626            sctr: 0,
8627            uctr: 0,
8628            draft_ctx: None,
8629            pending_tok: None,
8630            turn_ckpt: None,
8631            telem: SpecTelemetryCounters::default(),
8632            capture_at: None,
8633            boundary_captures: Vec::new(),
8634            ckpt_at: None,
8635        })
8636    }
8637
8638    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
8639    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
8640    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
8641    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
8642    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
8643    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
8644    /// worker always receives a fully-warm continuation session (committed = whole
8645    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
8646    /// boundary logits on the empty-suffix shape).
8647    ///
8648    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
8649    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
8650    /// request, and plain feeds a carried suffix via eager `decode_step` below
8651    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
8652    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
8653    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
8654    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
8655    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
8656    /// burst prime.
8657    ///
8658    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
8659    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
8660    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
8661    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
8662    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
8663    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
8664    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
8665    /// cold session draws from the identical row at counter 0 and then runs its rounds from
8666    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
8667    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
8668    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
8669    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
8670    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
8671    ///
8672    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
8673    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
8674    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
8675    /// and are never routed here.
8676    ///
8677    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
8678    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
8679    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
8680    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
8681    /// entry stays published for the next request.
8682    #[allow(clippy::too_many_arguments)]
8683    pub fn spec_session_from_restored(
8684        &self,
8685        e: &Engine,
8686        mut cache: Cache,
8687        prefix: Vec<u32>,
8688        suffix: &[u32],
8689        draft_k: &CudaSlice<u8>,
8690        draft_v: &CudaSlice<u8>,
8691        draft_k_tok_bytes: usize,
8692        draft_v_tok_bytes: usize,
8693        draft_len: usize,
8694        last_h: &[f32],
8695        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
8696        // when a suffix follows — the feed's own logits are the boundary then.
8697        boundary_logits: &[f32],
8698        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
8699        // ONE place instead of being half-applied by the worker.
8700        sampling: Option<SpecSampling>,
8701        require_anchor: bool,
8702        max_ctx: usize,
8703        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
8704        // prompt position to split the suffix feed at and capture the extended-entry
8705        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
8706        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
8707        // WHY: the prompt-end capture below includes the template's live generation header
8708        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
8709        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
8710        // diverged from every future prompt and the hit boundary FROZE at the first
8711        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
8712        republish_at: Option<usize>,
8713    ) -> Result<SpecSession, (Option<Cache>, String)> {
8714        let pos = prefix.len();
8715        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
8716            Err((Some(cache), msg))
8717        };
8718        if self.mtp.is_none() {
8719            return fail(cache, "no MTP head attached (nothing to draft with)".into());
8720        }
8721        if pos == 0 {
8722            return fail(cache, "empty committed prefix".into());
8723        }
8724        if cache.pos != pos {
8725            let msg = format!(
8726                "restored cache pos {} != restored prefix len {pos}",
8727                cache.pos
8728            );
8729            return fail(cache, msg);
8730        }
8731        if draft_len != pos {
8732            return fail(
8733                cache,
8734                format!("draft plane len {draft_len} != restored prefix len {pos}"),
8735            );
8736        }
8737        if pos + suffix.len() >= max_ctx {
8738            return fail(
8739                cache,
8740                format!(
8741                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
8742                    pos + suffix.len(),
8743                ),
8744            );
8745        }
8746        let mut scratch = match MtpScratch::new(
8747            e,
8748            &self.cfg,
8749            &self.plan,
8750            max_ctx,
8751            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
8752        ) {
8753            Ok(s) => s,
8754            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
8755        };
8756        if scratch.kv.ring.is_some() {
8757            return fail(
8758                cache,
8759                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
8760            );
8761        }
8762        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
8763            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
8764        {
8765            return fail(
8766                cache,
8767                format!(
8768                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
8769                     {}/{} bytes/token (stale entry across a format change)",
8770                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
8771                ),
8772            );
8773        }
8774        if pos > scratch.cap {
8775            return fail(
8776                cache,
8777                format!(
8778                    "draft plane rows {pos} exceed scratch capacity {}",
8779                    scratch.cap
8780                ),
8781            );
8782        }
8783        let kb = pos * draft_k_tok_bytes;
8784        let vb = pos * draft_v_tok_bytes;
8785        if draft_k.len() < kb || draft_v.len() < vb {
8786            return fail(
8787                cache,
8788                format!(
8789                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
8790                    draft_k.len(),
8791                    draft_v.len(),
8792                ),
8793            );
8794        }
8795        if kb > 0 {
8796            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
8797                return fail(cache, format!("draft K restore copy failed: {err}"));
8798            }
8799        }
8800        if vb > 0 {
8801            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
8802                return fail(cache, format!("draft V restore copy failed: {err}"));
8803            }
8804        }
8805        if let Err(err) = scratch.set_len(e, pos) {
8806            return fail(cache, format!("draft scratch len set failed: {err}"));
8807        }
8808        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
8809            // anchor upload failure is acceptance-only when a suffix feed follows (fill
8810            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
8811            // burst entry asserts committed + last_h + next_pred) — the caller says which.
8812            e.htod(last_h).ok()
8813        } else {
8814            None
8815        };
8816        if require_anchor && last_h_dev.is_none() {
8817            return fail(
8818                cache,
8819                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
8820            );
8821        }
8822        let mut committed = prefix;
8823        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
8824        // what the empty-suffix continuation assert in the burst entry requires.
8825        let next_pred;
8826        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
8827        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
8828        // drawing its own first token from the same row.
8829        let mut sctr = 0u32;
8830        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
8831        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
8832        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
8833        // after the suffix joins `committed` below.
8834        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
8835        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
8836        if !suffix.is_empty() {
8837            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
8838            // From here on the trunk cache mutates: failures return Err((None, _)) and
8839            // the worker serves the request cold-plain instead of reusing the carrier.
8840            let dirty =
8841                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
8842            let n_embd = self.cfg.n_embd as usize;
8843            let t = suffix.len();
8844            let mut h_rows = match e.uninit(t * n_embd) {
8845                Ok(b) => b,
8846                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
8847            };
8848            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
8849            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
8850            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
8851            let b_rel = republish_at
8852                .and_then(|abs| abs.checked_sub(pos))
8853                .filter(|&r| r > 0 && r < t);
8854            let mut feed_logits = Vec::new();
8855            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
8856                || e.frozen_cpu_experts_prefer_tokenwise_prime();
8857            let mut fed = 0usize;
8858            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
8859                if seg_end <= fed {
8860                    continue;
8861                }
8862                let seg = &suffix[fed..seg_end];
8863                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
8864                if batched {
8865                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
8866                    // queued after this segment ride `queued_after` so Step35 arm selection
8867                    // stays keyed to the request's end (tick-seg law).
8868                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
8869                        Ok((l, _h_seed, hiddens)) => {
8870                            if let Err(err) =
8871                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
8872                            {
8873                                return dirty(format!("suffix hidden copy: {err}"));
8874                            }
8875                            feed_logits = l;
8876                        }
8877                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
8878                    }
8879                } else {
8880                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
8881                    for (i, &tok) in seg.iter().enumerate() {
8882                        match self.decode_step_h(e, tok, &mut cache) {
8883                            Ok((l, h)) => {
8884                                if let Err(err) =
8885                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
8886                                {
8887                                    return dirty(format!("suffix hidden copy: {err}"));
8888                                }
8889                                feed_logits = l;
8890                            }
8891                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
8892                        }
8893                    }
8894                }
8895                fed = seg_end;
8896                if Some(seg_end) == b_rel {
8897                    // The stable pre-generation boundary: capture the extended-entry
8898                    // publication AND this session's own turn checkpoint here instead of at
8899                    // prompt-end (both would otherwise carry the volatile live-header tail
8900                    // the next re-render replaces). Failure silent, turn_ckpt convention.
8901                    debug_assert_eq!(
8902                        cache.pos,
8903                        pos + seg_end,
8904                        "stable-boundary capture off the feed split"
8905                    );
8906                    if spec_restore_republish_on() {
8907                        if let Ok(snap) = cache.snapshot(e) {
8908                            boundary_captures.push(SpecBoundaryCapture {
8909                                snap,
8910                                pos: pos + seg_end,
8911                                logits: feed_logits.clone(),
8912                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
8913                            });
8914                        }
8915                    }
8916                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
8917                        e.uninit(n_embd).and_then(|mut a| {
8918                            e.copy_view_into(
8919                                &mut a,
8920                                0,
8921                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
8922                                n_embd,
8923                            )?;
8924                            Ok(a)
8925                        });
8926                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
8927                        restored_turn_ckpt = Some(SpecCheckpoint {
8928                            snap,
8929                            pos: pos + seg_end,
8930                            last_h,
8931                        });
8932                    }
8933                }
8934            }
8935            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
8936            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
8937            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
8938            // with T). Fill failures are acceptance-only — truncate to the restored rows
8939            // and continue; the burst's own set_len keeps the invariant.
8940            let mtp = self.mtp.as_ref().expect("mtp checked above");
8941            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
8942            let embd_gpu = if spec_host_embd() {
8943                None
8944            } else {
8945                Some(
8946                    self.embd_gpu
8947                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
8948                )
8949            };
8950            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
8951            let fill_chunk = 4096usize;
8952            let mut filled = true;
8953            let mut start = 0usize;
8954            'fill: while start < t {
8955                let end = (start + fill_chunk).min(t);
8956                let tc = end - start;
8957                let Ok(mut phs) = e.zeros(tc * n_embd) else {
8958                    filled = false;
8959                    break 'fill;
8960                };
8961                let (src_lo, dst_off, n_copy) = if start == 0 {
8962                    (0, n_embd, (tc - 1) * n_embd)
8963                } else {
8964                    ((start - 1) * n_embd, 0, tc * n_embd)
8965                };
8966                if start == 0 {
8967                    if let Some(lh) = last_h_dev.as_ref() {
8968                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
8969                            filled = false;
8970                            break 'fill;
8971                        }
8972                    }
8973                }
8974                if n_copy > 0
8975                    && e.copy_view_into(
8976                        &mut phs,
8977                        dst_off,
8978                        &h_rows.slice(src_lo..src_lo + n_copy),
8979                        n_copy,
8980                    )
8981                    .is_err()
8982                {
8983                    filled = false;
8984                    break 'fill;
8985                }
8986                if self
8987                    .mtp_kv_fill_all(
8988                        e,
8989                        &suffix[start..end],
8990                        &phs,
8991                        pos + start,
8992                        &mut scratch,
8993                        embd_dev,
8994                    )
8995                    .is_err()
8996                {
8997                    filled = false;
8998                    break 'fill;
8999                }
9000                start = end;
9001            }
9002            if !filled {
9003                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9004                // so keep only the restored rows resident and let verify arbitrate.
9005                if let Err(err) = scratch.set_len(e, pos) {
9006                    return dirty(format!("scratch truncation after failed fill: {err}"));
9007                }
9008            }
9009            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9010            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9011            // finding (d)). Pre-lane, publication was armed only for COLD sessions
9012            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9013            // non-continuation burst — but a converted hit's first burst IS a continuation,
9014            // so a growing conversation learned exactly ONE boundary and turn 3 could never
9015            // hit a longer prefix than turn 2 did.
9016            //
9017            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9018            // line — the trunk is primed over the whole prompt, nothing is generated, and the
9019            // draft plane rows [0..prompt) are filled just above. That is a complete
9020            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9021            // publishes; the worker's existing publication sweep picks it up because it is
9022            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9023            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9024            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9025            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9026            // publication is an optimization, never a correctness dependency.
9027            //
9028            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9029            // entry's tail is the live generation header the next re-render replaces, so on a
9030            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9031            // the stable-boundary capture above IS this publication, minus the poisoned tail.
9032            if spec_restore_republish_on() && boundary_captures.is_empty() {
9033                debug_assert_eq!(
9034                    cache.pos,
9035                    pos + t,
9036                    "extended-entry capture must sit at the restored session's prompt end",
9037                );
9038                if let Ok(snap) = cache.snapshot(e) {
9039                    boundary_captures.push(SpecBoundaryCapture {
9040                        snap,
9041                        pos: pos + t,
9042                        logits: feed_logits.clone(),
9043                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9044                    });
9045                }
9046            }
9047            // continuation seed: the feed's boundary logits ARE the plain path's boundary
9048            // logits (same program), so greedy's argmax here is plain's first emitted token,
9049            // and the sampled draw is the cold sampled session's own first token.
9050            next_pred = Some(if sampled {
9051                let sp = sampling.expect("sampled implies a sampler");
9052                // `committed` is still the restored prefix here; the suffix joins it below —
9053                // so this is the last-N window over the WHOLE prompt, exactly the cold
9054                // session's own window at its first token.
9055                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9056                match sample_boundary_token(
9057                    e,
9058                    &feed_logits,
9059                    &sp,
9060                    &hist,
9061                    &mut sctr,
9062                    "restore-suffix-feed",
9063                ) {
9064                    Ok(t) => t,
9065                    // the trunk is already fed: hand nothing back, the worker serves the
9066                    // request cold-plain. Never fall back to an argmax — that would put a
9067                    // greedy token in a sampled stream to save a slow path.
9068                    Err(err) => {
9069                        return dirty(format!("boundary token draw failed: {err}"));
9070                    }
9071                }
9072            } else {
9073                argmax(&feed_logits) as u32
9074            });
9075            let mut lh = match e.uninit(n_embd) {
9076                Ok(b) => b,
9077                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9078            };
9079            if let Err(err) = e.copy_view_into(
9080                &mut lh,
9081                0,
9082                &h_rows.slice((t - 1) * n_embd..t * n_embd),
9083                n_embd,
9084            ) {
9085                return dirty(format!("boundary hidden copy: {err}"));
9086            }
9087            last_h_dev = Some(lh);
9088            committed.extend_from_slice(suffix);
9089        } else {
9090            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
9091            // ENTRY's boundary logits are the boundary row, and this is the token the cold
9092            // session emits from that same row. Owned here rather than in the worker so the
9093            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
9094            if boundary_logits.is_empty() {
9095                return fail(
9096                    cache,
9097                    "full-cover restore without the entry's boundary logits".into(),
9098                );
9099            }
9100            next_pred = Some(if sampled {
9101                let sp = sampling.expect("sampled implies a sampler");
9102                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
9103                match sample_boundary_token(
9104                    e,
9105                    boundary_logits,
9106                    &sp,
9107                    &hist,
9108                    &mut sctr,
9109                    "restore-full-cover",
9110                ) {
9111                    Ok(t) => t,
9112                    // nothing has been mutated on this shape — hand the carrier back and let
9113                    // the hit serve PLAIN (the banked pre-lane path).
9114                    Err(err) => {
9115                        return fail(cache, format!("boundary token draw failed: {err}"));
9116                    }
9117                }
9118            } else {
9119                argmax(boundary_logits) as u32
9120            });
9121        }
9122        Ok(SpecSession {
9123            cache,
9124            scratch,
9125            committed,
9126            last_h: last_h_dev,
9127            next_pred,
9128            sctr,
9129            uctr: 0,
9130            draft_ctx: None,
9131            pending_tok: None,
9132            // Stable-boundary capture from the split feed above (None on the legacy shape):
9133            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
9134            // affinity probe declined ("no turn checkpoint retained") and the conversation
9135            // fell back to the frozen prefix entry forever.
9136            turn_ckpt: restored_turn_ckpt,
9137            telem: SpecTelemetryCounters::default(),
9138            capture_at: None,
9139            boundary_captures,
9140            ckpt_at: None,
9141        })
9142    }
9143
9144    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
9145    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
9146    /// snapshot, or draft-KV row that only corrupts the following round.
9147    pub fn optipipe_compare_session_state(
9148        &self,
9149        e: &Engine,
9150        reference: &SpecSession,
9151        candidate: &SpecSession,
9152    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
9153        fn fail(what: &str) -> Box<dyn std::error::Error> {
9154            format!("optipipe state mismatch: {what}").into()
9155        }
9156        fn same_f32(a: &[f32], b: &[f32]) -> bool {
9157            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
9158        }
9159        fn compare_layers(
9160            es: &Engine,
9161            range: std::ops::Range<usize>,
9162            reference: &SpecSession,
9163            candidate: &SpecSession,
9164            report: &mut OptiForkStateIdentity,
9165        ) -> Result<(), Box<dyn std::error::Error>> {
9166            for il in range {
9167                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
9168                    (Some(a), Some(b)) => {
9169                        if a.len != b.len {
9170                            return Err(fail(&format!(
9171                                "layer {il} host KV len {} != {}",
9172                                a.len, b.len
9173                            )));
9174                        }
9175                        let ad = es.dtoh_i32(&a.len_d)?;
9176                        let bd = es.dtoh_i32(&b.len_d)?;
9177                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
9178                            return Err(fail(&format!(
9179                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
9180                                a.len,
9181                            )));
9182                        }
9183                        let kb = a.len * a.k_tok_bytes;
9184                        let vb = a.len * a.v_tok_bytes;
9185                        if kb > 0 {
9186                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
9187                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
9188                            if ak != bk {
9189                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
9190                                return Err(fail(&format!(
9191                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
9192                                    at / a.k_tok_bytes,
9193                                    at % a.k_tok_bytes,
9194                                    ak[at],
9195                                    bk[at],
9196                                )));
9197                            }
9198                        }
9199                        if vb > 0 {
9200                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
9201                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
9202                            if av != bv {
9203                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
9204                                return Err(fail(&format!(
9205                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
9206                                    at / a.v_tok_bytes,
9207                                    at % a.v_tok_bytes,
9208                                    av[at],
9209                                    bv[at],
9210                                )));
9211                            }
9212                        }
9213                        report.trunk_kv_bytes += kb + vb;
9214                    }
9215                    (None, None) => {}
9216                    _ => return Err(fail(&format!("layer {il} KV presence"))),
9217                }
9218                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
9219                    (Some(a), Some(b)) => {
9220                        let ac = es.dtoh(&a.conv_state)?;
9221                        let bc = es.dtoh(&b.conv_state)?;
9222                        if !same_f32(&ac, &bc) {
9223                            return Err(fail(&format!("layer {il} conv state")));
9224                        }
9225                        let as_ = es.dtoh(&a.ssm_state)?;
9226                        let bs = es.dtoh(&b.ssm_state)?;
9227                        if !same_f32(&as_, &bs) {
9228                            return Err(fail(&format!("layer {il} SSM state")));
9229                        }
9230                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
9231                    }
9232                    (None, None) => {}
9233                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
9234                }
9235            }
9236            Ok(())
9237        }
9238
9239        if reference.committed != candidate.committed {
9240            return Err(fail("committed token ids"));
9241        }
9242        if reference.cache.pos != candidate.cache.pos
9243            || reference.cache.max_ctx != candidate.cache.max_ctx
9244        {
9245            return Err(fail("cache pos/capacity"));
9246        }
9247        if reference.pending_tok != candidate.pending_tok
9248            || reference.next_pred != candidate.next_pred
9249            || reference.sctr != candidate.sctr
9250            || reference.uctr != candidate.uctr
9251        {
9252            return Err(fail("pending/prediction/counter tail"));
9253        }
9254
9255        let mut report = OptiForkStateIdentity::default();
9256        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
9257            let rt = crate::pp::PpNRt::get(e)?;
9258            for stage in 0..rt.n_stages() {
9259                let _scope = rt.enter(stage);
9260                compare_layers(
9261                    rt.engine(stage, e),
9262                    fence[stage]..fence[stage + 1],
9263                    reference,
9264                    candidate,
9265                    &mut report,
9266                )?;
9267            }
9268        } else {
9269            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
9270        }
9271
9272        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
9273            return Err(fail("draft scratch plane count"));
9274        }
9275        for index in 0..reference.scratch.plane_count() {
9276            let (a, _) = reference.scratch.plane(index);
9277            let (b, _) = candidate.scratch.plane(index);
9278            if a.len != b.len
9279                || a.kv_dim_k != b.kv_dim_k
9280                || a.kv_dim_v != b.kv_dim_v
9281                || a.k_tok_bytes != b.k_tok_bytes
9282                || a.v_tok_bytes != b.v_tok_bytes
9283                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
9284            {
9285                return Err(fail(&format!("draft scratch plane {index} length/layout")));
9286            }
9287            let kb = a.len * a.k_tok_bytes;
9288            let vb = a.len * a.v_tok_bytes;
9289            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
9290                return Err(fail(&format!("draft scratch plane {index} K bytes")));
9291            }
9292            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
9293                return Err(fail(&format!("draft scratch plane {index} V bytes")));
9294            }
9295            report.scratch_kv_bytes += kb + vb;
9296        }
9297
9298        match (&reference.last_h, &candidate.last_h) {
9299            (Some(a), Some(b)) => {
9300                let ah = e.dtoh(a)?;
9301                let bh = e.dtoh(b)?;
9302                if !same_f32(&ah, &bh) {
9303                    return Err(fail("last hidden/seed bytes"));
9304                }
9305                report.hidden_bytes = ah.len() * 4;
9306            }
9307            (None, None) => {}
9308            _ => return Err(fail("last hidden/seed presence")),
9309        }
9310        Ok(report)
9311    }
9312
9313    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
9314    /// retained prompt-end checkpoint, so a request whose prompt matches
9315    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
9316    ///
9317    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
9318    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
9319    /// restored from the device copy taken there, draft scratch length reset, `committed`
9320    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
9321    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
9322    /// every burst after it are identical to a cold run of the same token stream — the
9323    /// committed-tokens-authoritative contract.
9324    ///
9325    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
9326    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
9327    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
9328    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
9329    /// (the scratch KV, the resident embedding), none of which the rewind moves.
9330    ///
9331    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
9332    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
9333    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
9334    pub fn spec_rewind_to_checkpoint(
9335        &self,
9336        e: &Engine,
9337        sess: &mut SpecSession,
9338    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9339        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
9340            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
9341        }) {
9342            return Err(
9343                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
9344            );
9345        }
9346        let Some(ckpt) = sess.turn_ckpt.take() else {
9347            return Ok(None);
9348        };
9349        assert!(
9350            ckpt.pos <= sess.committed.len(),
9351            "checkpoint past committed ({} > {})",
9352            ckpt.pos,
9353            sess.committed.len()
9354        );
9355        // Restore through each layer's owning engine. A single primary-engine rollback is not
9356        // sufficient when the serving cache is stage-owned under cross-device PP.
9357        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
9358        debug_assert_eq!(
9359            sess.cache.pos, ckpt.pos,
9360            "rollback landed off the checkpoint"
9361        );
9362        sess.scratch.set_len(e, ckpt.pos)?;
9363        sess.committed.truncate(ckpt.pos);
9364        sess.last_h = Some(ckpt.last_h);
9365        sess.next_pred = None;
9366        sess.pending_tok = None;
9367        Ok(Some(ckpt.pos))
9368    }
9369
9370    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
9371    /// checkpoint without re-priming the checkpoint prefix.
9372    ///
9373    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
9374    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
9375    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
9376    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
9377    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
9378    ///
9379    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
9380    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
9381    pub fn spec_grow_and_rewind_to_checkpoint(
9382        &self,
9383        e: &Engine,
9384        sess: &mut SpecSession,
9385        target_cap: usize,
9386    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
9387        if target_cap <= sess.cache.max_ctx {
9388            return self.spec_rewind_to_checkpoint(e, sess);
9389        }
9390        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
9391            return Ok(None);
9392        };
9393        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
9394            return Err(format!(
9395                "checkpoint pos {} outside committed length {}",
9396                ckpt.pos,
9397                sess.committed.len(),
9398            )
9399            .into());
9400        }
9401        if ckpt.pos > target_cap {
9402            return Err(format!(
9403                "checkpoint pos {} exceeds grown capacity {target_cap}",
9404                ckpt.pos,
9405            )
9406            .into());
9407        }
9408
9409        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
9410        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
9411        crate::pp::restore_cache_checkpoint(
9412            e,
9413            self,
9414            Some(&sess.cache),
9415            &mut grown_cache,
9416            &ckpt.snap,
9417        )?;
9418
9419        if sess.scratch.plane_count() != grown_scratch.plane_count() {
9420            return Err("checkpoint draft plane count mismatch".into());
9421        }
9422        for index in 0..sess.scratch.plane_count() {
9423            let (src, _) = sess.scratch.plane(index);
9424            let (dst, _) = grown_scratch.plane_mut(index);
9425            if ckpt.pos > src.len
9426                || src.kv_dim_k != dst.kv_dim_k
9427                || src.kv_dim_v != dst.kv_dim_v
9428                || src.k_tok_bytes != dst.k_tok_bytes
9429                || src.v_tok_bytes != dst.v_tok_bytes
9430            {
9431                return Err(format!(
9432                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
9433                    ckpt.pos, src.len,
9434                )
9435                .into());
9436            }
9437            let kb = ckpt.pos * src.k_tok_bytes;
9438            let vb = ckpt.pos * src.v_tok_bytes;
9439            if kb > 0 {
9440                e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
9441            }
9442            if vb > 0 {
9443                e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
9444            }
9445        }
9446        grown_scratch.set_len(e, ckpt.pos)?;
9447        // The old scratch is dropped immediately after publication below. Bound its D2D reads
9448        // first; growth happens once per rewritten turn, outside the decode hot loop.
9449        e.stream().synchronize()?;
9450
9451        let ckpt = sess
9452            .turn_ckpt
9453            .take()
9454            .expect("checkpoint remained present through transactional grow");
9455        let pos = ckpt.pos;
9456        sess.cache = grown_cache;
9457        sess.scratch = grown_scratch;
9458        sess.committed.truncate(pos);
9459        sess.last_h = Some(ckpt.last_h);
9460        sess.next_pred = None;
9461        sess.pending_tok = None;
9462        sess.draft_ctx = None;
9463        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
9464        debug_assert!(
9465            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
9466            "grown draft rewind landed off checkpoint"
9467        );
9468        Ok(Some(pos))
9469    }
9470
9471    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
9472    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
9473    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
9474    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
9475    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
9476    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
9477    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
9478    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
9479    /// park-time flush is a future request whose sampler is not knowable here (residual
9480    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
9481    pub fn spec_flush_pending(
9482        &self,
9483        e: &Engine,
9484        sess: &mut SpecSession,
9485        sampling: Option<SpecSampling>,
9486    ) -> Result<(), Box<dyn std::error::Error>> {
9487        let Some(b) = sess.pending_tok.take() else {
9488            return Ok(());
9489        };
9490        if self.mtp.is_none() {
9491            return Err("pending carry requires an MTP head".into());
9492        }
9493        let n_embd = self.cfg.n_embd as usize;
9494        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9495        let embd_gpu = if spec_host_embd() {
9496            None
9497        } else {
9498            Some(
9499                self.embd_gpu
9500                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9501            )
9502        };
9503        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9504        let pos_b = sess.cache.pos;
9505        sess.scratch.set_len(e, pos_b)?;
9506        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
9507        sess.next_pred = Some(match sampling {
9508            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
9509                // window includes `b` itself: it is committed by this pass, and the pre-lane
9510                // code never counted a boundary token in the penalty history at all.
9511                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
9512                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
9513            }
9514            _ => argmax(&lg_b) as u32,
9515        });
9516        let anchor = sess
9517            .last_h
9518            .as_ref()
9519            .expect("pending carry requires last_h (the predecessor-row anchor)");
9520        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
9521        sess.last_h = Some(hb);
9522        sess.committed.push(b);
9523        Ok(())
9524    }
9525
9526    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
9527    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
9528    /// rounds through that same graph. Other model families keep their eager T=1 contract.
9529    fn spec_target_step_h(
9530        &self,
9531        e: &Engine,
9532        token: u32,
9533        cache: &mut Cache,
9534    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
9535        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
9536            return self.decode_step_h(e, token, cache);
9537        }
9538        let pos0 = cache.pos;
9539        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
9540        Ok((e.dtoh(&logits)?, hidden))
9541    }
9542
9543    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
9544    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
9545    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
9546    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
9547    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
9548    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
9549    /// dispatch sites cannot drift apart again.
9550    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
9551    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
9552    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
9553    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
9554    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
9555    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
9556    fn mtp_graph_capturable(&self) -> bool {
9557        self.mtp
9558            .as_ref()
9559            .map(|m| match &m.ffn {
9560                crate::hybrid::Ffn::Dense { .. } => true,
9561                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
9562            })
9563            .unwrap_or(false)
9564    }
9565
9566    fn batched_serving_numeric_class(&self) -> bool {
9567        self.plan
9568            .trunk_operations()
9569            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
9570    }
9571
9572    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
9573    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
9574    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
9575    /// keeping the engine's own version structural rather than name-based means a new
9576    /// checkpoint of the same shape inherits the default, and a different shape does not.
9577    fn vgraph_family_default(&self) -> bool {
9578        let has_linear = self
9579            .layers
9580            .iter()
9581            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
9582        let has_moe = self
9583            .layers
9584            .iter()
9585            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
9586        has_linear && has_moe
9587    }
9588
9589    fn sliding_gated_moe_batch_program(&self) -> bool {
9590        self.uses_sliding_gated_moe_program()
9591    }
9592
9593    fn gemma_batch_program(&self) -> bool {
9594        self.uses_gemma_program()
9595    }
9596
9597    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
9598    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
9599    /// session already exist.
9600    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
9601        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
9602            || !spec_devacc()
9603            || spec_replay_env_enabled()
9604            || spec_stream()
9605            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
9606            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
9607            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
9608            || std::env::var("MEMRA_SPEC_PMIN")
9609                .ok()
9610                .and_then(|v| v.parse::<f32>().ok())
9611                .unwrap_or(0.0)
9612                > 0.0
9613            || self.is_gemma4_e4b()
9614            || self.gemma_batch_program()
9615            || self.mtp.is_none()
9616            || !self.mtp_extra.is_empty()
9617        {
9618            return false;
9619        }
9620        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
9621            return false;
9622        };
9623        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
9624            return false;
9625        }
9626        crate::pp::PpNRt::get(e)
9627            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
9628            .unwrap_or(false)
9629    }
9630
9631    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
9632    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
9633    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
9634    #[allow(clippy::too_many_arguments)]
9635    pub fn generate_spec_session_pair(
9636        &self,
9637        e: &Engine,
9638        sess_a: &mut SpecSession,
9639        max_new_a: usize,
9640        k_a: usize,
9641        sess_b: &mut SpecSession,
9642        max_new_b: usize,
9643        k_b: usize,
9644    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
9645    {
9646        if !self.spec_pipe_available(e) {
9647            return Err("two-session speculative pipeline is outside its reduced matrix".into());
9648        }
9649        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
9650            return Err(
9651                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
9652            );
9653        }
9654        for sess in [&*sess_a, &*sess_b] {
9655            if sess.committed.is_empty()
9656                || sess.last_h.is_none()
9657                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
9658            {
9659                return Err("two-session speculative pipeline requires warm continuations".into());
9660            }
9661        }
9662
9663        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9664            && !spec_host_embd()
9665            && self.mtp_graph_capturable()
9666            && self.mtp_extra.is_empty()
9667            && !crate::model::full_prec_enabled();
9668        let graph_a = graph_ok && k_a + 2 < 96;
9669        let graph_b = graph_ok && k_b + 2 < 96;
9670        let was_tracking = e.ctx().is_event_tracking();
9671        if (graph_a || graph_b) && was_tracking {
9672            unsafe {
9673                e.ctx().disable_event_tracking();
9674            }
9675        }
9676
9677        static LOGGED: std::sync::Once = std::sync::Once::new();
9678        LOGGED.call_once(|| {
9679            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
9680        });
9681        let sync = std::sync::Arc::new(SpecPipeSync::new());
9682        let lane_a = SpecPipeLane {
9683            sync: sync.clone(),
9684            lane: 0,
9685        };
9686        let lane_b = SpecPipeLane { sync, lane: 1 };
9687        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
9688        let (result_a, result_b) = std::thread::scope(|scope| {
9689            let b = scope.spawn(move || {
9690                let mut finish = SpecPipeFinish::new(&lane_b);
9691                let sess_b = unsafe { sess_b_ptr.get_mut() };
9692                let result = e
9693                    .ctx()
9694                    .bind_to_thread()
9695                    .map_err(|err| err.to_string())
9696                    .and_then(|_| {
9697                        self.generate_spec_inner2(
9698                            e,
9699                            &[],
9700                            max_new_b,
9701                            k_b,
9702                            graph_b,
9703                            Some(sess_b),
9704                            None,
9705                            None,
9706                            None,
9707                            None,
9708                            Some(&lane_b),
9709                        )
9710                        .map_err(|err| err.to_string())
9711                    });
9712                finish.close(result.is_err());
9713                result
9714            });
9715            let mut finish = SpecPipeFinish::new(&lane_a);
9716            let result_a = self.generate_spec_inner2(
9717                e,
9718                &[],
9719                max_new_a,
9720                k_a,
9721                graph_a,
9722                Some(sess_a),
9723                None,
9724                None,
9725                None,
9726                None,
9727                Some(&lane_a),
9728            );
9729            finish.close(result_a.is_err());
9730            let result_b = b
9731                .join()
9732                .map_err(|_| "paired speculative session B panicked".to_string())
9733                .and_then(|r| r);
9734            (result_a, result_b)
9735        });
9736
9737        if (graph_a || graph_b) && was_tracking {
9738            unsafe {
9739                e.ctx().enable_event_tracking();
9740            }
9741        }
9742        let result_a = result_a?;
9743        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
9744        Ok((result_a, result_b))
9745    }
9746
9747    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
9748    /// message rendered through the chat template continuation). Returns (new tokens emitted,
9749    /// drafted, accepted); session.committed grows by suffix + emitted.
9750    pub fn generate_spec_session(
9751        &self,
9752        e: &Engine,
9753        sess: &mut SpecSession,
9754        suffix: &[u32],
9755        max_new: usize,
9756        k: usize,
9757    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9758        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
9759    }
9760
9761    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
9762    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
9763    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
9764    /// for the filtered target (feat/filtered-spec).
9765    ///
9766    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
9767    /// output — once right after the prime's first token, then once per round commit — so a
9768    /// streaming caller can flush text at round cadence instead of once per burst. The slices
9769    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
9770    /// timing only: token bytes, session state, and exactness are untouched.
9771    ///
9772    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
9773    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
9774    /// the caller's scheduler regains control without waiting the burst out. Burst size is
9775    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
9776    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
9777    /// drains and the defensive tail flush can land with nothing new committed).
9778    #[allow(clippy::too_many_arguments)]
9779    pub fn generate_spec_session_sampled(
9780        &self,
9781        e: &Engine,
9782        sess: &mut SpecSession,
9783        suffix: &[u32],
9784        max_new: usize,
9785        k: usize,
9786        sampling: Option<SpecSampling>,
9787        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9788    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9789        self.generate_spec_session_sampled_prime_split(
9790            e, sess, suffix, max_new, k, sampling, None, on_commit,
9791        )
9792    }
9793
9794    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
9795    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
9796    /// pass `None` and stay on the existing zero-prime path.
9797    #[allow(clippy::too_many_arguments)]
9798    pub fn generate_spec_session_sampled_prime_split(
9799        &self,
9800        e: &Engine,
9801        sess: &mut SpecSession,
9802        suffix: &[u32],
9803        max_new: usize,
9804        k: usize,
9805        sampling: Option<SpecSampling>,
9806        prime_split: Option<usize>,
9807        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9808    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9809        self.generate_spec_session_constrained_prime_split(
9810            e,
9811            sess,
9812            suffix,
9813            max_new,
9814            k,
9815            sampling,
9816            None,
9817            prime_split,
9818            on_commit,
9819        )
9820    }
9821
9822    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
9823    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
9824    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
9825    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
9826    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
9827    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
9828    /// may drop (drafter is unconstrained); that is measured, not hidden.
9829    #[allow(clippy::too_many_arguments)]
9830    pub fn generate_spec_session_constrained(
9831        &self,
9832        e: &Engine,
9833        sess: &mut SpecSession,
9834        suffix: &[u32],
9835        max_new: usize,
9836        k: usize,
9837        sampling: Option<SpecSampling>,
9838        constraint: Option<&mut dyn SpecConstraint>,
9839        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9840    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9841        self.generate_spec_session_constrained_prime_split(
9842            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
9843        )
9844    }
9845
9846    #[allow(clippy::too_many_arguments)]
9847    pub fn generate_spec_session_constrained_prime_split(
9848        &self,
9849        e: &Engine,
9850        sess: &mut SpecSession,
9851        suffix: &[u32],
9852        max_new: usize,
9853        k: usize,
9854        sampling: Option<SpecSampling>,
9855        constraint: Option<&mut dyn SpecConstraint>,
9856        prime_split: Option<usize>,
9857        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9858    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9859        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
9860            return Err(
9861                "constrained spec decode is greedy-only (worker routes sampled \
9862                        constrained to plain decode)"
9863                    .into(),
9864            );
9865        }
9866        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
9867        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
9868        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
9869        // serve continuation case — consume the carry in-loop with zero solo passes.
9870        if sess.pending_tok.is_some()
9871            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
9872        {
9873            self.spec_flush_pending(e, sess, sampling)?;
9874        }
9875
9876        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
9877        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
9878        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
9879        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9880            && !spec_host_embd()
9881            && self.mtp_graph_capturable()
9882            && self.mtp_extra.is_empty()
9883            && k + 2 < 96
9884            && !crate::model::full_prec_enabled();
9885        let was_tracking = e.ctx().is_event_tracking();
9886        if graph_draft && was_tracking {
9887            unsafe {
9888                e.ctx().disable_event_tracking();
9889            }
9890        }
9891        let r = self.generate_spec_inner2(
9892            e,
9893            suffix,
9894            max_new,
9895            k,
9896            graph_draft,
9897            Some(sess),
9898            sampling,
9899            constraint,
9900            on_commit,
9901            prime_split,
9902            None,
9903        );
9904        if graph_draft && was_tracking {
9905            unsafe {
9906                e.ctx().enable_event_tracking();
9907            }
9908        }
9909        let (out, d, a) = r?;
9910        Ok((out, d, a))
9911    }
9912
9913    pub fn generate_spec(
9914        &self,
9915        e: &Engine,
9916        prompt: &[u32],
9917        max_new: usize,
9918        k: usize,
9919    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9920        if crate::pp::pp_cuts(self.layers.len()).is_some()
9921            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
9922        {
9923            return Err("pipeline rewrite is not qualified for speculative decode".into());
9924        }
9925        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
9926            return Err("speculative rewrite is not qualified for this ModelPlan".into());
9927        }
9928        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
9929        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
9930        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
9931            && !spec_host_embd()
9932            && self.mtp_graph_capturable()
9933            && self.mtp_extra.is_empty()
9934            && k + 2 < 96
9935            && !crate::model::full_prec_enabled();
9936        if !graph_draft {
9937            return self.generate_spec_inner2(
9938                e, prompt, max_new, k, false, None, None, None, None, None, None,
9939            );
9940        }
9941        let was_tracking = e.ctx().is_event_tracking();
9942        if was_tracking {
9943            unsafe {
9944                e.ctx().disable_event_tracking();
9945            }
9946        }
9947        let r = self.generate_spec_inner2(
9948            e, prompt, max_new, k, true, None, None, None, None, None, None,
9949        );
9950        if was_tracking {
9951            unsafe {
9952                e.ctx().enable_event_tracking();
9953            }
9954        }
9955        r
9956    }
9957
9958    fn generate_spec_inner2(
9959        &self,
9960        e: &Engine,
9961        prompt: &[u32],
9962        max_new: usize,
9963        k: usize,
9964        graph_draft: bool,
9965        mut sess: Option<&mut SpecSession>,
9966        sampling: Option<SpecSampling>,
9967        mut constraint: Option<&mut dyn SpecConstraint>,
9968        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
9969        prime_split: Option<usize>,
9970        pipe: Option<&SpecPipeLane>,
9971    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
9972        assert!(k >= 1, "k must be >= 1");
9973        if let Some(p) = pipe {
9974            p.setup_begin()?;
9975        }
9976        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
9977        let mut flushed = 0usize;
9978        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
9979        // at the next round boundary (same exit as max_new reached — the session tail runs).
9980        // Initialized by the unconditional post-prime flush below.
9981        let mut keep_going;
9982        let mtp = self
9983            .mtp
9984            .as_ref()
9985            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
9986        let n_vocab = self.output.out_features();
9987        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
9988        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
9989        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
9990        let d_vocab = mtp
9991            .shared_head_head
9992            .as_ref()
9993            .unwrap_or(&self.output)
9994            .out_features();
9995        if !self.mtp_extra.is_empty() {
9996            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
9997                || self.plan.mtp_blocks.len() != self.mtp_head_count()
9998                || mtp.d2t.is_some()
9999            {
10000                return Err(
10001                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
10002                );
10003            }
10004            for (offset, head) in self.mtp_extra.iter().enumerate() {
10005                if head.d2t.is_some()
10006                    || head
10007                        .shared_head_head
10008                        .as_ref()
10009                        .unwrap_or(&self.output)
10010                        .out_features()
10011                        != d_vocab
10012                {
10013                    return Err(format!(
10014                        "embedded MTP head {} has incompatible draft vocabulary",
10015                        offset + 1
10016                    )
10017                    .into());
10018                }
10019            }
10020            eprintln!(
10021                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
10022                self.mtp_head_count()
10023            );
10024        }
10025        let n_embd = self.cfg.n_embd as usize;
10026        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
10027        // already committed (their state is in the caches); 0 = fresh single-shot call.
10028        let session_mode = sess.is_some();
10029        let max_ctx = match sess.as_ref() {
10030            Some(s) => s.cache.max_ctx,
10031            None => prompt.len() + max_new + k + 8,
10032        };
10033        let mut own_cache;
10034        let mut own_scratch;
10035        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10036        // (requested split, destination list). Single-shot per burst; fresh calls have none.
10037        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10038        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10039        // committed-length position; consumed one-shot like `capture_at`. None = legacy
10040        // prompt-end capture below.
10041        let mut ckpt_req: Option<usize> = None;
10042        let (
10043            cache,
10044            scratch,
10045            mut sess_tail,
10046            mut sess_draft_slot,
10047            mut sess_pending_slot,
10048            sess_ckpt_slot,
10049            sess_telem,
10050        ): (
10051            &mut Cache,
10052            &mut MtpScratch,
10053            Option<(
10054                &mut Vec<u32>,
10055                &mut Option<CudaSlice<f32>>,
10056                &mut Option<u32>,
10057                &mut u32,
10058                &mut u32,
10059            )>,
10060            Option<&mut Option<DraftGraphCtx>>,
10061            Option<&mut Option<u32>>,
10062            Option<&mut Option<SpecCheckpoint>>,
10063            Option<&SpecTelemetryCounters>,
10064        ) = match sess.take() {
10065            Some(sr) => {
10066                let SpecSession {
10067                    cache,
10068                    scratch,
10069                    committed,
10070                    last_h,
10071                    next_pred,
10072                    sctr: s_sctr,
10073                    uctr: s_uctr,
10074                    draft_ctx,
10075                    pending_tok,
10076                    turn_ckpt,
10077                    telem,
10078                    capture_at,
10079                    boundary_captures,
10080                    ckpt_at,
10081                } = sr;
10082                sess_capture = Some((capture_at.take(), boundary_captures));
10083                ckpt_req = ckpt_at.take();
10084                (
10085                    cache,
10086                    scratch,
10087                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
10088                    Some(draft_ctx),
10089                    Some(pending_tok),
10090                    Some(turn_ckpt),
10091                    Some(telem),
10092                )
10093            }
10094            None => {
10095                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
10096                // `Cache::new` verbatim.
10097                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
10098                // Persistent scratch = max_ctx rows (~2KB/token quantized).
10099                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
10100                (
10101                    &mut own_cache,
10102                    &mut own_scratch,
10103                    None,
10104                    None,
10105                    None,
10106                    None,
10107                    None,
10108                )
10109            }
10110        };
10111        if scratch.plane_count() != self.mtp_head_count() {
10112            return Err(format!(
10113                "MTP scratch/head count mismatch ({}/{})",
10114                scratch.plane_count(),
10115                self.mtp_head_count()
10116            )
10117            .into());
10118        }
10119        let base = cache.pos;
10120        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
10121        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
10122        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
10123        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
10124        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
10125        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
10126        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
10127        // acceptance-only — exactness is verify's job either way).
10128        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
10129        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
10130        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
10131        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
10132        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
10133        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
10134        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
10135        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
10136        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
10137        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
10138        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
10139        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
10140        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
10141        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
10142        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
10143        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
10144        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
10145        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
10146        // + fallback seam).
10147        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
10148        // bar — the retained verify-state commit proven equivalent to sequential serving —
10149        // was waiting on this arch running the serving batched verify class, which the
10150        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
10151        // replay-free commit consumes is now produced by the SAME serving-class verify that
10152        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
10153        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
10154        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
10155        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
10156        // rollback + A/B seam.
10157        let spec_replay = spec_replay_env_enabled();
10158        if constraint.is_some() && spec_replay {
10159            return Err(
10160                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
10161                        (legacy replay commits an unmasked bonus)"
10162                    .into(),
10163            );
10164        }
10165        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
10166        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
10167        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
10168        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
10169        if !refresh && !self.mtp_extra.is_empty() {
10170            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
10171        }
10172
10173        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
10174        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
10175        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
10176        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
10177        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
10178        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
10179        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
10180        // generation exactly where the last turn stopped — no prime at all. The stashed
10181        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
10182        // committed.last() by the same rule this entry applies to a cold prime's last row —
10183        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
10184        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
10185        // where the sampler and the session's Philox counters were live). `last_h` seeds the
10186        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
10187        let continuation = prompt.is_empty();
10188        if continuation {
10189            assert!(session_mode, "empty prompt requires a session");
10190            assert!(
10191                sess_tail
10192                    .as_ref()
10193                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
10194                        && lh.is_some()
10195                        && (np.is_some() || carried_pending.is_some())),
10196                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
10197            );
10198        }
10199        let mut prime_logits;
10200        let mut prompt_h: Option<CudaSlice<f32>> = None;
10201        let t_prime = std::time::Instant::now();
10202        let batched_prime = !continuation
10203            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
10204            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10205            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
10206        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
10207        if prime_split.is_some() && continuation {
10208            return Err("spec prime split requires a non-empty prime".into());
10209        }
10210        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
10211        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
10212        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
10213        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
10214        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
10215        // cannot honor (outside this prime's range) silently drops the capture — the
10216        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
10217        let ckpt_rel = if continuation {
10218            None
10219        } else {
10220            ckpt_req
10221                .and_then(|abs| abs.checked_sub(base))
10222                .filter(|&r| r > 0 && r < prompt.len())
10223        };
10224        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
10225        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
10226        // the legacy single-split program, byte-for-byte.
10227        let mut stops: Vec<usize> = Vec::new();
10228        for b in [prime_split, ckpt_rel].into_iter().flatten() {
10229            if !stops.contains(&b) {
10230                stops.push(b);
10231            }
10232        }
10233        stops.sort_unstable();
10234        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
10235        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
10236        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
10237        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
10238        if continuation {
10239            prime_logits = Vec::new();
10240        } else if !stops.is_empty() {
10241            if let Some(&first) = stops.first() {
10242                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
10243                    return Err(format!(
10244                        "spec prime split {first} is below PRIME_MIN_T {}",
10245                        crate::hybrid_forward::PRIME_MIN_T,
10246                    )
10247                    .into());
10248                }
10249            }
10250            // Mirror the plain worker's boundary stops exactly. Each segment is a
10251            // request-level prime (`queued_after` keeps Step35 arm selection independent of
10252            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
10253            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
10254            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
10255            // coherent prompt.
10256            let mut h_all = e.uninit(prompt.len() * n_embd)?;
10257            prime_logits = Vec::new();
10258            let mut prev = 0usize;
10259            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
10260                if seg_end <= prev {
10261                    continue;
10262                }
10263                let seg = &prompt[prev..seg_end];
10264                let is_final = seg_end == prompt.len();
10265                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
10266                    && (!is_final
10267                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
10268                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
10269                if batched_seg {
10270                    let (l, _, h_seg) =
10271                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
10272                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
10273                    prime_logits = l;
10274                } else {
10275                    for (i, &tok) in seg.iter().enumerate() {
10276                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
10277                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
10278                        prime_logits = l;
10279                    }
10280                }
10281                prev = seg_end;
10282                if is_final {
10283                    break;
10284                }
10285                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
10286                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
10287                // states are about to be advanced in place by the next segment, so this is
10288                // the ONLY moment the boundary's recurrent state exists. Capture iff the
10289                // worker requested exactly this stop (cold sessions only — `capture_at` is
10290                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
10291                // publication is an optimization, never a correctness dependency.
10292                if base == 0 {
10293                    if let Some((requested, slot)) = sess_capture.as_mut() {
10294                        // Publish at the requested miss-LCP stop (the shared-prefix class)
10295                        // AND at the stable-boundary stop (the next-turn re-render class,
10296                        // lane/frspec-multiturn-cache) — the same boundary set the plain
10297                        // prefill tick learns. Without the second entry, the turn after a
10298                        // cold re-park could only hit the OLDER lcp entry (the measured
10299                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
10300                        // rewound to 15222). Dedupe is the worker sweep's has_key.
10301                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
10302                            if let Ok(snap) = cache.snapshot(e) {
10303                                slot.push(SpecBoundaryCapture {
10304                                    snap,
10305                                    pos: seg_end,
10306                                    logits: prime_logits.clone(),
10307                                    // rows [0..seg_end) of h_all are primed — the following
10308                                    // segments append, never overwrite.
10309                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
10310                                });
10311                            }
10312                        }
10313                    }
10314                }
10315                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
10316                // same snapshot mechanics, installed post-prime in place of the prompt-end
10317                // capture the re-render class always diverged below.
10318                if ckpt_rel == Some(seg_end) {
10319                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10320                        e.uninit(n_embd).and_then(|mut a| {
10321                            e.copy_view_into(
10322                                &mut a,
10323                                0,
10324                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
10325                                n_embd,
10326                            )?;
10327                            Ok(a)
10328                        });
10329                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
10330                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
10331                            snap,
10332                            pos: base + seg_end,
10333                            last_h,
10334                        }),
10335                        _ => None,
10336                    });
10337                }
10338            }
10339            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
10340                eprintln!(
10341                    "[spec-prime] stops={stops:?} tail={}",
10342                    prompt.len() - stops.last().copied().unwrap_or(0)
10343                );
10344            }
10345            prompt_h = Some(h_all);
10346        } else if batched_prime {
10347            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
10348            prime_logits = l;
10349            prompt_h = Some(hiddens);
10350        } else {
10351            prime_logits = Vec::new();
10352            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
10353            for (i, &tok) in prompt.iter().enumerate() {
10354                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
10355                if let Some(ph) = prompt_h.as_mut() {
10356                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
10357                }
10358                prime_logits = l;
10359            }
10360        }
10361        e.stream().synchronize()?;
10362        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
10363        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
10364        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
10365        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
10366        // prime_split. The mid-prompt capture above already consumed the request if it matched.
10367        if !continuation && base == 0 {
10368            if let Some((requested, slot)) = sess_capture.as_mut() {
10369                if *requested == Some(prompt.len()) && slot.is_empty() {
10370                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
10371                    if let Ok(snap) = cache.snapshot(e) {
10372                        slot.push(SpecBoundaryCapture {
10373                            snap,
10374                            pos: prompt.len(),
10375                            logits: prime_logits.clone(),
10376                            last_h: prompt_h
10377                                .as_ref()
10378                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
10379                                .unwrap_or_default(),
10380                        });
10381                    }
10382                }
10383            }
10384        }
10385        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
10386        // prime-subtraction hack.
10387        crate::PRIME_NANOS.store(
10388            t_prime.elapsed().as_nanos() as u64,
10389            std::sync::atomic::Ordering::Relaxed,
10390        );
10391
10392        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10393        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
10394        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
10395        let host_embd = spec_host_embd();
10396        let embd_gpu = if host_embd {
10397            None
10398        } else {
10399            Some(
10400                self.embd_gpu
10401                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10402            )
10403        };
10404        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10405        if host_embd {
10406            eprintln!(
10407                "[spec] host-row embedding: {} bytes kept off HBM",
10408                self.embd.raw.len()
10409            );
10410        }
10411        let mut out: Vec<u32> = Vec::with_capacity(max_new);
10412        let mut total_drafted = 0usize;
10413        let mut total_accepted = 0usize;
10414
10415        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
10416        // The sampler config, the session's Philox counters and the penalty window are parsed
10417        // HERE, above the boundary-token selection, because the boundary token must be drawn
10418        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
10419        // selection, which is the whole mechanical reason the boundary token was an argmax:
10420        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
10421        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
10422        // below takes the argmax path it always took).
10423        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
10424        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
10425        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
10426        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
10427        let sp = sampling.unwrap_or_else(|| SpecSampling {
10428            temp: std::env::var("MEMRA_SPEC_TEMP")
10429                .ok()
10430                .and_then(|v| v.parse().ok())
10431                .unwrap_or(0.0),
10432            seed: std::env::var("MEMRA_SEED")
10433                .ok()
10434                .and_then(|v| v.parse().ok())
10435                .unwrap_or(42),
10436            top_k: std::env::var("MEMRA_TOP_K")
10437                .ok()
10438                .and_then(|v| v.parse().ok())
10439                .unwrap_or(0),
10440            top_p: std::env::var("MEMRA_TOP_P")
10441                .ok()
10442                .and_then(|v| v.parse().ok())
10443                .unwrap_or(1.0),
10444            min_p: std::env::var("MEMRA_MIN_P")
10445                .ok()
10446                .and_then(|v| v.parse().ok())
10447                .unwrap_or(0.0),
10448            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
10449                .ok()
10450                .and_then(|v| v.parse().ok())
10451                .unwrap_or(0),
10452            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
10453                .ok()
10454                .and_then(|v| v.parse().ok())
10455                .unwrap_or(1.0),
10456            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
10457                .ok()
10458                .and_then(|v| v.parse().ok())
10459                .unwrap_or(0.0),
10460            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
10461                .ok()
10462                .and_then(|v| v.parse().ok())
10463                .unwrap_or(0.0),
10464        });
10465        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
10466        let sampled = sp_temp > 0.0;
10467        // Counters resume from the session (burst continuity: randomness must never repeat
10468        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
10469        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
10470        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
10471        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
10472        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
10473        // for the penalized+filtered target). History = generated tokens, host-tracked window.
10474        let pen_on = sampled
10475            && sp.penalty_last_n > 0
10476            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
10477        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
10478        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
10479        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
10480        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
10481        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
10482        // which is what the API contract says and what the plain sampler's own `history` does.
10483        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
10484        let mut pen_hist: Vec<u32> = if pen_on {
10485            let sess_hist: &[u32] = if spec_pen_session_on() {
10486                sess_tail
10487                    .as_ref()
10488                    .map(|(c, ..)| c.as_slice())
10489                    .unwrap_or(&[])
10490            } else {
10491                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
10492            };
10493            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
10494        } else {
10495            Vec::new()
10496        };
10497        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
10498        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
10499        // request's own filtered/penalized target through the session's Philox stream
10500        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
10501        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
10502        // Emit it, then FEED it to establish the loop invariant below.
10503        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
10504        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
10505        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
10506        // prompt's last logits (plain constrained-greedy identity); a continuation without
10507        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
10508        // worker never resumes constrained sessions from the pool, so this cannot fire).
10509        if let Some(c) = constraint.as_deref_mut() {
10510            if continuation && carried_pending.is_none() {
10511                return Err("constrained spec continuation requires a carried pending \
10512                            (pool resume is unconstrained-only)"
10513                    .into());
10514            }
10515            if !continuation {
10516                c.mask_logits(&mut prime_logits)
10517                    .map_err(|e2| format!("constraint: {e2}"))?;
10518            }
10519        }
10520        let mut last_token = if let Some(b) = carried_pending {
10521            b
10522        } else if continuation {
10523            // A continuation's boundary token was DRAWN by the burst that stashed it (the
10524            // session tail below), or by `spec_session_from_restored` for a converted
10525            // prefix-cache hit — in both cases from the correct logits row with this same
10526            // session's Philox stream, which is why it can be consumed here as-is.
10527            sess_tail.as_ref().unwrap().2.unwrap()
10528        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
10529            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
10530        } else {
10531            // greedy (byte contract), the rollback door, or constrained (masked-argmax
10532            // identity — the worker routes sampled+constrained to the plain path, and this
10533            // function refuses the combination outright above).
10534            argmax(&prime_logits) as u32
10535        };
10536        if pen_on {
10537            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
10538            // emitted token into its penalty history, and pre-lane the burst's first token
10539            // was invisible to penalties forever (never pushed, and never in `committed`
10540            // until this burst's tail). Covers the carry/continuation seeds too — neither is
10541            // in `committed` yet.
10542            pen_hist.push(last_token);
10543        }
10544        if carried_pending.is_none() {
10545            out.push(last_token);
10546            // grammar advances with every emitted token (carried pendings were consumed
10547            // by the burst that emitted them).
10548            if let Some(c) = constraint.as_deref_mut() {
10549                c.consume(last_token)
10550                    .map_err(|e2| format!("constraint: {e2}"))?;
10551            }
10552        }
10553        if continuation {
10554            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
10555            // overhang so the chain's first append lands at slot base (== committed.len()).
10556            scratch.set_len(e, base)?;
10557        }
10558        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
10559        // concatenating to the full `out`). Called after the prime's first token and after each
10560        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
10561        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
10562        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
10563        fn flush_commit(
10564            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
10565            out: &[u32],
10566            flushed: &mut usize,
10567        ) -> bool {
10568            if let Some(f) = cb.as_mut() {
10569                let keep = f(&out[*flushed..]);
10570                *flushed = out.len();
10571                keep
10572            } else {
10573                true
10574            }
10575        }
10576        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
10577        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
10578        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
10579        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
10580        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
10581        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
10582        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
10583        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
10584        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
10585        // those, so their residual mass is p(x), correct by construction).
10586        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
10587            match &mtp.d2t {
10588                Some(map) => Some(e.htod_u32_v(map)?),
10589                None => None,
10590            }
10591        } else {
10592            None
10593        };
10594        let mut q_full_buf: Option<CudaSlice<f32>> = None;
10595        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
10596        // dspark sampled-admission walk); byte-identical to the closure it replaces.
10597        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
10598        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
10599        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
10600        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
10601        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
10602        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
10603        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
10604        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
10605        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
10606        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
10607        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
10608        let t_ent = std::time::Instant::now();
10609
10610        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
10611        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
10612        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
10613        // the one that matters (a history-rewriting client mutates what the session GENERATED,
10614        // so the next turn's prompt agrees with this one up to exactly here).
10615        //
10616        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
10617        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
10618        // hold exactly `base + prompt.len()` rows and nothing generated.
10619        //
10620        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
10621        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
10622        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
10623        // `<think>` block the client strips, so every later turn's diff diverged exactly one
10624        // token below the checkpoint and affinity declined 100% of the time. Measured on the
10625        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
10626        // whole mechanism inert while looking, from the outside, like a working
10627        // correctness-declines-safely path — hence the decline log carries the offsets.
10628        //
10629        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
10630        // state (the reason a spec session could not rewind before). The draft scratch needs no
10631        // copy: rows below the boundary are rewritten by the next turn's own fill.
10632        //
10633        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
10634        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
10635        // checkpoint rather than replacing it with a strictly worse one.
10636        //
10637        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
10638        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
10639        // fail the burst that is already running — so the error is swallowed, loud only under
10640        // MEMRA_DEBUG_SPEC.
10641        //
10642        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
10643        // posture above was DISPROVED for the think-posture template class — the prompt's own
10644        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
10645        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
10646        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
10647        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
10648        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
10649        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
10650        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
10651        if let Some(slot) = sess_ckpt_slot {
10652            if let Some(early) = ckpt_early {
10653                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10654                    eprintln!(
10655                        "[spec] stable-boundary turn checkpoint skipped; \
10656                               next turn re-primes in full"
10657                    );
10658                }
10659                *slot = early;
10660            } else if !continuation {
10661                let pos = cache.pos;
10662                debug_assert_eq!(
10663                    pos,
10664                    base + prompt.len(),
10665                    "turn checkpoint must sit at the prompt end, before the init feed"
10666                );
10667                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
10668                    if let Some(ph) = &prompt_h {
10669                        // hidden of the LAST primed row = the predecessor anchor at this
10670                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
10671                        // last_h, and what the next prime's fill reads for its first row).
10672                        let np = prompt.len();
10673                        e.uninit(n_embd).and_then(|mut a| {
10674                            e.copy_view_into(
10675                                &mut a,
10676                                0,
10677                                &ph.slice((np - 1) * n_embd..np * n_embd),
10678                                n_embd,
10679                            )?;
10680                            Ok(a)
10681                        })
10682                    } else {
10683                        Err("no prompt hiddens".into())
10684                    };
10685                match (cache.snapshot(e), anchor) {
10686                    (Ok(snap), Ok(last_h)) => {
10687                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
10688                    }
10689                    (s, a) => {
10690                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
10691                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
10692                            let err = s
10693                                .err()
10694                                .map(|e| e.to_string())
10695                                .or_else(|| a.err().map(|e| e.to_string()))
10696                                .unwrap_or_default();
10697                            eprintln!(
10698                                "[spec] turn checkpoint skipped ({err}); \
10699                                       next turn re-primes in full"
10700                            );
10701                        }
10702                    }
10703                }
10704            }
10705        }
10706        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
10707        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
10708        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
10709        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
10710        let mut last_pred = 0u32;
10711        let mut last_col_logits: Option<CudaSlice<f32>> = None;
10712        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
10713        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
10714        let mut init_logits_host: Option<Vec<f32>> = None;
10715        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
10716            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
10717            last_pred = argmax(&init_logits) as u32;
10718            if constraint.is_some() {
10719                init_logits_host = Some(init_logits.clone());
10720            }
10721            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
10722            if sampled {
10723                last_col_logits = Some(e.htod(&init_logits)?);
10724            }
10725            h
10726        } else {
10727            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
10728            let lh = sess_tail
10729                .as_ref()
10730                .unwrap()
10731                .1
10732                .as_ref()
10733                .expect("pending carry requires last_h");
10734            e.clone_dtod(lh)?
10735        };
10736        let t_init = t_ent.elapsed();
10737        let mut last_col_stats: Option<(f32, f32, f32)> = None;
10738        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
10739        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
10740        // stable pointer for the graph-draft round-start copy.
10741        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
10742        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
10743        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
10744        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
10745        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
10746        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
10747        // overwritten below).
10748        let mut fill_prev = e.clone_dtod(&h_seed0)?;
10749        {
10750            if let Some(ph) = &prompt_h {
10751                let np = prompt.len();
10752                e.copy_view_into(
10753                    &mut h_seed_buf,
10754                    0,
10755                    &ph.slice((np - 1) * n_embd..np * n_embd),
10756                    n_embd,
10757                )?;
10758            } else if continuation {
10759                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
10760                    if let Some(lh) = lh.as_ref() {
10761                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
10762                    }
10763                }
10764            }
10765        }
10766        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
10767        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
10768
10769        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
10770        let fork_mode = OptiForkGateMode::configured();
10771        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
10772        // the end. Metric normalization vs the reference engine: BOTH engines count
10773        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
10774        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
10775        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
10776        let mut st_drafted = vec![0usize; k];
10777        let mut st_accepted = vec![0usize; k];
10778        let mut st_len_hist = vec![0usize; k + 1];
10779        let mut st_full = 0usize;
10780        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
10781        // stop the draft chain early when the head's softmax confidence in its own pick drops
10782        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
10783        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
10784        let p_min = *PMIN.get_or_init(|| {
10785            std::env::var("MEMRA_SPEC_PMIN")
10786                .ok()
10787                .and_then(|v| v.parse().ok())
10788                .unwrap_or(0.0)
10789        });
10790        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
10791        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
10792        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
10793        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
10794        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
10795        // verify batch is not); the j==0 exemption stays for pending-less rounds.
10796        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
10797            .map(|v| v == "1")
10798            .unwrap_or(false);
10799
10800        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
10801        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
10802        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
10803        // cuBLAS path in an exotic head) falls back to the eager draft chain.
10804        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
10805        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
10806        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
10807        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
10808        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
10809        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
10810        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
10811        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
10812        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
10813            Some(c) => c,
10814            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
10815        };
10816        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
10817        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
10818        if sampled && dctx.g_q.len() < d_vocab {
10819            dctx.g_q = e.zeros(d_vocab)?;
10820            dctx.g_perturb = e.zeros(d_vocab)?;
10821        }
10822        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
10823        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
10824        // truncation (the correctness backstop) stops cutting every tight-schema round.
10825        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
10826        // shape, so a parked graph of the other shape is dropped and recaptured.
10827        let dmask_on = constraint
10828            .as_deref()
10829            .is_some_and(|c| c.draft_mask_enabled());
10830        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
10831        if dmask_on && dctx.g_dmask.len() < dmask_words {
10832            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
10833            dctx.graph = None; // the old capture baked the old (or no) mask pointer
10834            dctx.failed.clear_greedy();
10835            dctx.keeper.clear();
10836        }
10837        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
10838            dctx.graph = None;
10839            dctx.failed.clear_greedy();
10840            dctx.keeper.clear();
10841        }
10842        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
10843            let DraftGraphCtx {
10844                g_tok,
10845                g_pos,
10846                g_seed,
10847                g_p,
10848                g_dmask,
10849                ..
10850            } = &mut dctx;
10851            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
10852            // host uploads the position's real words, so the warmups stay grammar-free.
10853            if dmask_on {
10854                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
10855            }
10856            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
10857            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
10858            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
10859            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
10860            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
10861            // passes (and, in serve, other sessions) recycle those addresses and the replay then
10862            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
10863            let cap_res = e.capture_graph_retained(|e| {
10864                self.mtp_head_forward_cap(
10865                    e,
10866                    mtp,
10867                    g_tok,
10868                    g_pos,
10869                    g_seed,
10870                    g_p,
10871                    &mut *scratch,
10872                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
10873                    true,
10874                    embd_gpu.expect("graph draft requires resident embedding"),
10875                    embd_qt,
10876                    embd_rb,
10877                    d_vocab,
10878                    None,
10879                    None,
10880                    if dmask_on {
10881                        Some((g_dmask_ro, dmask_words))
10882                    } else {
10883                        None
10884                    },
10885                )
10886            });
10887            match cap_res {
10888                Ok((g, keep)) => {
10889                    scratch.set_len(e, base)?;
10890                    dctx.graph = Some(g);
10891                    dctx.graph_masked = dmask_on;
10892                    dctx.keeper = keep;
10893                }
10894                Err(err) => {
10895                    scratch.set_len(e, base)?;
10896                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
10897                    // silent. Once per flip — mark returns None on an already-failed ctx.
10898                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
10899                        eprintln!("{line}");
10900                    }
10901                }
10902            }
10903        }
10904        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
10905        // graph object, built only when sampled && graph-eligible — the greedy capture above is
10906        // untouched (and skipped when sampled: its graph would never be launched). Same head
10907        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
10908        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
10909        // once per round); the raw head logits land in the persistent g_q for the host's
10910        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
10911        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
10912        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
10913        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
10914        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
10915        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
10916        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
10917        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
10918        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
10919        // this compare misses at most ONCE per resumed request — the first burst recaptures
10920        // and every later burst in that request replays. A client that wants the parked graph
10921        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
10922        // stable across its whole conversation.
10923        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
10924        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
10925        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
10926        // force the eager draft (which computes stats/penalties per row).
10927        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
10928        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
10929        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
10930        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
10931        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
10932        // the request shape the vendor-default flip makes the majority).
10933        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
10934        let pure_temp = s_key.pure_temp();
10935        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
10936            dctx.graph_s = None;
10937            dctx.failed.clear_sampled();
10938            dctx.s_key = None;
10939            dctx.q_slots.clear();
10940            dctx.keeper_s.clear();
10941        }
10942        if graph_draft
10943            && sampled
10944            && pure_temp
10945            && dctx.graph_s.is_none()
10946            && !dctx.failed.sampled_failed()
10947        {
10948            let DraftGraphCtx {
10949                g_tok,
10950                g_pos,
10951                g_seed,
10952                g_p,
10953                g_ctr,
10954                g_perturb,
10955                g_q,
10956                ..
10957            } = &mut dctx;
10958            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
10959            let cap_res = e.capture_graph_retained(|e| {
10960                self.mtp_head_forward_cap(
10961                    e,
10962                    mtp,
10963                    g_tok,
10964                    g_pos,
10965                    g_seed,
10966                    g_p,
10967                    &mut *scratch,
10968                    p_min > 0.0,
10969                    true,
10970                    embd_gpu.expect("graph draft requires resident embedding"),
10971                    embd_qt,
10972                    embd_rb,
10973                    d_vocab,
10974                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
10975                    None,
10976                    None, // constrained spec is greedy-only — sampled never carries a hook
10977                )
10978            });
10979            match cap_res {
10980                Ok((g, keep)) => {
10981                    scratch.set_len(e, base)?;
10982                    for _ in 0..k {
10983                        dctx.q_slots.push(e.zeros(d_vocab)?);
10984                    }
10985                    dctx.graph_s = Some(g);
10986                    dctx.s_key = Some(s_key);
10987                    dctx.keeper_s = keep;
10988                }
10989                Err(err) => {
10990                    scratch.set_len(e, base)?;
10991                    // LOUD flip (audit Q2): same contract as the greedy capture above.
10992                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
10993                        eprintln!("{line}");
10994                    }
10995                }
10996            }
10997        }
10998        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
10999        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11000        // captured under this request's exact regime, and capture requires `pure_temp` — so a
11001        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
11002        // the graph arm, so it is asserted here rather than assumed: a future change that widens
11003        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11004        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
11005        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
11006        // rather than launching it; the launch site re-tests `pure_temp` independently.
11007        if sampled && !pure_temp && dctx.graph_s.is_some() {
11008            debug_assert!(
11009                false,
11010                "sampled draft graph parked under {:?} survived into a FILTERED request \
11011                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
11012                 softmax, so the verify's filtered q would test a distribution the draft was \
11013                 never sampled from",
11014                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11015            );
11016            eprintln!(
11017                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11018                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
11019                 EAGER — the key must carry every field that shapes q",
11020                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11021            );
11022            dctx.graph_s = None;
11023            dctx.s_key = None;
11024            dctx.q_slots.clear();
11025            dctx.keeper_s.clear();
11026        }
11027        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11028        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
11029        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
11030        // arms below print which chain actually ran, so the probe never restates the condition.
11031        if skey_probe() {
11032            eprintln!(
11033                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
11034                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
11035                sampled as u8,
11036                pure_temp as u8,
11037                sp_temp,
11038                sp.top_k,
11039                sp.top_p,
11040                sp.min_p,
11041                pen_on as u8,
11042                k,
11043                graph_draft as u8,
11044                dctx.graph_s.is_some() as u8,
11045                dctx.s_key,
11046            );
11047        }
11048        let t_cap = t_ent.elapsed();
11049        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
11050        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
11051        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
11052        // fill: the first chain step processes it and appends its entry at slot prompt.len().
11053        if let Some(ph) = &prompt_h {
11054            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
11055            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
11056            // global positions [base..base+tp). Fresh call: base==0, identical to before.
11057            scratch.set_len(e, base)?;
11058            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
11059            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
11060            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
11061            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
11062            let tp = prompt.len();
11063            let fill_chunk: usize = if crate::cache::swa_ring_on() {
11064                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
11065            } else {
11066                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
11067                // meaning one monolithic fill.
11068                std::env::var("MEMRA_PRIME_CHUNK")
11069                    .ok()
11070                    .and_then(|v| v.parse().ok())
11071                    .unwrap_or(4096)
11072            };
11073            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
11074            let mut start = 0usize;
11075            while start < tp {
11076                let end = (start + fill_chunk).min(tp);
11077                let tc = end - start;
11078                {
11079                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
11080                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
11081                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
11082                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
11083                    let mut phs = e.zeros(tc * n_embd)?;
11084                    let (src_lo, dst_off) = if start == 0 {
11085                        (0, n_embd)
11086                    } else {
11087                        ((start - 1) * n_embd, 0)
11088                    };
11089                    let n_copy = if start == 0 {
11090                        (tc - 1) * n_embd
11091                    } else {
11092                        tc * n_embd
11093                    };
11094                    if start == 0 {
11095                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11096                            if let Some(lh) = lh.as_ref() {
11097                                e.copy_into(&mut phs, 0, lh, n_embd)?;
11098                            }
11099                        }
11100                    }
11101                    if n_copy > 0 {
11102                        e.copy_view_into(
11103                            &mut phs,
11104                            dst_off,
11105                            &ph.slice(src_lo..src_lo + n_copy),
11106                            n_copy,
11107                        )?;
11108                    }
11109                    self.mtp_kv_fill_all(
11110                        e,
11111                        &prompt[start..end],
11112                        &phs,
11113                        base + start,
11114                        &mut *scratch,
11115                        embd_dev,
11116                    )?;
11117                }
11118                start = end;
11119            }
11120        }
11121        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
11122        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
11123        // (=1 brackets the whole call in run_spec.rs, prime included.)
11124        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
11125            unsafe extern "C" {
11126                fn cudaProfilerStart() -> i32;
11127            }
11128            unsafe {
11129                cudaProfilerStart();
11130            }
11131        }
11132        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
11133        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
11134        // consume each other's device outputs; the host drains the ring every M rounds. v1
11135        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
11136        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
11137        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
11138        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
11139        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
11140        let stream_on = crate::spec::spec_stream()
11141            && !sampled
11142            && !spec_replay
11143            && self.mtp_extra.is_empty()
11144            && constraint.is_none()
11145            && !session_mode
11146            && embd_gpu.is_some()
11147            && !crate::model::full_prec_enabled()
11148            && k + 2 < 96;
11149        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
11150        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
11151        if stream_on {
11152            let cap = e.capture_graph(|e| {
11153                for j in 0..k.max(1) {
11154                    self.mtp_head_forward_cap(
11155                        e,
11156                        mtp,
11157                        &mut dctx.g_tok,
11158                        &mut dctx.g_pos,
11159                        &mut dctx.g_seed,
11160                        &mut dctx.g_p,
11161                        &mut *scratch,
11162                        true,
11163                        true,
11164                        embd_gpu.expect("round stream requires resident embedding"),
11165                        embd_qt,
11166                        embd_rb,
11167                        d_vocab,
11168                        None,
11169                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
11170                        None, // round-stream requires constraint.is_none() (see stream_on)
11171                    )?;
11172                }
11173                Ok(())
11174            });
11175            match cap {
11176                Ok(g) => {
11177                    scratch.set_len(e, 0)?;
11178                    stream_graph = Some(g);
11179                }
11180                Err(err) => {
11181                    scratch.set_len(e, 0)?;
11182                    if debug_spec {
11183                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
11184                    }
11185                }
11186            }
11187        }
11188        let stream_active = stream_on && stream_graph.is_some();
11189        if debug_spec {
11190            eprintln!(
11191                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
11192                crate::spec::spec_stream(),
11193                dctx.graph.is_some(),
11194                stream_graph.is_some()
11195            );
11196        }
11197        let t_v_s = k + 1;
11198        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
11199        // module (extracted 2026-07-12; the gemma burst reuses them).
11200        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
11201        let crate::round_stream::StreamBufs {
11202            mut vtok_d,
11203            mut brk_d,
11204            mut pend_d,
11205            last_pred_d,
11206            mut pos_ctr,
11207            mut pos_start_d,
11208            mut ring_d,
11209            acc_d: mut stream_acc,
11210            m_rounds,
11211            k: _,
11212        } = sb;
11213        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
11214            Some(crate::round_stream::kv_len_ptr_table(
11215                e,
11216                cache,
11217                Some(&pos_ctr),
11218            )?)
11219        } else {
11220            None
11221        };
11222
11223        let t_fill = t_ent.elapsed();
11224        let mut round = 0usize;
11225        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
11226        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
11227        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
11228        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
11229        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
11230        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
11231        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
11232        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
11233        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
11234        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
11235        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
11236        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
11237        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
11238        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
11239        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
11240        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
11241        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
11242        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
11243        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
11244        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
11245        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
11246        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
11247        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
11248        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
11249        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
11250        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
11251        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
11252        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
11253        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
11254        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
11255            .ok()
11256            .and_then(|v| v.parse().ok());
11257        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
11258            4
11259        } else if self.cfg.n_embd as usize >= 2500 {
11260            2
11261        } else {
11262            1
11263        };
11264        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
11265        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
11266        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
11267        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
11268        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
11269            .ok()
11270            .and_then(|v| v.parse().ok())
11271            .unwrap_or(1024);
11272        let floor_at = |pos: usize| -> usize {
11273            if adapt_floor_env.is_some() || pos < floor_ctx {
11274                adapt_floor
11275            } else if adapt_floor >= 4 {
11276                1
11277            } else {
11278                adapt_floor
11279            }
11280        };
11281        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
11282        // fixed-K default path is untouched by this whole block.
11283        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
11284            .ok()
11285            .and_then(|v| v.parse().ok())
11286            .unwrap_or(7);
11287        let k_cap = k.min(cap_max).max(1);
11288        let mut kc = k_cap;
11289        let mut opti_fork: Option<OptiForkState> = None;
11290        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
11291        if fork_mode != OptiForkGateMode::Disabled {
11292            let fence = crate::pp::pp_cuts(self.layers.len());
11293            let refusal = if !session_mode {
11294                Some("not-session")
11295            } else if k != 1 || adapt {
11296                Some("requires-fixed-k1")
11297            } else if sampled || constraint.is_some() || spec_replay {
11298                Some("sampled-constrained-or-replay")
11299            } else if pipe.is_some() {
11300                Some("two-session-pipeline")
11301            } else if !spec_devacc() {
11302                Some("requires-device-accept")
11303            } else if stream_active || crate::spec::spec_stream() {
11304                Some("round-stream")
11305            } else if !self.mtp_extra.is_empty() {
11306                Some("multi-head-mtp")
11307            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
11308                Some("swa-ring")
11309            } else if crate::pp::pp_host_bounce_active() {
11310                Some("host-bounce")
11311            } else if fork_mode == OptiForkGateMode::Controller
11312                && cache.recur.iter().any(Option::is_some)
11313            {
11314                Some("controller-requires-zero-recurrent-state")
11315            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
11316                Some("requires-pp2")
11317            } else {
11318                None
11319            };
11320            if let Some(reason) = refusal {
11321                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11322                eprintln!("[opti-fork] refused reason={reason}");
11323            } else {
11324                let fence = fence.expect("validated PP-2 fence");
11325                let rt = crate::pp::PpNRt::get(e)?;
11326                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
11327                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
11328                let primary_supported =
11329                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
11330                if !rt.cross_device() || !primary_supported {
11331                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
11332                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
11333                } else {
11334                    // Both recurrent snapshots and both seed generations are allocated before
11335                    // the first fork, each through its owning PP stage. Allocation failure
11336                    // therefore happens before any optimistic state mutation can occur.
11337                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
11338                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
11339                    let fork = OptiForkState::new(
11340                        e,
11341                        cache,
11342                        fork_mode,
11343                        alternate_snapshot,
11344                        &h_seed_buf,
11345                        &fill_prev,
11346                        rt,
11347                        fence[1],
11348                        self.layers.len(),
11349                    )?;
11350                    eprintln!(
11351                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
11352                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
11353                        fence[1],
11354                        fork.logical_payload_bytes[0],
11355                        fork.logical_payload_bytes[1],
11356                        fork.controller.map_or(0.0, |policy| policy.threshold),
11357                    );
11358                    fork_snapshot = Some(current_snapshot);
11359                    opti_fork = Some(fork);
11360                }
11361            }
11362        }
11363        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
11364        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
11365        let mut snap = match fork_snapshot {
11366            Some(snapshot) => snapshot,
11367            None => cache.snapshot(e)?,
11368        };
11369        let mut carried_opti: Option<OptiControllerTicket> = None;
11370        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
11371        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
11372        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
11373            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
11374        } else {
11375            None
11376        };
11377        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
11378        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
11379        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
11380        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
11381        // pass of any kind). Verify still
11382        // checks every emitted token against the target -> exactness holds by construction; only
11383        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
11384        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
11385        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
11386        let mut pending: Option<u32> = carried_pending;
11387        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
11388        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
11389        // the verify accept readback). Printed once at loop end via spec-stats.
11390        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
11391        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
11392        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
11393        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
11394        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
11395        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
11396        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
11397        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
11398        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
11399        let mut ph_wait = 0f64;
11400        let mut ph_commit = 0f64;
11401        let mut ph_t = std::time::Instant::now();
11402        let mut ph_mark = |acc: &mut f64, on: bool| {
11403            if on {
11404                let now = std::time::Instant::now();
11405                *acc += (now - ph_t).as_secs_f64();
11406                ph_t = now;
11407            }
11408        };
11409        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
11410        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
11411        // arm holds it — the slab stash is live verify -> commit inside a round, and the
11412        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
11413        // the model (rebuilding per call re-captures the pool per prompt, which is the
11414        // measured way to lose more than the launches cost); the captured bodies are
11415        // cache-independent, every state read going through per-round refreshed pointer
11416        // tables. None = the eager walk, byte-identical.
11417        //
11418        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
11419        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
11420        // whenever the stream is live rather than relying on that refusal.
11421        // The lock is taken ONLY when the door is armed: with the flag off this whole block
11422        // is inert, so the default path cannot serialize two spec generations behind a mutex
11423        // it never reads.
11424        let vg_armed =
11425            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
11426        let mut vg_guard = if vg_armed && !stream_active {
11427            let mut g = self.dspark_vgraphs.lock().unwrap();
11428            if g.is_none() {
11429                // Size by the WIDEST verify this run can present, which is k+1 and NOT
11430                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
11431                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
11432                // panic in the sampled ON arm, measured before this line said k+1).
11433                let vt_cap = (k.max(k_cap) + 1).max(2);
11434                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
11435                if g.is_some() {
11436                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
11437                    // than trusting that a flag set means a pool built.
11438                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
11439                } else {
11440                    eprintln!(
11441                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
11442                         non-uniform state, or vt_cap < 2) — eager walk"
11443                    );
11444                }
11445            }
11446            Some(g)
11447        } else {
11448            None
11449        };
11450        // Capacity fail-safe: a round wider than the pool was built for must take the eager
11451        // walk, not slice the stash past its rows. The sizing above already covers every
11452        // round this run can present; this keeps a future caller (or a k that grows behind
11453        // the pool's back) on the byte-identical fallback instead of a panic.
11454        let vg_t_cap = vg_guard
11455            .as_ref()
11456            .and_then(|g| g.as_ref())
11457            .map(|g| g.t_capacity())
11458            .unwrap_or(0);
11459        if let Some(p) = pipe {
11460            p.setup_end();
11461        }
11462        while keep_going && out.len() < max_new {
11463            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
11464            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
11465            if let (true, Some(sg), Some(ptrs)) = (
11466                stream_active && round >= 1 && pending.is_some(),
11467                &stream_graph,
11468                &stream_ptrs,
11469            ) {
11470                if debug_spec {
11471                    static ONCE: std::sync::Once = std::sync::Once::new();
11472                    ONCE.call_once(|| {
11473                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
11474                    });
11475                }
11476                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
11477                e.set_u32_one(&mut pend_d, pending.unwrap())?;
11478                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
11479                for _mi in 0..m_rounds {
11480                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
11481                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
11482                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
11483                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
11484                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
11485                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11486                    sg.launch()?;
11487                    e.spec_assemble_verify(
11488                        &g_tokp2k,
11489                        &pend_d,
11490                        d2t_dev.as_ref(),
11491                        &mut vtok_d,
11492                        &mut brk_d,
11493                        p_min,
11494                        k,
11495                        pmin0,
11496                    )?;
11497                    let mut ck = VerifyCkpt::new(self.layers.len());
11498                    let dummy = vec![0u32; t_v_s];
11499                    let (tl_d, vx) = self.decode_step_t_core_stream(
11500                        e,
11501                        &dummy,
11502                        0,
11503                        &mut *cache,
11504                        embd_dev,
11505                        Some(&mut ck),
11506                        Some((&vtok_d, &pos_ctr)),
11507                        None,
11508                        None,
11509                        None,
11510                    )?;
11511                    for j in 0..t_v_s {
11512                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
11513                    }
11514                    e.spec_accept_greedy_dc(
11515                        &preds_d,
11516                        &vtok_d,
11517                        &last_pred_d,
11518                        &brk_d,
11519                        &mut stream_acc,
11520                    )?;
11521                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
11522                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
11523                    self.commit_verified_prefix_stream(
11524                        e,
11525                        &mut *cache,
11526                        &snap,
11527                        &ck,
11528                        &stream_acc,
11529                        1,
11530                        t_v_s,
11531                    )?;
11532                    e.spec_rollback_stream(
11533                        ptrs,
11534                        &pos_start_d,
11535                        &stream_acc,
11536                        1,
11537                        self.layers.len() + 1,
11538                    )?;
11539                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
11540                }
11541                e.stream().synchronize()?;
11542                let ring_h = e.dtoh_u32(&ring_d)?;
11543                let cnt = ring_h[0] as usize;
11544                for i in 0..cnt {
11545                    if out.len() < max_new {
11546                        out.push(ring_h[1 + i]);
11547                    }
11548                }
11549                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
11550                for il in 0..self.layers.len() {
11551                    if let Some(kvl) = cache.kv[il].as_mut() {
11552                        kvl.len = pos_h;
11553                    }
11554                }
11555                cache.pos = pos_h;
11556                scratch.kv.len = pos_h;
11557                pending = Some(ring_h[cnt]); // last drained token = the live bonus
11558                last_token = ring_h[cnt];
11559                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
11560                total_accepted += cnt.saturating_sub(m_rounds);
11561                if let Some(t) = sess_telem {
11562                    // totals only — the burst's per-round accept counts stayed on device
11563                    // (that is the point of the round-stream arm). pos_* untouched.
11564                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
11565                }
11566                round += m_rounds;
11567                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
11568                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11569                continue;
11570            }
11571            let pipe_draft = match pipe {
11572                Some(p) => Some(p.draft_begin(round)?),
11573                None => None,
11574            };
11575            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
11576            let mut current_opti = carried_opti.take();
11577            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
11578                match opti_fork.as_mut() {
11579                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
11580                    None => None,
11581                    Some(_) => None,
11582                }
11583            } else {
11584                None
11585            };
11586            if current_opti.is_none() {
11587                if let Some(fork) = opti_fork.as_ref() {
11588                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
11589                } else {
11590                    cache.snapshot_into(e, &mut snap)?;
11591                }
11592            } else if snap.pos != pos {
11593                return Err(format!(
11594                    "optipipe carried snapshot pos {} != current pos {pos}",
11595                    snap.pos
11596                )
11597                .into());
11598            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
11599            ph_mark(&mut ph_rest, phase_on);
11600
11601            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
11602            // p-min semantics (both paths): stop the chain early when the head's confidence in
11603            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
11604            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
11605            let base0 = if pending.is_some() { 1usize } else { 0usize };
11606            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
11607            // accepted run + 1 (the gemma law — see the setup block above the loop).
11608            let k_this = if adapt { kc } else { k };
11609            let mut draft: Vec<u32> = Vec::with_capacity(k);
11610            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
11611            let mut controller_draft_prob: Option<f32> = None;
11612            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
11613            if let Some(ticket) = current_opti.as_mut() {
11614                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
11615                if ticket.verify_tokens[0] != carried_pending {
11616                    return Err(format!(
11617                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
11618                        ticket.verify_tokens[0],
11619                    )
11620                    .into());
11621                }
11622                draft.push(ticket.verify_tokens[1]);
11623                controller_draft_prob = Some(ticket.draft_prob);
11624                controller_eager_state = ticket
11625                    .take_eager_seed()
11626                    .map(|seed| (ticket.verify_tokens[1], seed));
11627            } else {
11628                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
11629                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
11630                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
11631                // rejected drafts and p-min extras via the len mechanism).
11632                scratch.set_len(e, pos + base0 - 1)?;
11633                if pen_on {
11634                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
11635                    // device dedup: `penalty_last_n` is usize::MAX for any serve request with
11636                    // a penalty, so without the cap this grew with the whole session.
11637                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
11638                    let w0 = pen_hist.len().saturating_sub(win);
11639                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
11640                }
11641                if sampled {
11642                    draft_logits.clear();
11643                    draft_stats.clear();
11644                }
11645                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
11646                // position's mask is computed on that clone and advanced by the PROPOSED token. The
11647                // real state moves only on emission (verify's job), so the emitted stream is
11648                // unchanged — the mask only removes tokens the verify would have truncated anyway.
11649                let mut dmask_live = dmask_on;
11650                if dmask_live {
11651                    let t_c = std::time::Instant::now();
11652                    constraint
11653                        .as_deref_mut()
11654                        .unwrap()
11655                        .draft_begin()
11656                        .map_err(|e2| format!("constraint: {e2}"))?;
11657                    dm_clone_ns += t_c.elapsed().as_nanos();
11658                    dm_rounds += 1;
11659                }
11660                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
11661                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
11662                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
11663                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
11664                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11665                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
11666                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11667                    for j in 0..k_this {
11668                        // per-position mask upload (contents only — the graph's baked pointer is
11669                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
11670                        // mask node degrades to a no-op ban instead of needing a second graph.
11671                        if dmask_live
11672                            && !upload_draft_mask(
11673                                e,
11674                                constraint.as_deref_mut().unwrap(),
11675                                &mut dctx.g_dmask,
11676                                mtp.d2t.as_ref(),
11677                                d_vocab,
11678                                dmask_words,
11679                            )?
11680                        {
11681                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
11682                            // genuinely miss the legal set): neutralize the captured mask node and
11683                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
11684                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11685                            dmask_live = false;
11686                        }
11687                        gr.launch()?;
11688                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11689                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11690                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
11691                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
11692                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
11693                        // replay's embed node, and the MMU fault kills the CUDA context for the
11694                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
11695                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
11696                        // buffer (g_seed = the verify-side handoff vs head-side compute).
11697                        if (idx as usize) >= d_vocab {
11698                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
11699                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
11700                            // seed, untouched since the round-start copy — the pair discriminates
11701                            // "seed arrived poisoned" from "head forward produced NaN".
11702                            let seed_h = e.dtoh(&dctx.g_seed)?;
11703                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11704                            let in_h = e.dtoh(&h_seed_buf)?;
11705                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
11706                            return Err(format!(
11707                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11708                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
11709                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
11710                             the embed row (#87 trap)"
11711                            )
11712                            .into());
11713                        }
11714                        // trimmed draft vocab -> target token id (identity when no d2t map)
11715                        let d = match &mtp.d2t {
11716                            Some(map) => map[idx as usize],
11717                            None => idx,
11718                        };
11719                        let draft_p = if p_min > 0.0
11720                            || opti_fork
11721                                .as_ref()
11722                                .is_some_and(|fork| fork.controller.is_some())
11723                        {
11724                            Some(e.dtoh(&dctx.g_p)?[0])
11725                        } else {
11726                            None
11727                        };
11728                        if j == 0 {
11729                            controller_draft_prob = draft_p;
11730                        }
11731                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11732                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11733                                break;
11734                            }
11735                        }
11736                        draft.push(d);
11737                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
11738                        // index the argmax wrote — patch the persistent token buffer (4B htod).
11739                        if d != idx {
11740                            e.set_u32_one(&mut dctx.g_tok, d)?;
11741                        }
11742                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
11743                        // unmasked drafting for the remaining positions (verify still arbitrates).
11744                        // speculative advance; a chain the grammar can no longer follow (EOS
11745                        // proposed) ends here. The captured mask node always runs, so a dead chain
11746                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
11747                        if dmask_live
11748                            && !constraint
11749                                .as_deref_mut()
11750                                .unwrap()
11751                                .draft_advance(d)
11752                                .map_err(|e2| format!("constraint: {e2}"))?
11753                        {
11754                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
11755                            break;
11756                        }
11757                    }
11758                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
11759                // legal ONLY in the regime it was captured in. The condition used to read
11760                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
11761                // which it could not, because the key omitted the filters. Both halves are now
11762                // enforced: the key drops a stale graph, and this site refuses to launch one.
11763                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
11764                    if skey_probe() {
11765                        eprintln!(
11766                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
11767                             top_p={} min_p={} s_key_parked={:?}",
11768                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11769                        );
11770                    }
11771                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
11772                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
11773                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
11774                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
11775                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
11776                    // stream. Host sctr advances in lockstep (computed, no readback needed).
11777                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
11778                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
11779                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
11780                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
11781                    for j in 0..k_this {
11782                        gr.launch()?;
11783                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
11784                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
11785                        // counts the p-min-discarded token too)
11786                        // q retention: ONE async D2D of the persistent head-logits buffer into this
11787                        // round's slot j (stream-ordered after the replay, before the next one).
11788                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
11789                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
11790                        // #87 SENTINEL TRAP (see the greedy graph arm above).
11791                        if (idx as usize) >= d_vocab {
11792                            let seed_h = e.dtoh(&dctx.g_seed)?;
11793                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11794                            return Err(format!(
11795                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
11796                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
11797                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
11798                             (#87 trap)"
11799                            )
11800                            .into());
11801                        }
11802                        let d = match &mtp.d2t {
11803                            Some(map) => map[idx as usize],
11804                            None => idx,
11805                        };
11806                        draft_idx.push(idx);
11807                        if p_min > 0.0 {
11808                            let p = e.dtoh(&dctx.g_p)?[0];
11809                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11810                                break;
11811                            }
11812                        }
11813                        draft.push(d);
11814                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
11815                        if d != idx {
11816                            e.set_u32_one(&mut dctx.g_tok, d)?;
11817                        }
11818                    }
11819                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
11820                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
11821                    for j in 0..draft.len().max(draft_idx.len()) {
11822                        let rows0 = e.htod_i32(&[0])?;
11823                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11824                        e.filter_stats(
11825                            &dctx.q_slots[j],
11826                            d_vocab,
11827                            &rows0,
11828                            &mut th_d,
11829                            &mut z_d,
11830                            &mut mx_d,
11831                            d_vocab,
11832                            1,
11833                            sp_temp,
11834                            sp.top_k,
11835                            sp.top_p,
11836                            sp.min_p,
11837                        )?;
11838                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
11839                    }
11840                } else {
11841                    if skey_probe() && sampled {
11842                        eprintln!(
11843                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
11844                             top_p={} min_p={} s_key_parked={:?}",
11845                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
11846                        );
11847                    }
11848                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
11849                    let chain_heads = !self.mtp_extra.is_empty();
11850                    let mut e_tok = last_token;
11851                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
11852                    let mut chain_tokens = if chain_heads {
11853                        vec![last_token]
11854                    } else {
11855                        Vec::new()
11856                    };
11857                    let mut chain_seeds = if chain_heads {
11858                        vec![e.clone_dtod(&h_seed_buf)?]
11859                    } else {
11860                        Vec::new()
11861                    };
11862                    for j in 0..k_this {
11863                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
11864                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
11865                        let mtp_pos = pos + base0 + j;
11866                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
11867                        // A position with no legal draft-vocab row drops to unmasked drafting for
11868                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
11869                        if dmask_live {
11870                            dmask_live = upload_draft_mask(
11871                                e,
11872                                constraint.as_deref_mut().unwrap(),
11873                                &mut dctx.g_dmask,
11874                                mtp.d2t.as_ref(),
11875                                d_vocab,
11876                                dmask_words,
11877                            )?;
11878                        }
11879                        let mask = if dmask_live {
11880                            Some((&dctx.g_dmask, dmask_words))
11881                        } else {
11882                            None
11883                        };
11884                        let (dl_d, h_nextn) = if chain_heads {
11885                            if debug_spec {
11886                                eprintln!(
11887                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
11888                                    mtp_chain_head_index(j, self.mtp_head_count()),
11889                                    chain_tokens.len(),
11890                                );
11891                            }
11892                            self.mtp_chain_forward_dev(
11893                                e,
11894                                &chain_tokens,
11895                                &chain_seeds,
11896                                &mut *scratch,
11897                                pos + base0 - 1,
11898                                embd_dev,
11899                                mask,
11900                            )?
11901                        } else {
11902                            self.mtp_head_forward_dev(
11903                                e,
11904                                mtp,
11905                                e_tok,
11906                                &d_seed,
11907                                &mut *scratch,
11908                                mtp_pos,
11909                                embd_dev,
11910                                mask,
11911                            )?
11912                        };
11913                        let tok_d = if sampled {
11914                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
11915                            // the filtered softmax (filters off => th=0, exact v1 semantics).
11916                            if perturb_buf.is_none() {
11917                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
11918                            }
11919                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
11920                            if pen_on {
11921                                let h = pen_hist_d.as_ref().unwrap();
11922                                let nh = h.len();
11923                                e.penalize_logits(
11924                                    &mut q_row,
11925                                    h,
11926                                    nh,
11927                                    sp.penalty_repeat,
11928                                    sp.penalty_freq,
11929                                    sp.penalty_present,
11930                                    d_vocab,
11931                                )?;
11932                            }
11933                            let rows0 = e.htod_i32(&[0])?;
11934                            let (mut th_d, mut z_d, mut mx_d) =
11935                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
11936                            e.filter_stats(
11937                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
11938                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
11939                            )?;
11940                            let (th, z, mx) =
11941                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
11942                            let pb = perturb_buf.as_mut().unwrap();
11943                            e.gumbel_perturb_filtered(
11944                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
11945                            )?;
11946                            sctr += 1;
11947                            draft_logits.push(q_row);
11948                            draft_stats.push((mx, th, z));
11949                            e.argmax_token_device(pb, d_vocab)?
11950                        } else {
11951                            e.argmax_token_device(&dl_d, d_vocab)?
11952                        };
11953                        let idx = e.dtoh_u32_one(&tok_d)?;
11954                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
11955                        // here because the eager chain's operands are all readable: dl_d (the head
11956                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
11957                        if (idx as usize) >= d_vocab {
11958                            let dl_h = e.dtoh(&dl_d)?;
11959                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
11960                            let seed_h = if chain_heads {
11961                                e.dtoh(chain_seeds.last().unwrap())?
11962                            } else {
11963                                e.dtoh(&d_seed)?
11964                            };
11965                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
11966                            return Err(format!(
11967                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
11968                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
11969                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
11970                             embed row (#87 trap)"
11971                            )
11972                            .into());
11973                        }
11974                        let d = match &mtp.d2t {
11975                            Some(map) => map[idx as usize],
11976                            None => idx,
11977                        };
11978                        if sampled {
11979                            draft_idx.push(idx);
11980                        }
11981                        let draft_p = if p_min > 0.0
11982                            || opti_fork
11983                                .as_ref()
11984                                .is_some_and(|fork| fork.controller.is_some())
11985                        {
11986                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
11987                            Some(e.dtoh(&p_d)?[0])
11988                        } else {
11989                            None
11990                        };
11991                        if j == 0 {
11992                            controller_draft_prob = draft_p;
11993                        }
11994                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
11995                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
11996                                break;
11997                            }
11998                        }
11999                        draft.push(d);
12000                        if chain_heads {
12001                            chain_tokens.push(d);
12002                            chain_seeds.push(h_nextn);
12003                        } else {
12004                            e_tok = d;
12005                            d_seed = h_nextn;
12006                        }
12007                        // speculative advance; a chain the grammar can no longer follow (EOS
12008                        // proposed) ends here — the prefix already proposed still rides verify.
12009                        if dmask_live
12010                            && !constraint
12011                                .as_deref_mut()
12012                                .unwrap()
12013                                .draft_advance(d)
12014                                .map_err(|e2| format!("constraint: {e2}"))?
12015                        {
12016                            break;
12017                        }
12018                    }
12019                    if !chain_heads
12020                        && opti_fork
12021                            .as_ref()
12022                            .is_some_and(|fork| fork.controller.is_some())
12023                    {
12024                        controller_eager_state = Some((e_tok, d_seed));
12025                    }
12026                }
12027            }
12028            let k_round = draft.len();
12029            if let Some(p) = pipe {
12030                p.draft_end(round);
12031            }
12032            drop(pipe_draft);
12033
12034            ph_mark(&mut ph_draft, phase_on);
12035            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
12036            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
12037            let verify_tokens: Vec<u32> = match pending {
12038                Some(b) => {
12039                    let mut v = Vec::with_capacity(k_round + 1);
12040                    v.push(b);
12041                    v.extend_from_slice(&draft);
12042                    v
12043                }
12044                None => draft.clone(),
12045            };
12046            let base = if pending.is_some() { 1 } else { 0 };
12047            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
12048            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
12049            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
12050                Some(ticket.take_ckpt())
12051            } else if spec_replay {
12052                None
12053            } else {
12054                Some(VerifyCkpt::new(self.layers.len()))
12055            };
12056            let controller_can_probe = base == 1
12057                && k_round == 1
12058                && out.len().saturating_add(2) < max_new
12059                && controller_draft_prob.is_some()
12060                && opti_fork
12061                    .as_ref()
12062                    .and_then(|fork| fork.controller.as_ref())
12063                    .is_some_and(|policy| !policy.breaker_tripped);
12064            let mut successor_attempt: Option<OptiControllerTicket> = None;
12065            let mut rejected_probe: Option<(f32, u32)> = None;
12066            let mut controller_prepared: Option<OptiControllerPrepared> = None;
12067            if controller_can_probe {
12068                // Prepare d2/q and, on admission, d3 before either current verify half is
12069                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
12070                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
12071                // the primary stream after N stage 1 would serialize the supposed pipeline.
12072                let eager_pos = scratch.kv.len + 1;
12073                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
12074                    e,
12075                    mtp,
12076                    &mut dctx,
12077                    &mut *scratch,
12078                    d_vocab,
12079                    &mut controller_eager_state,
12080                    eager_pos,
12081                    embd_dev,
12082                )?;
12083                let first_probability = controller_draft_prob
12084                    .ok_or("optipipe controller probe lost first-token probability")?;
12085                let q_proxy = first_probability * pending_probability;
12086                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12087                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12088                let admitted = opti_fork
12089                    .as_ref()
12090                    .and_then(|fork| fork.controller.as_ref())
12091                    .ok_or("optipipe controller policy disappeared")?
12092                    .admit(q_proxy);
12093                if admitted {
12094                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12095                    let eager_pos = scratch.kv.len + 1;
12096                    let (optimistic_draft, optimistic_draft_probability) = self
12097                        .opti_controller_draft_step(
12098                            e,
12099                            mtp,
12100                            &mut dctx,
12101                            &mut *scratch,
12102                            d_vocab,
12103                            &mut controller_eager_state,
12104                            eager_pos,
12105                            embd_dev,
12106                        )?;
12107                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12108                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
12109                        debug_assert_eq!(token, optimistic_draft);
12110                        seed
12111                    });
12112                    controller_prepared = Some(OptiControllerPrepared {
12113                        verify_tokens: [optimistic_pending, optimistic_draft],
12114                        draft_prob: optimistic_draft_probability,
12115                        eager_seed,
12116                        q_proxy,
12117                        scratch_len: scratch.kv.len,
12118                    });
12119                } else {
12120                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12121                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12122                    rejected_probe = Some((q_proxy, optimistic_pending));
12123                    eprintln!(
12124                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
12125                        opti_fork
12126                            .as_ref()
12127                            .and_then(|fork| fork.controller.as_ref())
12128                            .expect("controller policy")
12129                            .threshold,
12130                    );
12131                }
12132            }
12133            let fork_attempt = match fork_generation.take() {
12134                Some(generation) if base == 1 && k_round == 1 => Some(generation),
12135                Some(generation) => {
12136                    opti_fork
12137                        .as_mut()
12138                        .expect("fork generation without fork state")
12139                        .retire(generation)?;
12140                    None
12141                }
12142                None => None,
12143            };
12144            let (tlogits_d, vx) = if let Some(p) = pipe {
12145                self.decode_step_t_core_pipelined(
12146                    e,
12147                    &verify_tokens,
12148                    pos,
12149                    &mut *cache,
12150                    embd_dev,
12151                    ckpt.as_mut(),
12152                    p,
12153                    round,
12154                )?
12155            } else if controller_can_probe {
12156                let fence = opti_fork
12157                    .as_ref()
12158                    .ok_or("optipipe controller probe lost fork state")?
12159                    .fence;
12160                let boundary = match current_opti.as_mut() {
12161                    Some(ticket) => ticket.take_boundary(),
12162                    None => self.verify_stage0_issue(
12163                        e,
12164                        &verify_tokens,
12165                        pos,
12166                        &mut *cache,
12167                        embd_dev,
12168                        ckpt.as_mut(),
12169                        None,
12170                        &fence,
12171                        Some(true),
12172                        None,
12173                    )?,
12174                };
12175                if let Some(prepared) = controller_prepared.take() {
12176                    let generation = {
12177                        let fork = opti_fork
12178                            .as_mut()
12179                            .ok_or("optipipe controller admission lost fork state")?;
12180                        let generation = fork.reserve_successor()?;
12181                        let rt = fork.rt;
12182                        let snapshot_fence = fork.fence;
12183                        opti_snapshot_one_stage_owned_into(
12184                            e,
12185                            cache,
12186                            rt,
12187                            &snapshot_fence,
12188                            0,
12189                            fork.successor_snapshot_mut(),
12190                        )?;
12191                        generation
12192                    };
12193                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
12194                    let successor_boundary = self.verify_stage0_issue(
12195                        e,
12196                        &prepared.verify_tokens,
12197                        pos + verify_tokens.len(),
12198                        &mut *cache,
12199                        embd_dev,
12200                        Some(&mut successor_ckpt),
12201                        None,
12202                        &fence,
12203                        Some(false),
12204                        None,
12205                    )?;
12206                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12207                    let fork = opti_fork
12208                        .as_ref()
12209                        .ok_or("optipipe controller ticket lost fork state")?;
12210                    successor_attempt = Some(fork.controller_ticket(
12211                        generation,
12212                        successor_boundary,
12213                        successor_ckpt,
12214                        prepared.verify_tokens,
12215                        prepared.draft_prob,
12216                        prepared.eager_seed,
12217                        prepared.q_proxy,
12218                        prepared.scratch_len,
12219                    ));
12220                    eprintln!(
12221                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
12222                         verify={:?}",
12223                        generation.id,
12224                        prepared.q_proxy,
12225                        fork.controller.expect("controller policy").threshold,
12226                        prepared.verify_tokens,
12227                    );
12228                }
12229                let result = self.verify_stage1_finish(
12230                    e,
12231                    boundary,
12232                    &mut *cache,
12233                    ckpt.as_mut(),
12234                    None,
12235                    &fence,
12236                    successor_attempt.is_none(),
12237                )?;
12238                if let Some(ticket) = current_opti.as_mut() {
12239                    ticket.settle();
12240                }
12241                if successor_attempt.is_some() {
12242                    let fork = opti_fork
12243                        .as_mut()
12244                        .ok_or("optipipe successor snapshot lost fork state")?;
12245                    let rt = fork.rt;
12246                    let snapshot_fence = fork.fence;
12247                    opti_snapshot_one_stage_owned_into(
12248                        e,
12249                        cache,
12250                        rt,
12251                        &snapshot_fence,
12252                        1,
12253                        fork.successor_snapshot_mut(),
12254                    )?;
12255                    // Publish N only after both independent successor-state queues are complete.
12256                    fork.rt.publish_to(1, &e.stream())?;
12257                }
12258                result
12259            } else if let Some(ticket) = current_opti.as_mut() {
12260                let fork = opti_fork
12261                    .as_mut()
12262                    .ok_or("optipipe carried controller ticket lost fork state")?;
12263                let boundary = ticket.take_boundary();
12264                let result = self.verify_stage1_finish(
12265                    e,
12266                    boundary,
12267                    &mut *cache,
12268                    ckpt.as_mut(),
12269                    None,
12270                    &fork.fence,
12271                    true,
12272                )?;
12273                ticket.settle();
12274                result
12275            } else if let Some(generation) = fork_attempt {
12276                let fork = opti_fork
12277                    .as_mut()
12278                    .expect("fork generation without fork state");
12279                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
12280                let action = fork.mode.action(generation.id);
12281                let boundary = self.verify_stage0_issue(
12282                    e,
12283                    &verify_tokens,
12284                    pos,
12285                    &mut *cache,
12286                    embd_dev,
12287                    ckpt.as_mut(),
12288                    None,
12289                    &fork.fence,
12290                    Some(true),
12291                    None,
12292                )?;
12293                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12294                let mut ticket = fork.ticket(generation, boundary);
12295                if action == OptiForkAction::Abort {
12296                    return Err(format!(
12297                        "optipipe forced abort with generation {} stage0 in flight",
12298                        generation.id,
12299                    )
12300                    .into());
12301                }
12302                fork.reconcile(
12303                    e,
12304                    &mut *cache,
12305                    &mut *scratch,
12306                    &snap,
12307                    &mut h_seed_buf,
12308                    &mut fill_prev,
12309                    generation,
12310                    action,
12311                    verify_tokens[0],
12312                )?;
12313                let result = if action == OptiForkAction::Hit {
12314                    let boundary = ticket.take_boundary();
12315                    self.verify_stage1_finish(
12316                        e,
12317                        boundary,
12318                        &mut *cache,
12319                        ckpt.as_mut(),
12320                        None,
12321                        &fork.fence,
12322                        true,
12323                    )?
12324                } else {
12325                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
12326                    // verify only after E_restart published the restored stage-0 state.
12327                    self.decode_step_t_core(
12328                        e,
12329                        &verify_tokens,
12330                        pos,
12331                        &mut *cache,
12332                        embd_dev,
12333                        ckpt.as_mut(),
12334                    )?
12335                };
12336                ticket.settle();
12337                debug_assert_eq!(ticket.generation, generation);
12338                fork.retire(generation)?;
12339                result
12340            } else {
12341                // The serial verify every non-fork round takes — the MTP route's
12342                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
12343                // a pool above, and then the walk replays the captured trunk instead of
12344                // re-issuing it launch by launch.
12345                let vg_round = if verify_tokens.len() <= vg_t_cap {
12346                    vg_guard.as_mut().and_then(|g| g.as_mut())
12347                } else {
12348                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
12349                        // The commit reads this flag to pick its arm; a round that declines
12350                        // the pool must not inherit a stale `true` from the round before it.
12351                        g.round_slab = false;
12352                    }
12353                    None
12354                };
12355                self.decode_step_t_core_vg(
12356                    e,
12357                    &verify_tokens,
12358                    pos,
12359                    &mut *cache,
12360                    embd_dev,
12361                    ckpt.as_mut(),
12362                    vg_round,
12363                )?
12364            };
12365            let pipe_accept = match pipe {
12366                Some(p) => Some(p.accept_begin(round)?),
12367                None => None,
12368            };
12369
12370            ph_mark(&mut ph_verify, phase_on);
12371            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
12372            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
12373            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
12374            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
12375            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
12376            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
12377            // (== the bonus), so every index shifts by `base` and last_pred is unused.
12378            let t_v = verify_tokens.len();
12379            let mut preds: Vec<u32> = Vec::new();
12380            if !sampled {
12381                for j in 0..t_v {
12382                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
12383                }
12384                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
12385                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
12386                // next round's last_token = the next chain's embed lookup. Catch it at the
12387                // source with the column named — an all-NaN VERIFY column implicates the
12388                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
12389                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
12390                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
12391                    let mut probe = e.zeros(n_vocab)?;
12392                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
12393                    let col_h = e.dtoh(&probe)?;
12394                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
12395                    return Err(format!(
12396                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
12397                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
12398                         — the stage-split verify produced a poisoned column (#87 trap)",
12399                        preds[bad]
12400                    )
12401                    .into());
12402                }
12403            }
12404            ph_mark(&mut ph_wait, phase_on);
12405            let t_pred = |j: usize| -> u32 {
12406                if j == 0 && base == 0 {
12407                    last_pred
12408                } else {
12409                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
12410                    // used to call this from the sampled arm and panicked the worker; it now goes
12411                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
12412                    // out-of-range pred is a real bug, not something to paper over.
12413                    debug_assert!(
12414                        !sampled,
12415                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
12416                    );
12417                    preds[base + j - 1]
12418                }
12419            };
12420            let mut devacc_seeded = false;
12421            let mut devacc_acc: Option<CudaSlice<u32>> = None;
12422            let (n_acc, bonus) = if !sampled {
12423                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
12424                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
12425                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
12426                // gated on token identity vs the host walk (the arms below are bit-equal rules).
12427                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
12428                {
12429                    let draft_d = e.htod_u32_v(&draft)?;
12430                    let mut acc_out = e.alloc_u32_zeroed(2)?;
12431                    e.spec_accept_greedy(
12432                        &preds_d,
12433                        &draft_d,
12434                        last_pred,
12435                        base,
12436                        k_round,
12437                        &mut acc_out,
12438                    )?;
12439                    devacc_acc = Some(acc_out.clone());
12440                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
12441                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
12442                    // non-replay commit arms skip their host-offset seed copies (guarded below);
12443                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
12444                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
12445                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
12446                    // the update lands after the arms (devacc_seeded guard below).
12447                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
12448                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
12449                    // unified rule; full accept rewrites the verify-left value). Host mirrors
12450                    // update after the readback; commit_verified_prefix skips its len_d writes.
12451                    if let Some(successor) = successor_attempt.as_ref() {
12452                        opti_fork
12453                            .as_mut()
12454                            .ok_or("optipipe successor reconcile lost fork state")?
12455                            .queue_actual_reconcile(
12456                                e,
12457                                &snap,
12458                                &acc_out,
12459                                successor.verify_tokens[0],
12460                                base,
12461                            )?;
12462                    } else if let Some(ptrs) = &kv_len_ptrs {
12463                        let saved: Vec<i32> = (0..self.layers.len())
12464                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
12465                            .collect();
12466                        let saved_d = e.htod_i32(&saved)?;
12467                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
12468                    }
12469                    devacc_seeded = true;
12470                    let ab = e.dtoh_u32(&acc_out)?;
12471                    (ab[0] as usize, ab[1])
12472                } else {
12473                    let mut n_acc = 0usize;
12474                    for j in 0..k_round {
12475                        if t_pred(j) == draft[j] {
12476                            n_acc += 1;
12477                        } else {
12478                            break;
12479                        }
12480                    }
12481                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
12482                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
12483                    (n_acc, t_pred(n_acc))
12484                }
12485            } else {
12486                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
12487                if col_buf.is_none() {
12488                    col_buf = Some(e.zeros(n_vocab)?);
12489                }
12490                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
12491                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
12492                let mut pj = vec![0f32; k_round.max(1)];
12493                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
12494                if k_round > 0 {
12495                    let mut ids: Vec<u32> = Vec::new();
12496                    let mut rows: Vec<i32> = Vec::new();
12497                    for j in 0..k_round {
12498                        if j > 0 || base == 1 {
12499                            ids.push(draft[j]);
12500                            rows.push((base + j) as i32 - 1);
12501                        }
12502                    }
12503                    if !ids.is_empty() {
12504                        let nr = rows.len();
12505                        // penalties: materialize the used columns into one contiguous penalized
12506                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
12507                        // penalties: materialize used columns contiguously, penalize all rows in
12508                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
12509                        let p_rows: Vec<i32> = if pen_on {
12510                            (0..nr as i32).collect()
12511                        } else {
12512                            rows.clone()
12513                        };
12514                        if pen_on {
12515                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
12516                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
12517                            }
12518                            let pc = pcol_buf.as_mut().unwrap();
12519                            for (i2, &r) in rows.iter().enumerate() {
12520                                let c = r as usize;
12521                                e.copy_view_into(
12522                                    pc,
12523                                    i2 * n_vocab,
12524                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
12525                                    n_vocab,
12526                                )?;
12527                            }
12528                            let h = pen_hist_d.as_ref().unwrap();
12529                            let nh = h.len();
12530                            e.penalize_logits_rows(
12531                                pc,
12532                                h,
12533                                nh,
12534                                sp.penalty_repeat,
12535                                sp.penalty_freq,
12536                                sp.penalty_present,
12537                                n_vocab,
12538                                nr,
12539                            )?;
12540                        }
12541                        let p_src: &CudaSlice<f32> = if pen_on {
12542                            pcol_buf.as_ref().unwrap()
12543                        } else {
12544                            &tlogits_d
12545                        };
12546                        let rowsd = e.htod_i32(&p_rows)?;
12547                        let (mut th_d, mut z_d, mut mx_d) =
12548                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
12549                        e.filter_stats(
12550                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
12551                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12552                        )?;
12553                        let idsd = e.htod_u32_v(&ids)?;
12554                        let mut outd = e.zeros(nr)?;
12555                        e.softmax_gather_filtered(
12556                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
12557                            sp_temp,
12558                        )?;
12559                        let outv = e.dtoh(&outd)?;
12560                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
12561                        let mut oi = 0usize;
12562                        for j in 0..k_round {
12563                            if j > 0 || base == 1 {
12564                                pj[j] = outv[oi];
12565                                oi += 1;
12566                            }
12567                        }
12568                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
12569                    }
12570                    if base == 0 {
12571                        let lc: &CudaSlice<f32> = if pen_on {
12572                            if col_buf.is_none() {
12573                                col_buf = Some(e.zeros(n_vocab)?);
12574                            }
12575                            let cb = col_buf.as_mut().unwrap();
12576                            e.copy_into(
12577                                cb,
12578                                0,
12579                                last_col_logits
12580                                    .as_ref()
12581                                    .expect("sampled: last_col_logits unset"),
12582                                n_vocab,
12583                            )?;
12584                            let h = pen_hist_d.as_ref().unwrap();
12585                            let nh = h.len();
12586                            e.penalize_logits(
12587                                cb,
12588                                h,
12589                                nh,
12590                                sp.penalty_repeat,
12591                                sp.penalty_freq,
12592                                sp.penalty_present,
12593                                n_vocab,
12594                            )?;
12595                            col_buf.as_ref().unwrap()
12596                        } else {
12597                            last_col_logits
12598                                .as_ref()
12599                                .expect("sampled: last_col_logits unset")
12600                        };
12601                        let rows0 = e.htod_i32(&[0])?;
12602                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12603                        e.filter_stats(
12604                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12605                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12606                        )?;
12607                        let idsd = e.htod_u32_v(&[draft[0]])?;
12608                        let mut outd = e.zeros(1)?;
12609                        e.softmax_gather_filtered(
12610                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
12611                        )?;
12612                        pj[0] = e.dtoh(&outd)?[0];
12613                        last_col_stats =
12614                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12615                    }
12616                }
12617                // q source: the graph arm retained the head logits in the persistent q_slots;
12618                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
12619                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
12620                // computes them post-replay — graph engages only filter/penalty-free, so the
12621                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
12622                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
12623                    &dctx.q_slots
12624                } else {
12625                    &draft_logits
12626                };
12627                let mut n_acc = 0usize;
12628                for j in 0..k_round {
12629                    let (qmx, qth, qz) = draft_stats[j];
12630                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
12631                    let rowsd = e.htod_i32(&[0])?;
12632                    let thd = e.htod(&[qth])?;
12633                    let zd = e.htod(&[qz])?;
12634                    let _ = qmx;
12635                    let mut outd = e.zeros(1)?;
12636                    e.softmax_gather_filtered(
12637                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
12638                        sp_temp,
12639                    )?;
12640                    let qj = e.dtoh(&outd)?[0];
12641                    let u = host_u01(sp_seed, uctr);
12642                    uctr += 1;
12643                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
12644                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
12645                    // exactness signature (see `skey_probe`). Impossible when the draft was
12646                    // drawn from the same filtered distribution the verify reconstructs here;
12647                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
12648                    if skey_probe() && qj == 0.0 {
12649                        eprintln!(
12650                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
12651                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
12652                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
12653                        );
12654                    }
12655                    if accept {
12656                        n_acc += 1;
12657                    } else {
12658                        break;
12659                    }
12660                }
12661                let bonus = if n_acc == k_round {
12662                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
12663                    let col = base + k_round - 1;
12664                    let cb = col_buf.as_mut().unwrap();
12665                    e.copy_view_into(
12666                        cb,
12667                        0,
12668                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12669                        n_vocab,
12670                    )?;
12671                    if pen_on {
12672                        let h = pen_hist_d.as_ref().unwrap();
12673                        let nh = h.len();
12674                        e.penalize_logits(
12675                            cb,
12676                            h,
12677                            nh,
12678                            sp.penalty_repeat,
12679                            sp.penalty_freq,
12680                            sp.penalty_present,
12681                            n_vocab,
12682                        )?;
12683                    }
12684                    if perturb_buf.is_none() {
12685                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12686                    }
12687                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
12688                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
12689                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
12690                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
12691                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
12692                    // last gathered column, in both base arms. `th` is a threshold in e-units of
12693                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
12694                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
12695                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
12696                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
12697                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
12698                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
12699                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
12700                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
12701                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
12702                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
12703                    // and row_max is unused once nothing is masked), so this fix is a byte-level
12704                    // no-op for the untruncated serve default. One extra one-block filter_stats
12705                    // per full-accept round is the whole cost.
12706                    let (mx, th) = {
12707                        let rows0 = e.htod_i32(&[0])?;
12708                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12709                        let cb0 = col_buf.as_ref().unwrap();
12710                        e.filter_stats(
12711                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
12712                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
12713                        )?;
12714                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
12715                    };
12716                    let pb = perturb_buf.as_mut().unwrap();
12717                    let cb2 = col_buf.as_ref().unwrap();
12718                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
12719                    sctr += 1;
12720                    let td = e.argmax_token_device(pb, n_vocab)?;
12721                    e.dtoh_u32_one(&td)?
12722                } else {
12723                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
12724                    let cb = col_buf.as_mut().unwrap();
12725                    if n_acc > 0 || base == 1 {
12726                        let col = base + n_acc - 1;
12727                        e.copy_view_into(
12728                            cb,
12729                            0,
12730                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
12731                            n_vocab,
12732                        )?;
12733                    } else {
12734                        let lc = last_col_logits.as_ref().unwrap();
12735                        e.copy_into(cb, 0, lc, n_vocab)?;
12736                    }
12737                    if pen_on {
12738                        let h = pen_hist_d.as_ref().unwrap();
12739                        let nh = h.len();
12740                        e.penalize_logits(
12741                            cb,
12742                            h,
12743                            nh,
12744                            sp.penalty_repeat,
12745                            sp.penalty_freq,
12746                            sp.penalty_present,
12747                            n_vocab,
12748                        )?;
12749                    }
12750                    let cb2 = col_buf.as_ref().unwrap();
12751                    let sc = sctr;
12752                    sctr += 1;
12753                    // p-stats for the reject column: from col_stats when the col was gathered,
12754                    // else (j==0&&base==0) from last_col_stats.
12755                    let p_stats = if n_acc > 0 || base == 1 {
12756                        // col index within the gathered set == number of gathered cols before n_acc
12757                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
12758                        col_stats.get(gi).copied().unwrap_or_else(|| {
12759                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
12760                        })
12761                    } else {
12762                        last_col_stats.expect("sampled: last_col_stats unset at reject")
12763                    };
12764                    let q_stats = draft_stats[n_acc];
12765                    if let Some(map) = &d2t_dev {
12766                        if q_full_buf.is_none() {
12767                            q_full_buf = Some(e.zeros(n_vocab)?);
12768                        }
12769                        let qf = q_full_buf.as_mut().unwrap();
12770                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
12771                        let qf2 = q_full_buf.as_ref().unwrap();
12772                        e.residual_sample_filtered(
12773                            cb2,
12774                            Some(qf2),
12775                            n_vocab,
12776                            sp_temp,
12777                            sp_seed,
12778                            sc,
12779                            p_stats,
12780                            q_stats,
12781                            &mut sample_tok,
12782                        )?;
12783                    } else {
12784                        e.residual_sample_filtered(
12785                            cb2,
12786                            Some(&q_bufs[n_acc]),
12787                            n_vocab,
12788                            sp_temp,
12789                            sp_seed,
12790                            sc,
12791                            p_stats,
12792                            q_stats,
12793                            &mut sample_tok,
12794                        )?;
12795                    }
12796                    e.dtoh_u32(&sample_tok)?[0]
12797                };
12798                (n_acc, bonus)
12799            };
12800            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
12801            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
12802            // ordering). Walk the accepted drafts through the grammar in commit order; the
12803            // first illegal token truncates acceptance at its slot, and that slot's emission
12804            // is recomputed as the MASKED argmax of the target's own verify column — token-
12805            // identical to constrained plain greedy decode (an unmasked argmax that is
12806            // grammar-legal IS the masked argmax: masking only removes competitors). The
12807            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
12808            // measured in acceptance numbers, never hidden.
12809            let (n_acc, bonus) = match constraint.as_deref_mut() {
12810                None => (n_acc, bonus),
12811                Some(c) => {
12812                    fn ce(e2: String) -> Box<dyn std::error::Error> {
12813                        format!("constraint: {e2}").into()
12814                    }
12815                    let mut na = n_acc;
12816                    let mut cut = false;
12817                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
12818                        if c.is_allowed(d).map_err(ce)? {
12819                            c.consume(d).map_err(ce)?;
12820                        } else {
12821                            na = j;
12822                            cut = true;
12823                            dm_cut_tokens += n_acc - j;
12824                            break;
12825                        }
12826                    }
12827                    if cut {
12828                        dm_cuts += 1;
12829                    }
12830                    let mut bo = bonus;
12831                    if cut || !c.is_allowed(bo).map_err(ce)? {
12832                        let mut row = if na == 0 && base == 0 {
12833                            init_logits_host
12834                                .clone()
12835                                .ok_or("constraint: init logits missing (round-0 cut)")?
12836                        } else {
12837                            e.dtoh_view(
12838                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
12839                            )?
12840                        };
12841                        c.mask_logits(&mut row).map_err(ce)?;
12842                        bo = argmax(&row) as u32;
12843                    }
12844                    c.consume(bo).map_err(ce)?;
12845                    (na, bo)
12846                }
12847            };
12848            let mut successor_valid = false;
12849            if let Some((q_proxy, expected_d2)) = rejected_probe {
12850                let v_n = n_acc == 1 && bonus == expected_d2;
12851                eprintln!(
12852                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
12853                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
12854                );
12855            }
12856            if let Some(successor) = successor_attempt.as_ref() {
12857                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
12858                let generation = successor.generation;
12859                let q_proxy = successor.q_proxy;
12860                let expected_pending = successor.verify_tokens[0];
12861                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
12862                let fork = opti_fork
12863                    .as_mut()
12864                    .ok_or("optipipe successor resolution lost fork state")?;
12865                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
12866                if successor_valid {
12867                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12868                } else {
12869                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12870                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12871                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
12872                }
12873                let breaker_tripped = fork
12874                    .controller
12875                    .as_mut()
12876                    .expect("controller policy")
12877                    .resolve(successor_valid);
12878                if breaker_tripped {
12879                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12880                }
12881                eprintln!(
12882                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
12883                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
12884                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
12885                    generation.id, successor_valid, !successor_valid, breaker_tripped,
12886                );
12887                if !successor_valid {
12888                    let mut successor = successor_attempt
12889                        .take()
12890                        .expect("controller successor disappeared on miss");
12891                    successor.settle();
12892                    fork.retire(generation)?;
12893                }
12894            }
12895            total_drafted += k_round;
12896            total_accepted += n_acc;
12897            if let Some(t) = sess_telem {
12898                // Greedy, rejection-sampling, and grammar truncation all converge here after
12899                // the accept decision is already on host. Fixed-size relaxed atomics only.
12900                t.record_round(k_round, n_acc);
12901            }
12902            if spec_stats {
12903                st_len_hist[k_round] += 1;
12904                for j in 0..k_round {
12905                    st_drafted[j] += 1;
12906                }
12907                for j in 0..n_acc {
12908                    st_accepted[j] += 1;
12909                }
12910                if n_acc == k_round {
12911                    st_full += 1;
12912                }
12913            }
12914
12915            if debug_spec {
12916                eprintln!(
12917                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
12918                    out.len(),
12919                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
12920                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
12921                    // the GPU worker thread — a debug flag that killed the exact regime you would
12922                    // set it to investigate. See `debug_t_pred0`.
12923                    debug_t_pred0(sampled, base, last_pred, &preds)
12924                );
12925            }
12926
12927            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
12928            let commit_started = std::time::Instant::now();
12929            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
12930            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
12931            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
12932            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
12933            for j in 0..n_acc {
12934                if !session_mode && out.len() >= max_new {
12935                    break;
12936                }
12937                out.push(draft[j]);
12938            }
12939            if pen_on {
12940                pen_hist.extend_from_slice(&draft[0..n_acc]);
12941                pen_hist.push(bonus);
12942            }
12943            let bonus_emitted = session_mode || out.len() < max_new;
12944            if bonus_emitted {
12945                out.push(bonus);
12946            }
12947            last_token = bonus;
12948
12949            // --- 5. ROLLBACK + advance (§C) ---
12950            if n_acc == k_round && !spec_replay {
12951                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
12952                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
12953                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
12954                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
12955                // last_pred is dead in the pending path (t_pred reads verify col 0).
12956                //
12957                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
12958                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
12959                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
12960                // trunk hidden (the last verify column). set_len first: a p-min break may have
12961                // left one extra chain append at that slot. Partial accepts need NO fill (the
12962                // chain already covered every accepted position; round-start set_len truncates).
12963                let mut vh_seed = e.zeros(n_embd)?;
12964                e.copy_view_into(
12965                    &mut vh_seed,
12966                    0,
12967                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
12968                    n_embd,
12969                )?;
12970                if refresh {
12971                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
12972                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
12973                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
12974                    // the full stack (vx) is already resident from the verify. Replaces both the
12975                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
12976                    // (draft attention quality); exactness stays the verify's job.
12977                    scratch.set_len(e, pos)?;
12978                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
12979                    // (hidden of the last committed row before this verify batch).
12980                    let mut vxs = e.zeros(t_v * n_embd)?;
12981                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
12982                    if t_v > 1 {
12983                        e.copy_view_into(
12984                            &mut vxs,
12985                            n_embd,
12986                            &vx.slice(0..(t_v - 1) * n_embd),
12987                            (t_v - 1) * n_embd,
12988                        )?;
12989                    }
12990                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
12991                } else {
12992                    scratch.set_len(e, pos + base + k_round - 1)?;
12993                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
12994                    let mut hp = e.zeros(n_embd)?;
12995                    if t_v >= 2 {
12996                        e.copy_view_into(
12997                            &mut hp,
12998                            0,
12999                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
13000                            n_embd,
13001                        )?;
13002                    } else {
13003                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
13004                    }
13005                    self.mtp_kv_fill_all(
13006                        e,
13007                        &[draft[k_round - 1]],
13008                        &hp,
13009                        pos + base + k_round - 1,
13010                        &mut *scratch,
13011                        embd_dev,
13012                    )?;
13013                }
13014                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
13015                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
13016                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
13017                // col). Saves one MTP-block pass per round on top of the pairing fix.
13018                if !devacc_seeded {
13019                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
13020                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
13021                }
13022                pending = Some(bonus);
13023                if debug_spec {
13024                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
13025                }
13026            } else if !spec_replay && base + n_acc >= 1 {
13027                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
13028                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
13029                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
13030                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
13031                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
13032                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
13033                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
13034                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
13035                // accept (never compounds: the next verify recomputes true hiddens for all
13036                // committed columns).
13037                let j = base + n_acc;
13038                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
13039                // column stash was written into the graphs ctx's persistent slabs as in-graph
13040                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
13041                // commit must take the slab twin (same semantics, slab-addressed sources). The
13042                // ctx states which of the two this round produced via `round_slab`; trusting the
13043                // flag rather than the env keeps a round that fell back to the eager walk (a
13044                // capture that declined, a t the pool never captured) on the cols arm.
13045                let slab_commit = vg_guard
13046                    .as_ref()
13047                    .and_then(|g| g.as_ref())
13048                    .map(|g| g.round_slab)
13049                    .unwrap_or(false);
13050                if slab_commit {
13051                    self.dspark_commit_prefix_slab(
13052                        e,
13053                        &mut *cache,
13054                        &snap,
13055                        vg_guard
13056                            .as_ref()
13057                            .and_then(|g| g.as_ref())
13058                            .expect("slab_commit implies a graphs ctx"),
13059                        j,
13060                    )?;
13061                } else {
13062                    self.commit_verified_prefix(
13063                        e,
13064                        &mut *cache,
13065                        &snap,
13066                        ckpt.as_ref().unwrap(),
13067                        j,
13068                        devacc_seeded,
13069                        if devacc_seeded {
13070                            devacc_acc.as_ref().map(|a| (a, base, t_v))
13071                        } else {
13072                            None
13073                        },
13074                    )?;
13075                }
13076                let mut seed = e.zeros(n_embd)?;
13077                e.copy_view_into(
13078                    &mut seed,
13079                    0,
13080                    &vx.slice((j - 1) * n_embd..j * n_embd),
13081                    n_embd,
13082                )?;
13083                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
13084                // branch); without it the chain entries stand and only the tail truncates. Either
13085                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
13086                // (persistent mode), rope pos+j+1 (chain convention).
13087                if refresh {
13088                    scratch.set_len(e, pos)?;
13089                    let mut vxs = e.zeros(j * n_embd)?;
13090                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13091                    if j > 1 {
13092                        e.copy_view_into(
13093                            &mut vxs,
13094                            n_embd,
13095                            &vx.slice(0..(j - 1) * n_embd),
13096                            (j - 1) * n_embd,
13097                        )?;
13098                    }
13099                    self.mtp_kv_fill_all(
13100                        e,
13101                        &verify_tokens[0..j],
13102                        &vxs,
13103                        pos,
13104                        &mut *scratch,
13105                        embd_dev,
13106                    )?;
13107                } else {
13108                    scratch.set_len(e, pos + j)?;
13109                }
13110                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
13111                // bonus's predecessor (verify col j-1); no pseudo pass.
13112                if !devacc_seeded {
13113                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
13114                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
13115                }
13116                pending = Some(bonus);
13117                if debug_spec {
13118                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
13119                }
13120            } else if !spec_replay {
13121                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
13122                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
13123                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
13124                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
13125                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
13126                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
13127                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
13128                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
13129                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
13130                cache.rollback(e, &snap, 0)?;
13131                scratch.set_len(e, pos)?;
13132                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13133                pending = Some(bonus);
13134                if debug_spec {
13135                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
13136                }
13137            } else {
13138                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
13139                // this round survives, only possible before the first pending exists, ~round 0):
13140                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
13141                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
13142                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
13143                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
13144                // trunk hidden.
13145                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
13146                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
13147                if let Some(b) = pending.take() {
13148                    replay.push(b);
13149                }
13150                replay.extend_from_slice(&draft[0..n_acc]);
13151                replay.push(bonus);
13152                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
13153                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
13154                // last col exactly as before (byte-identical to the old _h_emb_dev call).
13155                let (rl_d, rx) = if self.batched_serving_numeric_class() {
13156                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
13157                    let mut hidden = e.uninit(replay.len() * n_embd)?;
13158                    for (row, &token) in replay.iter().enumerate() {
13159                        let (row_logits, row_hidden) =
13160                            self.spec_target_step_h(e, token, &mut *cache)?;
13161                        logits.extend_from_slice(&row_logits);
13162                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
13163                    }
13164                    (e.htod(&logits)?, hidden)
13165                } else {
13166                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
13167                };
13168                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
13169                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
13170                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
13171                last_pred = e.dtoh_u32(&preds_d)?[0];
13172                if sampled {
13173                    let lr0 = replay.len();
13174                    let lc = last_col_logits
13175                        .as_mut()
13176                        .expect("sampled: last_col_logits unset");
13177                    e.copy_view_into(
13178                        lc,
13179                        0,
13180                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
13181                        n_vocab,
13182                    )?;
13183                }
13184                let lr = replay.len();
13185                if lr >= 2 {
13186                    e.copy_view_into(
13187                        &mut h_seed_buf,
13188                        0,
13189                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
13190                        n_embd,
13191                    )?;
13192                } else {
13193                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
13194                    // last_token, whose own-row hidden fill_prev still holds.
13195                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
13196                }
13197                // the bonus is COMMITTED here — it becomes the last committed row.
13198                let mut rh_last = e.zeros(n_embd)?;
13199                e.copy_view_into(
13200                    &mut rh_last,
13201                    0,
13202                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
13203                    n_embd,
13204                )?;
13205                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
13206                if debug_spec {
13207                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
13208                }
13209            }
13210            if devacc_seeded {
13211                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
13212                // consumed the old value (both slots carry the same value in every non-replay arm).
13213                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
13214            }
13215            if successor_valid {
13216                let optimistic_scratch_len = successor_attempt
13217                    .as_ref()
13218                    .expect("valid controller successor disappeared")
13219                    .scratch_len;
13220                // The normal current-round commit refreshed/truncated the logical scratch tail.
13221                // Its optimistic successor row was already written physically, so restoring only
13222                // the retained logical length makes that row live for the carried round.
13223                scratch.set_len(e, optimistic_scratch_len)?;
13224            }
13225            if let Some(current) = current_opti.take() {
13226                opti_fork
13227                    .as_mut()
13228                    .ok_or("optipipe current retirement lost fork state")?
13229                    .retire(current.generation)?;
13230            }
13231            if successor_valid {
13232                let successor = successor_attempt
13233                    .take()
13234                    .expect("valid controller successor disappeared before promotion");
13235                let generation = successor.generation;
13236                opti_fork
13237                    .as_mut()
13238                    .ok_or("optipipe successor promotion lost fork state")?
13239                    .promote_successor_snapshot(&mut snap, generation);
13240                carried_opti = Some(successor);
13241            }
13242            if anatomy_on {
13243                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
13244                // only for this diagnostic so it does not disappear into the following draft's
13245                // first token readback.
13246                e.stream().synchronize()?;
13247                ph_commit += commit_started.elapsed().as_secs_f64();
13248            }
13249            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
13250            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
13251            // final position — the floor's position key reads the committed depth). Burst
13252            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
13253            // like gemma's burst arm.
13254            if adapt {
13255                let fl_now = floor_at(cache.pos);
13256                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
13257            }
13258            ph_mark(&mut ph_rest, phase_on);
13259            if let Some(p) = pipe {
13260                p.accept_end(round);
13261            }
13262            drop(pipe_accept);
13263            round += 1;
13264            // sse-cadence: this round's accepted drafts + bonus are committed (out is
13265            // append-only past step 4) — flush at round cadence.
13266            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
13267        }
13268        if let Some(mut ticket) = carried_opti.take() {
13269            opti_fork
13270                .as_mut()
13271                .ok_or("optipipe tail drain lost fork state")?
13272                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
13273        }
13274        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
13275        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
13276        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
13277
13278        if spec_stats {
13279            let per_slot: Vec<String> = (0..k)
13280                .map(|j| {
13281                    if st_drafted[j] > 0 {
13282                        format!(
13283                            "{}/{}={:.3}",
13284                            st_accepted[j],
13285                            st_drafted[j],
13286                            st_accepted[j] as f64 / st_drafted[j] as f64
13287                        )
13288                    } else {
13289                        "0/0".into()
13290                    }
13291                })
13292                .collect();
13293            let acc = if total_drafted > 0 {
13294                total_accepted as f64 / total_drafted as f64
13295            } else {
13296                0.0
13297            };
13298            eprintln!(
13299                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
13300                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
13301                       tok_per_round={:.3}",
13302                per_slot.join(" "),
13303                (total_accepted + round) as f64 / round.max(1) as f64
13304            );
13305        }
13306        if constraint.is_some() {
13307            eprintln!(
13308                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
13309                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
13310                dm_clone_ns as f64 / 1e6,
13311                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
13312            );
13313        }
13314        if phase_on {
13315            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
13316            eprintln!(
13317                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
13318                ph_draft * 1e3,
13319                ph_draft / tot * 100.0,
13320                ph_verify * 1e3,
13321                ph_verify / tot * 100.0,
13322                ph_wait * 1e3,
13323                ph_wait / tot * 100.0,
13324                ph_rest * 1e3,
13325                ph_rest / tot * 100.0
13326            );
13327        }
13328        if anatomy_on {
13329            let rounds_f = round.max(1) as f64;
13330            let other = (ph_rest - ph_commit).max(0.0);
13331            eprintln!(
13332                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
13333                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
13334                ph_draft * 1e3 / rounds_f,
13335                ph_verify * 1e3 / rounds_f,
13336                ph_wait * 1e3 / rounds_f,
13337                ph_commit * 1e3 / rounds_f,
13338                other * 1e3 / rounds_f,
13339            );
13340        }
13341        let _pipe_tail = pipe.map(|p| p.primary());
13342        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
13343        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
13344        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
13345        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
13346        if let Some(slot) = sess_draft_slot.take() {
13347            *slot = Some(dctx);
13348        }
13349        let t_rounds = t_ent.elapsed();
13350        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
13351            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
13352            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
13353            // HERE, where the sampler, the session Philox counters and the penalty window are
13354            // all live and the boundary logits row still exists — that is the "make the state
13355            // available" half of the fix; the consuming burst then just emits it. `sctr` is
13356            // written to the session BELOW the draws so the advance is never lost.
13357            *next_pred_slot = Some(last_pred);
13358            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
13359            let mut stashed_pending = false;
13360            if let Some(b) = pending.take() {
13361                if !sampled {
13362                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
13363                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
13364                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
13365                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
13366                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
13367                    // OUT of `committed` (cache rows == committed); the consuming call
13368                    // prepends it once its verify commits the row. next_pred is unknowable
13369                    // without the commit pass — None; callers gate on pending_tok too.
13370                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
13371                    if let Some(slot) = sess_pending_slot.take() {
13372                        *slot = Some(b);
13373                    }
13374                    *next_pred_slot = None;
13375                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
13376                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
13377                    *last_h = Some(e.clone_dtod(&fill_prev)?);
13378                    stashed_pending = true;
13379                } else {
13380                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
13381                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
13382                    let pos_b = cache.pos;
13383                    scratch.set_len(e, pos_b)?;
13384                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
13385                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
13386                    // itself — the prediction AFTER the bonus never materialized; it would have
13387                    // been the next round's verify col 0). The commit's logits ARE that
13388                    // prediction — so they are also the row the next burst's boundary token
13389                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
13390                    *next_pred_slot = Some(if sample_boundary {
13391                        sample_boundary_token(
13392                            e,
13393                            &lg_b,
13394                            &sp,
13395                            &pen_hist,
13396                            &mut sctr,
13397                            "burst-tail-commit",
13398                        )?
13399                    } else {
13400                        argmax(&lg_b) as u32
13401                    });
13402                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
13403                    *last_h = Some(hb);
13404                }
13405            } else {
13406                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
13407                *last_h = Some(e.clone_dtod(&fill_prev)?);
13408                if sample_boundary {
13409                    // No pending to commit, so the boundary row is the one `last_pred` was
13410                    // argmaxed from and the sampled path keeps it on device: the init feed's
13411                    // logits when the burst ran zero rounds, else the legacy-replay path's
13412                    // last verify column (both predict the token AFTER the last committed
13413                    // row). It is retained precisely because round 0's accept test needs it,
13414                    // so the draw costs no extra D2H of the [n_vocab] row.
13415                    match last_col_logits.as_ref() {
13416                        Some(lc) => {
13417                            *next_pred_slot = Some(sample_boundary_token_dev(
13418                                e,
13419                                lc,
13420                                n_vocab,
13421                                &sp,
13422                                &pen_hist,
13423                                &mut sctr,
13424                                "burst-tail-nopending",
13425                            )?);
13426                        }
13427                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
13428                        // burst always feeds or replays, so the row exists — but if it ever
13429                        // is, the stream takes a greedy token and SAYS so rather than
13430                        // silently regressing to the pre-lane behaviour.
13431                        None => eprintln!(
13432                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
13433                             (reason: no retained boundary logits row)"
13434                        ),
13435                    }
13436                }
13437            }
13438            *sctr_slot = sctr;
13439            *uctr_slot = uctr;
13440            committed.extend_from_slice(prompt);
13441            if let Some(cb) = carried_pending {
13442                // the consumed carry's cache row landed in round 0's verify (every pending
13443                // round commits col 0) — it joins `committed` here, in sequence order.
13444                committed.push(cb);
13445            }
13446            if stashed_pending {
13447                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
13448                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
13449                // 18446744073709551615 out of range for slice of length 0", killing the
13450                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
13451                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
13452                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
13453                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
13454                // did). So a burst that stashes a pending without emitting anything of its own —
13455                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
13456                // guard skipping every token under a tight budget — arrives here with
13457                // out.len() == 0 and stashed_pending == true.
13458                //
13459                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
13460                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
13461                // just above is already accounted. Saturating, not a min/assert: an empty `out`
13462                // here is a legitimate burst shape, not a corrupt state.
13463                let emitted = out.len().saturating_sub(1);
13464                committed.extend_from_slice(&out[..emitted]);
13465            } else {
13466                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
13467            }
13468            debug_assert_eq!(
13469                cache.pos,
13470                committed.len(),
13471                "session invariant: cache rows == committed tokens"
13472            );
13473            if setup_trace {
13474                e.stream().synchronize()?; // bound the async tail fill in the trace
13475                let t_tail = t_ent.elapsed();
13476                eprintln!(
13477                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
13478                    t_init.as_secs_f64() * 1e3,
13479                    (t_cap - t_init).as_secs_f64() * 1e3,
13480                    (t_fill - t_cap).as_secs_f64() * 1e3,
13481                    (t_rounds - t_fill).as_secs_f64() * 1e3,
13482                    (t_tail - t_rounds).as_secs_f64() * 1e3,
13483                    t_tail.as_secs_f64() * 1e3,
13484                    out.len(),
13485                    continuation
13486                );
13487            }
13488            return Ok((out, total_drafted, total_accepted));
13489        }
13490        out.truncate(max_new);
13491        Ok((out, total_drafted, total_accepted))
13492    }
13493
13494    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
13495    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
13496    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
13497    pub fn extract_dspark_anchors(
13498        &self,
13499        e: &Engine,
13500        tokens: &[u32],
13501        anchor_positions: &[usize],
13502        gamma: usize,
13503        top_k: usize,
13504        chunk: usize,
13505        temperature: f32,
13506    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
13507        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
13508            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
13509        }
13510        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
13511            return Err("DSpark anchor positions must be sorted and unique".into());
13512        }
13513        for &position in anchor_positions {
13514            if position == 0 || position + gamma >= tokens.len() {
13515                return Err(format!(
13516                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
13517                    tokens.len()
13518                )
13519                .into());
13520            }
13521        }
13522
13523        let n_vocab = self.output.out_features();
13524        let n_embd = self.cfg.n_embd as usize;
13525        let mut cache =
13526            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
13527        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13528        let embd_gpu = if spec_host_embd() {
13529            None
13530        } else {
13531            Some(
13532                self.embd_gpu
13533                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13534            )
13535        };
13536        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
13537
13538        struct PendingRecord {
13539            position: usize,
13540            hidden: Option<Vec<f32>>,
13541            tokens: Vec<u32>,
13542            target_top_ids: Vec<Option<Vec<u32>>>,
13543            target_top_logits: Vec<Option<Vec<f32>>>,
13544            target_top_probs: Vec<Option<Vec<f32>>>,
13545            target_tail_probs: Vec<Option<f32>>,
13546        }
13547
13548        let mut pending: Vec<PendingRecord> = anchor_positions
13549            .iter()
13550            .map(|&position| PendingRecord {
13551                position,
13552                hidden: None,
13553                tokens: tokens[position..=position + gamma].to_vec(),
13554                target_top_ids: vec![None; gamma],
13555                target_top_logits: vec![None; gamma],
13556                target_top_probs: vec![None; gamma],
13557                target_tail_probs: vec![None; gamma],
13558            })
13559            .collect();
13560
13561        let mut start = 0usize;
13562        while start < tokens.len() {
13563            let end = (start + chunk).min(tokens.len());
13564            let chunk_tokens = &tokens[start..end];
13565            let (target_logits, hidden_rows) =
13566                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
13567            for record in &mut pending {
13568                let hidden_position = record.position - 1;
13569                if hidden_position >= start && hidden_position < end {
13570                    let local = hidden_position - start;
13571                    record.hidden = Some(
13572                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
13573                    );
13574                }
13575                for slot in 0..gamma {
13576                    let target_row = record.position + slot;
13577                    if target_row < start || target_row >= end {
13578                        continue;
13579                    }
13580                    let local = target_row - start;
13581                    let logits =
13582                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
13583                    let (ids, top_logits, probs, tail) =
13584                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
13585                    record.target_top_ids[slot] = Some(ids);
13586                    record.target_top_logits[slot] = Some(top_logits);
13587                    record.target_top_probs[slot] = Some(probs);
13588                    record.target_tail_probs[slot] = Some(tail);
13589                }
13590            }
13591            start = end;
13592        }
13593
13594        pending
13595            .into_iter()
13596            .map(|record| {
13597                let hidden = record
13598                    .hidden
13599                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
13600                let target_top_ids =
13601                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
13602                let target_top_logits = flatten_dspark_rows(
13603                    record.target_top_logits,
13604                    record.position,
13605                    "target logits",
13606                )?;
13607                let target_top_probs =
13608                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
13609                let target_tail_probs = record
13610                    .target_tail_probs
13611                    .into_iter()
13612                    .enumerate()
13613                    .map(|(slot, value)| {
13614                        value.ok_or_else(|| {
13615                            format!("missing DSpark tail at {} slot {slot}", record.position)
13616                        })
13617                    })
13618                    .collect::<Result<Vec<_>, _>>()?;
13619                Ok(DsparkAnchorRecord {
13620                    position: record.position,
13621                    hidden,
13622                    tokens: record.tokens,
13623                    target_top_ids,
13624                    target_top_logits,
13625                    target_top_probs,
13626                    target_tail_probs,
13627                })
13628            })
13629            .collect()
13630    }
13631
13632    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
13633    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
13634    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
13635    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
13636    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
13637    /// quant-induced head/hidden-state mismatch from text drift.
13638    ///
13639    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
13640    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
13641    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
13642    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
13643    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
13644    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
13645    ///              conditions on the corpus — deterministic and arm-comparable by design.
13646    ///
13647    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
13648    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
13649    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
13650    ///
13651    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
13652    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
13653    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
13654    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
13655    /// agreement vs this path — not usable as a training-data source).
13656    pub fn replay_acceptance(
13657        &self,
13658        e: &Engine,
13659        tokens: &[u32],
13660        k: usize,
13661        stride: usize,
13662        chunk: usize,
13663        mut hdump: Option<&mut std::fs::File>,
13664    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
13665        assert!(k >= 1 && stride >= 1 && chunk >= 2);
13666        let mtp = self
13667            .mtp
13668            .as_ref()
13669            .expect("replay_acceptance requires an MTP head");
13670        let n_vocab = self.output.out_features();
13671        let d_vocab = mtp
13672            .shared_head_head
13673            .as_ref()
13674            .unwrap_or(&self.output)
13675            .out_features();
13676        let n_embd = self.cfg.n_embd as usize;
13677        let t_total = tokens.len();
13678        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
13679        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
13680        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
13681        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
13682        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
13683        let embd_gpu = if spec_host_embd() {
13684            None
13685        } else {
13686            Some(
13687                self.embd_gpu
13688                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
13689            )
13690        };
13691        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
13692
13693        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
13694        let mut bg: Vec<u32> = vec![0; t_total + 1];
13695        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
13696        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
13697        let mut seed_buf = e.zeros(n_embd)?;
13698        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
13699        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
13700        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
13701        let mut s = 0usize;
13702        while s < t_total {
13703            let cend = (s + chunk).min(t_total);
13704            let tc = cend - s;
13705            let ch = &tokens[s..cend];
13706            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
13707            //    the chunk's true hiddens.
13708            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
13709            for j in 0..tc {
13710                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
13711            }
13712            let preds = e.dtoh_u32(&preds_d)?;
13713            for j in 0..tc {
13714                bg[s + j + 1] = preds[j];
13715            }
13716            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
13717            // checkpoint-quality metric (position j's logits score the GOLD next token).
13718            if nll_on {
13719                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
13720                if jmax > 0 {
13721                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
13722                    let rows: Vec<i32> = (0..jmax as i32).collect();
13723                    let idsd = e.htod_u32_v(&ids)?;
13724                    let rowsd = e.htod_i32(&rows)?;
13725                    let mut outd = e.zeros(jmax)?;
13726                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
13727                    for pr in e.dtoh(&outd)? {
13728                        nll_sum += -((pr.max(1e-30)) as f64).ln();
13729                        nll_cnt += 1;
13730                    }
13731                }
13732            }
13733            if let Some(f) = hdump.as_deref_mut() {
13734                use std::io::Write;
13735                let host: Vec<f32> = e.dtoh(&vx)?;
13736                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
13737                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
13738                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
13739                for v in &host[..tc * n_embd] {
13740                    let b = v.to_bits();
13741                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
13742                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
13743                }
13744                f.write_all(&bytes)?;
13745            }
13746            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
13747            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
13748            // per token saved; the forced trunk pass + hdump is all the mode needs).
13749            let chainless = stride > t_total;
13750            if chainless {
13751                e.copy_view_into(
13752                    &mut prev_last_h,
13753                    0,
13754                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
13755                    n_embd,
13756                )?;
13757                s = cend;
13758                continue;
13759            }
13760            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
13761            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
13762            let mut vxs = e.zeros(tc * n_embd)?;
13763            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
13764            if tc > 1 {
13765                e.copy_view_into(
13766                    &mut vxs,
13767                    n_embd,
13768                    &vx.slice(0..(tc - 1) * n_embd),
13769                    (tc - 1) * n_embd,
13770                )?;
13771            }
13772            scratch.set_len(e, s)?;
13773            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13774            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
13775            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
13776            //    truncates those approximate appends before they can ever be read.
13777            let ps: Vec<usize> = (s..cend)
13778                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
13779                .collect();
13780            for &p in ps.iter().rev() {
13781                scratch.set_len(e, p)?;
13782                if p == s {
13783                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
13784                } else {
13785                    e.copy_view_into(
13786                        &mut seed_buf,
13787                        0,
13788                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
13789                        n_embd,
13790                    )?;
13791                }
13792                let mut e_tok = tokens[p];
13793                let mut d_seed = e.clone_dtod(&seed_buf)?;
13794                let chain_heads = !self.mtp_extra.is_empty();
13795                let mut chain_tokens = if chain_heads {
13796                    vec![tokens[p]]
13797                } else {
13798                    Vec::new()
13799                };
13800                let mut chain_seeds = if chain_heads {
13801                    vec![e.clone_dtod(&seed_buf)?]
13802                } else {
13803                    Vec::new()
13804                };
13805                let mut drafts: Vec<u32> = Vec::with_capacity(k);
13806                for j in 0..k {
13807                    let (dl_d, h_nextn) = if chain_heads {
13808                        self.mtp_chain_forward_dev(
13809                            e,
13810                            &chain_tokens,
13811                            &chain_seeds,
13812                            &mut scratch,
13813                            p,
13814                            embd_dev,
13815                            None,
13816                        )?
13817                    } else {
13818                        self.mtp_head_forward_dev(
13819                            e,
13820                            mtp,
13821                            e_tok,
13822                            &d_seed,
13823                            &mut scratch,
13824                            p + 1 + j,
13825                            embd_dev,
13826                            None,
13827                        )?
13828                    };
13829                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
13830                    let idx = e.dtoh_u32_one(&tok_d)?;
13831                    let d = match &mtp.d2t {
13832                        Some(map) => map[idx as usize],
13833                        None => idx,
13834                    };
13835                    drafts.push(d);
13836                    if chain_heads {
13837                        chain_tokens.push(d);
13838                        chain_seeds.push(h_nextn);
13839                    } else {
13840                        e_tok = d;
13841                        d_seed = h_nextn;
13842                    }
13843                }
13844                // targets may live in a LATER chunk's bg — resolved after the walk.
13845                rows.push((p, drafts, Vec::new()));
13846            }
13847            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
13848            //    expect scratch.len == cend with exact rows).
13849            scratch.set_len(e, s)?;
13850            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
13851            e.copy_view_into(
13852                &mut prev_last_h,
13853                0,
13854                &vx.slice((tc - 1) * n_embd..tc * n_embd),
13855                n_embd,
13856            )?;
13857            s = cend;
13858        }
13859        for (p, drafts, targets) in rows.iter_mut() {
13860            for j in 0..drafts.len() {
13861                targets.push(bg[*p + 1 + j]);
13862            }
13863        }
13864        rows.sort_by_key(|r| r.0);
13865        if nll_cnt > 0 {
13866            let mean = nll_sum / nll_cnt as f64;
13867            println!(
13868                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
13869                mean.exp()
13870            );
13871        }
13872        Ok((rows, bg))
13873    }
13874}
13875
13876#[cfg(test)]
13877mod vg_debt_tests {
13878    use super::dspark_vg_debt_projection;
13879
13880    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
13881    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
13882    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
13883    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
13884    /// impossible must zero the debt.
13885    #[test]
13886    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
13887        const MIB: usize = 1 << 20;
13888        let d = dspark_vg_debt_projection;
13889        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
13890        assert_eq!(d(0, 256, 0, None), 0);
13891        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
13892        assert_eq!(d(10, 0, 500 * MIB, None), 0);
13893        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
13894        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
13895        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
13896
13897        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
13898        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
13899        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
13900
13901        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
13902        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
13903        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
13904        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
13905
13906        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
13907        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
13908        assert_eq!(debt, 250 * (40 * MIB));
13909        assert!(
13910            debt > 3 * (1536 * MIB),
13911            "real growth must dwarf SPEC_SHRINK_RESERVE"
13912        );
13913
13914        // a shrinking/recycled reading never becomes a negative charge.
13915        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
13916        // a stale observation at the same capture count falls back to bootstrap.
13917        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
13918    }
13919}
13920
13921#[cfg(test)]
13922mod mtp_chain_tests {
13923    use super::mtp_chain_head_index;
13924
13925    #[test]
13926    fn embedded_step_heads_cycle_in_declared_order() {
13927        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
13928        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
13929    }
13930
13931    #[test]
13932    fn standalone_draft_remains_single_head() {
13933        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
13934    }
13935}
13936
13937#[cfg(test)]
13938mod tp_verified_prefix_tests {
13939    use super::rewind_tp_kv_verified_prefix;
13940    use crate::tp::ResidentTpKvCache;
13941
13942    fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
13943        let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
13944        let transaction = cache.begin_transaction().unwrap();
13945        let target = cache.append_target(transaction, committed).unwrap();
13946        cache.publish_append(transaction, target).unwrap();
13947        let target = cache.commit_target(transaction, committed).unwrap();
13948        cache.publish_finalize(transaction, target).unwrap();
13949        cache
13950    }
13951
13952    #[test]
13953    fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
13954        let mut layers = vec![Some(cache_with_committed_len(5)), None];
13955        rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
13956        let cache = layers[0].as_ref().unwrap();
13957        assert_eq!(cache.committed_len(), 3);
13958        assert_eq!(cache.staged_len(), 3);
13959    }
13960
13961    #[test]
13962    fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
13963        let mut layers = vec![Some(cache_with_committed_len(1))];
13964        let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
13965            .unwrap_err()
13966            .to_string();
13967        assert!(error.contains("changed shape"), "unexpected error: {error}");
13968    }
13969}
13970
13971#[cfg(test)]
13972mod dspark_sparse_tests {
13973    use super::dspark_sparse_softmax_topk;
13974
13975    #[test]
13976    fn topk_keeps_full_softmax_mass_and_stable_ties() {
13977        let logits = [1.0f32, 3.0, 3.0, -2.0];
13978        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
13979        assert_eq!(ids, vec![1, 2]);
13980        assert_eq!(top_logits, vec![3.0, 3.0]);
13981        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
13982        let expected = 1.0 / denominator;
13983        assert!((probs[0] - expected).abs() < 1.0e-6);
13984        assert!((probs[1] - expected).abs() < 1.0e-6);
13985        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
13986        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
13987    }
13988}
13989
13990#[cfg(test)]
13991mod spec_replay_env_tests {
13992    use super::spec_replay_env_on;
13993
13994    #[test]
13995    fn replay_requires_literal_one() {
13996        assert!(!spec_replay_env_on(None));
13997        assert!(!spec_replay_env_on(Some("")));
13998        assert!(!spec_replay_env_on(Some("0")));
13999        assert!(!spec_replay_env_on(Some("true")));
14000        assert!(!spec_replay_env_on(Some("2")));
14001        assert!(spec_replay_env_on(Some("1")));
14002    }
14003}
14004
14005#[cfg(test)]
14006mod telem_tests {
14007    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
14008
14009    #[test]
14010    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
14011        let counters = SpecTelemetryCounters::default();
14012        for mask in [
14013            [true, true, true],
14014            [true, true, false],
14015            [true, false, false],
14016            [false, false, false],
14017        ] {
14018            let accepted = mask.iter().take_while(|&&value| value).count();
14019            counters.record_round(mask.len(), accepted);
14020        }
14021
14022        let snapshot = counters.snapshot();
14023        assert_eq!(
14024            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
14025            (4, 12, 6)
14026        );
14027        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
14028        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
14029        assert_eq!(snapshot.tau(), 1.5);
14030        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14031        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
14032    }
14033
14034    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
14035    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
14036    #[test]
14037    fn delta_isolates_burst_contribution() {
14038        let mut t = SpecTelemetry::default();
14039        // "previous request": 2 rounds of k=3, accepts 3 then 1.
14040        for (kr, na) in [(3usize, 3usize), (3, 1)] {
14041            t.rounds += 1;
14042            t.drafted += kr as u64;
14043            t.accepted += na as u64;
14044            for j in 0..kr {
14045                t.pos_drafted[j] += 1;
14046            }
14047            for j in 0..na {
14048                t.pos_accepted[j] += 1;
14049            }
14050        }
14051        let before = t;
14052        // "this burst": 1 round k=3, accepts 2.
14053        t.rounds += 1;
14054        t.drafted += 3;
14055        t.accepted += 2;
14056        for j in 0..3 {
14057            t.pos_drafted[j] += 1;
14058        }
14059        for j in 0..2 {
14060            t.pos_accepted[j] += 1;
14061        }
14062        let d = t.delta_since(&before);
14063        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
14064        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
14065        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
14066        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
14067    }
14068
14069    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
14070    /// aggregation invariant.
14071    #[test]
14072    fn merge_accumulates_fieldwise() {
14073        let mut agg = SpecTelemetry::default();
14074        let mut d1 = SpecTelemetry {
14075            rounds: 2,
14076            drafted: 6,
14077            accepted: 4,
14078            ..Default::default()
14079        };
14080        d1.pos_drafted[0] = 2;
14081        d1.pos_accepted[0] = 2;
14082        let mut d2 = SpecTelemetry {
14083            rounds: 1,
14084            drafted: 3,
14085            accepted: 1,
14086            ..Default::default()
14087        };
14088        d2.pos_drafted[0] = 1;
14089        d2.pos_accepted[0] = 1;
14090        d2.pos_drafted[1] = 1;
14091        agg.merge(&d1);
14092        agg.merge(&d2);
14093        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
14094        assert_eq!(agg.pos_drafted[0], 3);
14095        assert_eq!(agg.pos_accepted[0], 3);
14096        assert_eq!(agg.pos_drafted[1], 1);
14097        assert_eq!(agg.pos_accepted[1], 0);
14098    }
14099
14100    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
14101    /// public metrics surface and must never publish a u64-wrapped garbage value.
14102    #[test]
14103    fn delta_saturates_never_wraps() {
14104        let small = SpecTelemetry {
14105            rounds: 1,
14106            drafted: 2,
14107            accepted: 1,
14108            ..Default::default()
14109        };
14110        let big = SpecTelemetry {
14111            rounds: 5,
14112            drafted: 15,
14113            accepted: 9,
14114            ..Default::default()
14115        };
14116        let d = small.delta_since(&big);
14117        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
14118    }
14119}
14120
14121#[cfg(test)]
14122mod opti_fork_tests {
14123    use super::{
14124        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
14125    };
14126
14127    #[test]
14128    fn controller_threshold_and_three_miss_breaker_are_exact() {
14129        let mut policy = OptiControllerPolicy {
14130            threshold: 0.7,
14131            consecutive_misses: 0,
14132            breaker_tripped: false,
14133        };
14134        assert!(!policy.admit(0.699_999));
14135        assert!(policy.admit(0.7));
14136        assert!(!policy.resolve(false));
14137        assert!(!policy.resolve(false));
14138        assert!(policy.resolve(false));
14139        assert!(policy.breaker_tripped);
14140        assert!(!policy.admit(1.0));
14141        assert!(
14142            !policy.resolve(true),
14143            "a resolved hit cannot re-arm a tripped request"
14144        );
14145        assert!(policy.breaker_tripped);
14146    }
14147
14148    #[test]
14149    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
14150        let mut policy = OptiControllerPolicy {
14151            threshold: 0.0,
14152            consecutive_misses: 0,
14153            breaker_tripped: false,
14154        };
14155        for _ in 0..16 {
14156            assert!(policy.admit(0.0));
14157            assert!(!policy.resolve(false));
14158        }
14159        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
14160            assert!(
14161                !policy.admit(invalid),
14162                "invalid q proxy must fail closed: {invalid}"
14163            );
14164        }
14165        assert!(!policy.breaker_tripped);
14166        assert_eq!(policy.consecutive_misses, 0);
14167    }
14168
14169    #[test]
14170    fn alternating_mode_flips_by_generation_not_round_parity() {
14171        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
14172        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
14173        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
14174        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
14175    }
14176
14177    #[test]
14178    fn live_generation_cannot_be_overwritten() {
14179        let mut tracker = OptiForkGenerationTracker::default();
14180        let g0 = tracker.reserve().unwrap();
14181        let g1 = tracker.reserve().unwrap();
14182        let err = tracker.reserve().unwrap_err().to_string();
14183        assert!(
14184            err.contains("still owns generation 0"),
14185            "unexpected error: {err}"
14186        );
14187        tracker.retire(g0).unwrap();
14188        let g2 = tracker.reserve().unwrap();
14189        assert_eq!((g2.id, g2.slot), (2, 0));
14190        tracker.retire(g1).unwrap();
14191        tracker.retire(g2).unwrap();
14192    }
14193
14194    #[test]
14195    fn teardown_rejects_a_stale_generation_tag() {
14196        let mut tracker = OptiForkGenerationTracker::default();
14197        let g0 = tracker.reserve().unwrap();
14198        tracker.retire(g0).unwrap();
14199        let err = tracker.retire(g0).unwrap_err().to_string();
14200        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
14201    }
14202}
14203
14204#[cfg(test)]
14205mod draft_graph_fallback_tests {
14206    use super::DraftGraphFallback;
14207
14208    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
14209    #[test]
14210    fn flip_is_loud_once_and_memoized_after() {
14211        let mut f = DraftGraphFallback::default();
14212        let line = f
14213            .mark_greedy("out of memory")
14214            .expect("first flip must return the warn line");
14215        assert!(
14216            line.contains("WARN"),
14217            "flip line must be warn-level: {line}"
14218        );
14219        assert!(
14220            line.contains("out of memory"),
14221            "flip line must carry the reason: {line}"
14222        );
14223        assert!(f.greedy_failed());
14224        // re-marking an already-failed graph is the memoization: quiet, still failed.
14225        assert!(f.mark_greedy("out of memory").is_none());
14226        assert!(f.greedy_failed());
14227        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
14228        assert!(!f.sampled_failed());
14229        let line_s = f
14230            .mark_sampled("capture unsupported")
14231            .expect("sampled flip is its own flip");
14232        assert!(
14233            line_s.contains("sampled"),
14234            "sampled flip names itself: {line_s}"
14235        );
14236        assert!(f.mark_sampled("capture unsupported").is_none());
14237    }
14238
14239    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
14240    /// and says so exactly when there was something to reset.
14241    #[test]
14242    fn reset_on_resume_clears_flags_and_logs_once() {
14243        let mut f = DraftGraphFallback::default();
14244        // clean session: resume is silent, nothing to reset.
14245        assert!(f.reset_on_resume().is_none());
14246        f.mark_greedy("oom").unwrap();
14247        f.mark_sampled("oom").unwrap();
14248        let note = f
14249            .reset_on_resume()
14250            .expect("a set flag must produce the reset note");
14251        assert!(
14252            note.contains("greedy+sampled"),
14253            "note names what was reset: {note}"
14254        );
14255        assert!(
14256            !f.greedy_failed() && !f.sampled_failed(),
14257            "both flags cleared"
14258        );
14259        // and the NEXT failure after a reset is a fresh flip — loud again.
14260        assert!(f.mark_greedy("oom again").is_some());
14261        let note2 = f.reset_on_resume().expect("greedy-only reset");
14262        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
14263    }
14264
14265    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
14266    /// they precede a fresh capture attempt whose own failure re-flips loudly.
14267    #[test]
14268    fn shape_change_clears_are_silent() {
14269        let mut f = DraftGraphFallback::default();
14270        f.mark_greedy("oom").unwrap();
14271        f.clear_greedy();
14272        assert!(!f.greedy_failed());
14273        f.mark_sampled("oom").unwrap();
14274        f.clear_sampled();
14275        assert!(!f.sampled_failed());
14276        // after a silent clear there is nothing left for resume to report.
14277        assert!(f.reset_on_resume().is_none());
14278    }
14279}
14280
14281/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
14282///
14283/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
14284/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
14285/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
14286/// than remembered.
14287#[cfg(test)]
14288mod sampled_graph_key_tests {
14289    use super::{SampledGraphKey, debug_t_pred0};
14290
14291    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
14292    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
14293        (k.seed, k.temp_bits, k.k)
14294    }
14295
14296    fn pure_temp_key() -> SampledGraphKey {
14297        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
14298        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
14299    }
14300
14301    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
14302    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
14303    #[test]
14304    fn vendor_filters_change_the_key() {
14305        let parked = pure_temp_key();
14306        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
14307        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
14308        assert_eq!(
14309            legacy_key(&parked),
14310            legacy_key(&vendor),
14311            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
14312        );
14313        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
14314        assert!(parked.pure_temp());
14315        assert!(!vendor.pure_temp());
14316    }
14317
14318    /// Each distribution-shaping field alone is enough to drop the parked graph.
14319    #[test]
14320    fn every_filter_field_is_keyed() {
14321        let base = pure_temp_key();
14322        for (what, other) in [
14323            (
14324                "top_k",
14325                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
14326            ),
14327            (
14328                "top_p",
14329                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
14330            ),
14331            (
14332                "min_p",
14333                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
14334            ),
14335            (
14336                "penalties",
14337                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
14338            ),
14339        ] {
14340            assert_ne!(base, other, "{what} must be part of the key");
14341            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
14342            assert_eq!(
14343                legacy_key(&base),
14344                legacy_key(&other),
14345                "{what} was invisible to the pre-fix key",
14346            );
14347        }
14348    }
14349
14350    /// The baked constants stay keyed (this half was always right — regression cover for it).
14351    #[test]
14352    fn baked_constants_stay_keyed() {
14353        let base = pure_temp_key();
14354        assert_ne!(
14355            base,
14356            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
14357            "seed"
14358        );
14359        assert_ne!(
14360            base,
14361            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
14362            "temp"
14363        );
14364        assert_ne!(
14365            base,
14366            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
14367            "k"
14368        );
14369        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
14370        assert_eq!(
14371            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
14372            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
14373        );
14374    }
14375
14376    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
14377    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
14378    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
14379    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
14380    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
14381    ///
14382    /// This test is the other end of that argument, asserted here rather than remembered in a
14383    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
14384    /// would silently become the unsound thing it is documented not to be.
14385    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
14386    #[test]
14387    fn seed_alone_still_rekeys_the_draft_graph() {
14388        let parked = pure_temp_key();
14389        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
14390        assert_ne!(
14391            parked, reseeded,
14392            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
14393             decision not to compare seed rests on exactly this",
14394        );
14395        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
14396        // because of a filter difference.
14397        assert!(parked.pure_temp() && reseeded.pure_temp());
14398    }
14399
14400    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
14401    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
14402    /// agree on the regime, so a graph that survives the drop is legal to launch.
14403    #[test]
14404    fn equal_keys_agree_on_the_regime() {
14405        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14406        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
14407        assert_eq!(a, b);
14408        assert_eq!(a.pure_temp(), b.pure_temp());
14409        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
14410        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
14411        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
14412        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
14413    }
14414
14415    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
14416    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
14417    #[test]
14418    fn debug_print_survives_the_sampled_arm() {
14419        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
14420        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
14421        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
14422        // round 0 without a pending bonus still reports last_pred, in both arms.
14423        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
14424        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
14425        // greedy keeps the real prediction it always printed.
14426        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
14427        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
14428    }
14429}