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/// step35 dcw draft-chain door (lane/step37-draft-graph-20260829; default OFF until the gate
32/// battery banks receipts, per the new-flags law). ON routes the step35 MTP block's draft
33/// attention through the WINDOWED device-counter family (`append_kv_quantized_dcw` +
34/// `fa_decode_dcw`, the step TP graph arc's kernels), which derives the SWA view entirely
35/// from device state (len_d, base_d, window): exactly the view offset the old capture
36/// refusal said `fa_decode_dc` could not express. BOTH draft modes switch together: eager
37/// and captured run the ONE launcher at the ONE bucket (min(cap, window)), so graph-vs-eager
38/// draft parity holds by construction (the `mtp_full_attn_dc` precedent). OFF keeps today's
39/// serving byte-for-byte: the host-len eager arm (`mtp_step35_attn`) plus the named capture
40/// refusal. Rollback seam: unset the env (or set 0) and restart; no state survives.
41fn step35_draft_dcw_on() -> bool {
42    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
43    *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() == Ok("1"))
44}
45
46fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
47    let Some(raw) = value else {
48        return Ok(8);
49    };
50    let width = raw
51        .parse::<usize>()
52        .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
53    if !(2..=8).contains(&width) {
54        return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
55    }
56    Ok(width)
57}
58
59#[cfg(test)]
60mod prime_trows_width_tests {
61    #[test]
62    fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
63        assert_eq!(super::parse_prime_trows_width(None), Ok(8));
64        assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
65        assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
66        for invalid in ["", "1", "9", "32", "wide"] {
67            let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
68            assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
69            assert!(err.contains("2..=8"), "{err}");
70        }
71    }
72}
73
74/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
75/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
76/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
77/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
78/// target arrays are `[gamma, top_k]` in row-major order.
79pub struct DsparkAnchorRecord {
80    pub position: usize,
81    pub hidden: Vec<f32>,
82    pub tokens: Vec<u32>,
83    pub target_top_ids: Vec<u32>,
84    pub target_top_logits: Vec<f32>,
85    pub target_top_probs: Vec<f32>,
86    pub target_tail_probs: Vec<f32>,
87}
88
89fn dspark_sparse_softmax_topk(
90    logits: &[f32],
91    top_k: usize,
92    temperature: f32,
93) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
94    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
95        return Err("invalid DSpark sparse-softmax shape or temperature".into());
96    }
97    if logits.iter().any(|value| !value.is_finite()) {
98        return Err("DSpark target logits contain a non-finite value".into());
99    }
100    let mut ranked: Vec<(u32, f32)> = logits
101        .iter()
102        .copied()
103        .enumerate()
104        .map(|(index, value)| (index as u32, value))
105        .collect();
106    let compare = |left: &(u32, f32), right: &(u32, f32)| {
107        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
108    };
109    ranked.select_nth_unstable_by(top_k - 1, compare);
110    ranked[..top_k].sort_unstable_by(compare);
111
112    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
113    let inv_temperature = 1.0f64 / temperature as f64;
114    let denominator: f64 = logits
115        .iter()
116        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
117        .sum();
118    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
119    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
120    let top_probs: Vec<f32> = top_logits
121        .iter()
122        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
123        .collect();
124    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
125    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
126    Ok((ids, top_logits, top_probs, tail))
127}
128
129fn flatten_dspark_rows<T>(
130    rows: Vec<Option<Vec<T>>>,
131    position: usize,
132    label: &str,
133) -> Result<Vec<T>, Box<dyn std::error::Error>> {
134    let mut flattened = Vec::new();
135    for (slot, row) in rows.into_iter().enumerate() {
136        flattened.extend(
137            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
138        );
139    }
140    Ok(flattened)
141}
142
143/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
144/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
145/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
146/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
147/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
148/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
149/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
150/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
151/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
152/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
153/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
154pub(crate) fn head_rows_on() -> bool {
155    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
156    crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
157}
158
159/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
160/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
161pub(crate) fn spec_verify_eager_on() -> bool {
162    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
163    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
164}
165
166pub(crate) fn spec_verify_tcol_on() -> bool {
167    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
168    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
169}
170
171/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
172/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
173/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
174/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
175/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
176/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
177/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
178/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
179pub(crate) fn prime_trows_on() -> bool {
180    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
181    *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
182}
183
184pub(crate) fn tcol_ffn_on() -> bool {
185    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
186    crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
187}
188
189pub(crate) fn spec_hpost() -> bool {
190    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
191    *H.get_or_init(|| {
192        std::env::var("MEMRA_SPEC_HPOST")
193            .map(|v| v != "0")
194            .unwrap_or(false)
195    })
196}
197
198/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
199/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
200/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
201/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
202/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
203/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
204/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
205/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
206/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
207pub(crate) fn spec_lean() -> bool {
208    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
209    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
210    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
211    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
212    *L.get_or_init(|| {
213        std::env::var("MEMRA_SPEC_LEAN")
214            .map(|v| v != "0")
215            .unwrap_or(true)
216    })
217}
218
219/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
220/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
221/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
222/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
223/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
224/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
225///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
226///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
227///     t-loop == chained T=1 steps);
228/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
229///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
230/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
231pub(crate) fn spec_m2() -> bool {
232    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
233    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
234    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
235    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
236    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
237    *M.get_or_init(|| {
238        std::env::var("MEMRA_SPEC_M2")
239            .map(|v| v != "0")
240            .unwrap_or(true)
241    })
242}
243pub(crate) fn spec_stream() -> bool {
244    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
245    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
246}
247pub(crate) fn spec_stream_m() -> usize {
248    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
249    *M.get_or_init(|| {
250        std::env::var("MEMRA_SPEC_STREAM_M")
251            .ok()
252            .and_then(|v| v.parse().ok())
253            .unwrap_or(4)
254    })
255}
256pub(crate) fn spec_devacc() -> bool {
257    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
258    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
259}
260/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
261/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
262/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
263/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
264/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
265/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
266/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
267/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
268/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
269/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
270pub(crate) fn dspark_defer_readback_on() -> bool {
271    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
272    *ON.get_or_init(|| {
273        std::env::var("MEMRA_DSPARK_DEFER_READBACK")
274            .map(|v| v != "0")
275            .unwrap_or(true)
276    })
277}
278/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
279/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
280/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
281/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
282/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
283/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
284/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
285pub(crate) fn state_copy_batch_on() -> bool {
286    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
287    *ON.get_or_init(|| {
288        std::env::var("MEMRA_STATE_COPY_BATCH")
289            .map(|v| v != "0")
290            .unwrap_or(true)
291    })
292}
293/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
294/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
295/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
296/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
297/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
298///
299/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
300/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
301/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
302/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
303/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
304/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
305/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
306/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
307/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
308/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
309/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
310/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
311/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
312/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
313/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
314/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
315/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
316/// ratification on the serve-surface battery.
317pub(crate) fn dspark_verify_graph_on() -> bool {
318    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
319    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
320}
321/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
322/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
323///
324/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
325/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
326/// on this route. The MTP spec round is that caller.
327///
328/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
329/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
330/// the host is never waiting for the device, it is spending its own time launching the trunk.
331/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
332/// 8-10 ms per burst).
333///
334/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
335///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
336///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
337///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
338///     comes from per-round phase totals, which are internal to each boot).
339/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
340/// the round off the host and onto the device, which is the whole point.
341///
342/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
343/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
344/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
345/// at every K, kernel-check ALL GREEN.
346///
347/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
348/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
349/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
350/// opt in with `=1` once it has its own interleave. Also never armed together with
351/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
352pub(crate) fn spec_verify_graph_env() -> Option<bool> {
353    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
354    *ON.get_or_init(
355        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
356            Ok("1") => Some(true),
357            Ok("0") => Some(false),
358            _ => None,
359        },
360    )
361}
362/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
363/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
364/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
365/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
366/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
367/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
368/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
369/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
370/// 256-token run vs the serve session's thousands of rounds), and the two
371/// instruments must keep their own measured dispositions rather than share one flag.
372pub(crate) fn dspark_verify_graph_serve_on() -> bool {
373    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
374    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
375}
376/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
377/// pool's memory policy STATED instead of silently unbounded. The keyspace is
378/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
379/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
380/// on the q38 export — so the default (256) never engages there; the knob is the
381/// safety valve for a future export with a wider ladder. At the ceiling the pool
382/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
383/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
384/// cols-stashed layers inside one commit). No eviction by design: destroying a live
385/// exec graph re-opens the stale-address class the indirect tables exist to close,
386/// and the bounded keyspace makes reclaim worthless.
387pub(crate) fn dspark_vg_cap() -> usize {
388    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
389    *CAP.get_or_init(|| {
390        std::env::var("MEMRA_DSPARK_VG_MAX")
391            .ok()
392            .and_then(|v| v.parse().ok())
393            .unwrap_or(256)
394    })
395}
396
397/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
398/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
399/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
400/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
401/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
402/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
403///
404/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
405/// and proves nothing about another export): the debt is remaining capture slots x the
406/// MARGINAL bytes a capture adds to this device's graph mem pool.
407///
408/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
409/// version of this used the mean (`reserved / captures`) and the live serve log showed why
410/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
411/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
412/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
413/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
414/// boot can refuse admissions that would have fit, which is a worse defect than the
415/// under-charge this accounting exists to remove. The marginal reading prices what an
416/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
417/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
418/// tracks real growth on one that does.
419///
420/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
421/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
422/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
423/// the same direction as the old rule without the 255x extrapolation.
424///
425/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
426/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
427/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
428/// debt is 0 there too.
429pub fn dspark_vg_debt_projection(
430    captures: usize,
431    cap: usize,
432    reserved_bytes: usize,
433    prev: Option<(usize, usize)>,
434) -> usize {
435    if captures == 0 || cap == 0 {
436        return 0;
437    }
438    let remaining = cap.saturating_sub(captures);
439    if remaining == 0 {
440        return 0;
441    }
442    match prev {
443        // marginal growth between two observations of the same pool
444        Some((c0, r0)) if captures > c0 => {
445            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
446            remaining.saturating_mul(marginal)
447        }
448        // bootstrap: at most one more pool's worth
449        _ => remaining
450            .saturating_mul(reserved_bytes / captures)
451            .min(reserved_bytes),
452    }
453}
454/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
455/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
456/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
457/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
458/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
459/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
460/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
461/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
462/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
463/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
464/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
465/// empty partial the combine never reads, so the shared n_splits_max stride changes no
466/// bytes) and re-gated e2e by this lane's battery.
467pub(crate) fn dspark_fa_rows_on() -> bool {
468    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
469    *ON.get_or_init(|| {
470        std::env::var("MEMRA_DSPARK_FA_ROWS")
471            .map(|v| v != "0")
472            .unwrap_or(true)
473    })
474}
475
476/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
477///
478/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
479/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
480/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
481/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
482/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
483/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
484/// the flag crashed precisely the regime it exists to investigate.
485///
486/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
487/// indexing (an out-of-range pred there is a real bug and must still be loud).
488fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
489    if base == 0 {
490        return last_pred.to_string();
491    }
492    match preds.get(base - 1) {
493        Some(p) => p.to_string(),
494        // sampled: the greedy per-column argmax was never run for this round.
495        None => {
496            debug_assert!(
497                sampled,
498                "greedy spec: preds[{}] missing at base {base}",
499                base - 1
500            );
501            "n/a".to_string()
502        }
503    }
504}
505
506/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
507///
508/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
509/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
510/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
511/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
512/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
513/// not believe in — and `u * 0 < p` then accepts it unconditionally.
514///
515/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
516/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
517pub(crate) fn skey_probe() -> bool {
518    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
519    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
520}
521
522/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
523/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
524/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
525/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
526/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
527/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
528/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
529/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
530/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
531pub trait SpecConstraint {
532    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
533    /// masked argmax).
534    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
535    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
536    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
537    /// Is `tok` consumable in the CURRENT state?
538    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
539    /// Advance the state with an emitted token.
540    fn consume(&mut self, tok: u32) -> Result<(), String>;
541
542    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
543    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
544    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
545    // loose, research/constrained-full-20260803). These three methods let the engine mask the
546    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
547    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
548    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
549    // stays the correctness backstop and the emitted stream is unchanged by construction
550    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
551    // argmax; a cut slot is recomputed as the masked argmax either way).
552    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
553
554    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
555    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
556    fn draft_mask_enabled(&self) -> bool {
557        false
558    }
559    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
560    /// slot. Called once per spec round, before the first draft position.
561    fn draft_begin(&mut self) -> Result<(), String> {
562        Ok(())
563    }
564    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
565    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
566    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
567        Ok(None)
568    }
569    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
570    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
571    /// engine stops drafting; the token already pushed still goes through verify.
572    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
573        Ok(false)
574    }
575}
576
577/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
578/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
579/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
580/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
581/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
582/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
583/// verify emits the masked argmax as usual).
584fn upload_draft_mask(
585    e: &Engine,
586    c: &mut dyn SpecConstraint,
587    dst: &mut CudaSlice<u32>,
588    d2t: Option<&Vec<u32>>,
589    d_vocab: usize,
590    words: usize,
591) -> Result<bool, Box<dyn std::error::Error>> {
592    let Some(tw) = c
593        .draft_mask_words()
594        .map_err(|e2| format!("constraint: {e2}"))?
595    else {
596        return Ok(false);
597    };
598    let bit = |t: usize| -> bool {
599        let w = t >> 5;
600        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
601    };
602    let mut buf = vec![0u32; words];
603    match d2t {
604        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
605        Some(map) => {
606            for (i, &t) in map.iter().enumerate().take(d_vocab) {
607                if bit(t as usize) {
608                    buf[i >> 5] |= 1u32 << (i & 31);
609                }
610            }
611        }
612        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
613        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
614        None => {
615            let n = tw.len().min(words);
616            buf[..n].copy_from_slice(&tw[..n]);
617        }
618    }
619    if buf.iter().all(|w| *w == 0) {
620        return Ok(false);
621    }
622    e.htod_u32_into(dst, &buf)?;
623    Ok(true)
624}
625
626/// Keep the full token-embedding table in host memory and upload only the rows needed by each
627/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
628/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
629/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
630pub(crate) fn spec_host_embd() -> bool {
631    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
632    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
633}
634
635/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
636/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
637/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
638/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
639/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
640/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
641/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
642/// run-spec K=1..8 + acceptance identity arbitrate e2e).
643pub(crate) fn spec_fused_t() -> bool {
644    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
645    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
646    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
647    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
648    *F.get_or_init(|| {
649        std::env::var("MEMRA_SPEC_FUSED_T")
650            .map(|v| v != "0")
651            .unwrap_or(true)
652    })
653}
654
655/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
656/// Only call this on such buffers — the lean contract is "identical bytes by construction".
657/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
658///
659/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
660/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
661/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
662/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
663/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
664/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
665/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
666/// the vendor-default sampled shape we actually serve.
667pub(crate) fn guard_vocab_token(
668    tok: u32,
669    n_vocab: usize,
670    what: &str,
671) -> Result<u32, Box<dyn std::error::Error>> {
672    if (tok as usize) >= n_vocab {
673        return Err(format!(
674            "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
675             the device argmax's init sentinel in place; refusing to dereference the embed \
676             row (#87 trap)"
677        )
678        .into());
679    }
680    Ok(tok)
681}
682
683/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
684///
685/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
686/// head but not where it entered. With the scan armed the verify walk syncs and reads back
687/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
688/// round's row and position. Off by default and never on a serving path: it costs one host
689/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
690/// reproducing under the scan is itself a datum, not an all-clear).
691///
692/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
693/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
694pub(crate) fn spec_nan_scan() -> bool {
695    spec_nan_scan_level() > 0
696}
697
698/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
699/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
700/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
701/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
702/// MoE produced it, and those are different bugs with different fixes.
703pub(crate) fn spec_nan_scan_level() -> u8 {
704    static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
705    *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
706        Ok("1") => 1,
707        Ok("2") => 2,
708        _ => 0,
709    })
710}
711
712/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
713/// producer (layer index, walk arm) so the error line is the localization.
714/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
715///
716/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
717/// because the level-1 residual scan below sat only on the non-fused tail: the fused
718/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
719/// silently read as "clean". A poisoned residual therefore first reported at the next
720/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
721/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
722/// ran" is distinguishable from "it ran and was innocent".
723/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
724///
725/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
726/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
727/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
728/// head of each block is host-checkable straight out of the byte plane.
729///
730/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
731/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
732/// implicates the shared KV history those rows walk, not per-column staging. "The attention
733/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
734/// different owners, and nothing measured so far separates them. A first-corrupt-row index
735/// also dates the corruption against the prime/decode boundary.
736///
737/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
738/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
739/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
740/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
741pub(crate) fn kv_plane_scan_on() -> bool {
742    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
743    *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
744}
745
746fn kv_plane_scan_rounds() -> usize {
747    static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
748    *R.get_or_init(|| {
749        std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
750            .ok()
751            .and_then(|v| v.parse().ok())
752            .unwrap_or(2)
753    })
754}
755
756/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
757/// scale every `stride` bytes. Returns None when every block scale is finite.
758fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
759    if stride == 0 {
760        return None;
761    }
762    for (i, blk) in bytes.chunks_exact(stride).enumerate() {
763        let raw = u16::from_le_bytes([blk[0], blk[1]]);
764        if half_is_non_finite(raw) {
765            return Some((i, raw));
766        }
767    }
768    None
769}
770
771/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
772fn half_is_non_finite(raw: u16) -> bool {
773    (raw & 0x7C00) == 0x7C00
774}
775
776/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
777/// receipt line, or None when the layer is out of scope or every scale is finite.
778pub(crate) fn scan_kv_plane(
779    e: &crate::Engine,
780    distributed: &memra_kv::ResidentTpKvCache,
781    il: usize,
782    pos0: usize,
783) -> Result<(), Box<dyn std::error::Error>> {
784    // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
785    // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
786    // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
787    // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
788    // be poisoned and report a clean history it never looked at.
789    static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
790    static LAST_POS: std::sync::atomic::AtomicUsize =
791        std::sync::atomic::AtomicUsize::new(usize::MAX);
792    if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
793        ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
794    }
795    if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
796        return Ok(());
797    }
798    let staged = distributed.staged_len();
799    if staged == 0 {
800        return Ok(());
801    }
802    // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
803    // read the same whether the history was clean or the scan never ran once. Bounded so a
804    // 45-layer walk cannot flood the log.
805    static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
806    let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
807    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
808    if seen < 4 {
809        eprintln!(
810            "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
811             ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
812        );
813    }
814    for rank in 0..distributed.ranks().len() {
815        let Some(rc) = distributed.rank(rank) else {
816            continue;
817        };
818        // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
819        let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
820        let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
821        let kbad = first_bad_scale(&kbytes, 34);
822        let vbad = first_bad_scale(&vbytes, 24);
823        if kbad.is_some() || vbad.is_some() {
824            let row = |b: Option<(usize, u16)>, tok: usize| {
825                b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
826                    .unwrap_or_else(|| "clean".into())
827            };
828            eprintln!(
829                "[kv-plane] layer {il} rank {rank} pos0={pos0} staged={staged}                  K={} V={} - the attended KV history is ALREADY non-finite, so a non-finite                  attention output here is a symptom and not the origin",
830                row(kbad, ktb),
831                row(vbad, vtb)
832            );
833            return Ok(());
834        }
835    }
836    Ok(())
837}
838
839pub(crate) fn verify_arm_receipt(
840    arm: &str,
841    il: usize,
842    pos0: usize,
843    t: usize,
844    staged: Option<usize>,
845) {
846    static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
847    if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
848        return;
849    }
850    eprintln!(
851        "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
852        staged.map(|v| v as i64).unwrap_or(-1),
853        crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
854    );
855}
856
857pub(crate) fn nan_scan_rows(
858    e: &Engine,
859    buf: &CudaSlice<f32>,
860    rows: usize,
861    cols: usize,
862    what: &str,
863) -> Result<(), Box<dyn std::error::Error>> {
864    // The readback is also the ATTRIBUTION point for an asynchronous fault: a
865    // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
866    // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
867    // died somewhere" into "it died at or before this layer, on this row, at this position".
868    let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
869        format!(
870            "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
871                     this point in the walk"
872        )
873        .into()
874    })?;
875    if host.len() < rows * cols {
876        return Err(format!(
877            "nan-scan {what}: buffer holds {} < {rows}x{cols}",
878            host.len()
879        )
880        .into());
881    }
882    // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
883    // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
884    // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
885    // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
886    // per-column staging can appear in one. Report the whole map.
887    let mut per_row: Vec<usize> = Vec::with_capacity(rows);
888    let mut first_bad: Option<(usize, usize)> = None;
889    for r in 0..rows {
890        let row = &host[r * cols..(r + 1) * cols];
891        let bad = row.iter().filter(|v| !v.is_finite()).count();
892        per_row.push(bad);
893        if bad > 0 && first_bad.is_none() {
894            first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
895        }
896    }
897    if let Some((r0, c0)) = first_bad {
898        let map: String = per_row
899            .iter()
900            .map(|&b| if b == 0 { '.' } else { 'X' })
901            .collect();
902        return Err(format!(
903            "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
904             counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
905             implicates shared state (the KV history this layer reads); one row bad implicates \
906             per-column staging."
907        )
908        .into());
909    }
910    Ok(())
911}
912
913fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
914    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
915}
916
917/// Scratch KV for the MTP block (one full-attn layer).
918///
919/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
920/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
921/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
922/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
923/// engine's "mtp_update" design). Entries come from two sources:
924///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
925///     hidden chain-approximate — the reference engine accepts the same);
926///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
927///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
928/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
929/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
930/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
931/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
932/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
933/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
934/// committed row across turns (the predecessor-pairing seed + fill anchor).
935/// Per-request sampling config for the sampled-spec serve path.
936#[derive(Clone, Copy, Debug)]
937pub struct SpecSampling {
938    pub temp: f32,
939    pub seed: u64,
940    pub top_k: i32,            // 0 = off
941    pub top_p: f32,            // 1.0 = off
942    pub min_p: f32,            // 0.0 = off
943    pub penalty_last_n: usize, // 0 = penalties off
944    pub penalty_repeat: f32,
945    pub penalty_freq: f32,
946    pub penalty_present: f32,
947}
948
949impl SpecSampling {
950    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
951    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
952    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
953    /// key their penalty arms off this.
954    pub fn pen_on(&self) -> bool {
955        self.penalty_last_n > 0
956            && (self.penalty_repeat != 1.0
957                || self.penalty_freq != 0.0
958                || self.penalty_present != 0.0)
959    }
960}
961
962/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
963/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
964/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
965/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
966/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
967/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
968/// is a distributional bug, not a style problem).
969pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
970    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
971    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
972    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
973    for _ in 0..10 {
974        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
975        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
976        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
977        c0 = n0;
978        c1 = n1;
979        c2 = n2;
980        c3 = n3;
981        k0 = k0.wrapping_add(0x9E3779B9);
982        k1 = k1.wrapping_add(0xBB67AE85);
983    }
984    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
985}
986
987/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
988/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
989pub const SPEC_TELEM_POS: usize = 8;
990
991/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
992/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
993/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
994/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
995/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
996/// in NEITHER drafted nor accepted.
997#[derive(Clone, Copy, Default, Debug)]
998pub struct SpecTelemetry {
999    /// verify rounds completed (a round-stream burst counts each of its M rounds).
1000    pub rounds: u64,
1001    /// tokens drafted / accepted across all rounds.
1002    pub drafted: u64,
1003    pub accepted: u64,
1004    /// how often draft position j (0-based within a round's chain) was offered / accepted.
1005    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1006    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1007    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1008    pub pos_drafted: [u64; SPEC_TELEM_POS],
1009    pub pos_accepted: [u64; SPEC_TELEM_POS],
1010}
1011
1012impl SpecTelemetry {
1013    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1014    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1015    /// a wrapped counter.
1016    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1017        let mut d = SpecTelemetry {
1018            rounds: self.rounds.saturating_sub(prev.rounds),
1019            drafted: self.drafted.saturating_sub(prev.drafted),
1020            accepted: self.accepted.saturating_sub(prev.accepted),
1021            ..Default::default()
1022        };
1023        for j in 0..SPEC_TELEM_POS {
1024            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1025            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1026        }
1027        d
1028    }
1029    /// Fieldwise `self += d` — the worker's per-model aggregation.
1030    pub fn merge(&mut self, d: &SpecTelemetry) {
1031        self.rounds += d.rounds;
1032        self.drafted += d.drafted;
1033        self.accepted += d.accepted;
1034        for j in 0..SPEC_TELEM_POS {
1035            self.pos_drafted[j] += d.pos_drafted[j];
1036            self.pos_accepted[j] += d.pos_accepted[j];
1037        }
1038    }
1039
1040    /// Mean accepted draft-prefix length per verify round (tau).
1041    pub fn tau(&self) -> f64 {
1042        if self.rounds > 0 {
1043            self.accepted as f64 / self.rounds as f64
1044        } else {
1045            0.0
1046        }
1047    }
1048}
1049
1050/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1051/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1052/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1053struct SpecTelemetryCounters {
1054    rounds: AtomicU64,
1055    drafted: AtomicU64,
1056    accepted: AtomicU64,
1057    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1058    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1059}
1060
1061impl Default for SpecTelemetryCounters {
1062    fn default() -> Self {
1063        Self {
1064            rounds: AtomicU64::new(0),
1065            drafted: AtomicU64::new(0),
1066            accepted: AtomicU64::new(0),
1067            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1068            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1069        }
1070    }
1071}
1072
1073impl SpecTelemetryCounters {
1074    fn record_round(&self, drafted: usize, accepted: usize) {
1075        debug_assert!(accepted <= drafted);
1076        self.rounds.fetch_add(1, Ordering::Relaxed);
1077        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1078        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1079        for counter in self.pos_drafted.iter().take(drafted) {
1080            counter.fetch_add(1, Ordering::Relaxed);
1081        }
1082        for counter in self.pos_accepted.iter().take(accepted) {
1083            counter.fetch_add(1, Ordering::Relaxed);
1084        }
1085    }
1086
1087    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1088    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1089    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1090        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1091        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1092        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1093    }
1094
1095    fn snapshot(&self) -> SpecTelemetry {
1096        SpecTelemetry {
1097            rounds: self.rounds.load(Ordering::Relaxed),
1098            drafted: self.drafted.load(Ordering::Relaxed),
1099            accepted: self.accepted.load(Ordering::Relaxed),
1100            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1101            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1102        }
1103    }
1104}
1105
1106pub struct SpecSession {
1107    pub(crate) cache: Cache,
1108    pub(crate) scratch: MtpScratch,
1109    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1110    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1111    /// session must count them. Callers render output from this, not from their own echo.
1112    pub committed: Vec<u32>,
1113    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1114    pub(crate) last_h: Option<CudaSlice<f32>>,
1115    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1116    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1117    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1118    pub next_pred: Option<u32>,
1119    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1120    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1121    pub sctr: u32,
1122    pub uctr: u32,
1123    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1124    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1125    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1126    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1127    /// research/spec-serving-20260801). None before the first turn; error paths drop it
1128    /// (next burst recaptures — serve retires errored sessions anyway).
1129    pub(crate) draft_ctx: Option<DraftGraphCtx>,
1130    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1131    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1132    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1133    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1134    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1135    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1136    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1137    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1138    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1139    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1140    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1141    pub pending_tok: Option<u32>,
1142    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1143    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1144    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1145    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1146    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1147    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1148    /// accounting the loop already does — no syncs, no allocation. NOTE a
1149    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1150    /// diff with [`SpecTelemetry::delta_since`] around each burst.
1151    telem: SpecTelemetryCounters,
1152    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1153    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1154    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1155    /// prime, result lands in `boundary_captures`.
1156    pub capture_at: Option<usize>,
1157    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1158    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1159    /// publication just isn't available for that request. Plural since
1160    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1161    /// split (the shared-prefix class) and the stable pre-generation boundary (the
1162    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1163    /// prefill tick publishes/checkpoints.
1164    pub boundary_captures: Vec<SpecBoundaryCapture>,
1165    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1166    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1167    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1168    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1169    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1170    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1171    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1172    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1173    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1174    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1175    /// prompt-end capture.
1176    pub ckpt_at: Option<usize>,
1177}
1178impl SpecSession {
1179    /// Context capacity of the session's caches (the server's ContextFull guard).
1180    pub fn cache_max_ctx(&self) -> usize {
1181        self.cache.max_ctx
1182    }
1183    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1184    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1185    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1186    /// the prime boundary), so no copy was taken at prime time.
1187    pub fn cache_ref(&self) -> &Cache {
1188        &self.cache
1189    }
1190    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1191    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1192    /// like the trunk KV — draft rows below the prompt end are append-only for the
1193    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1194    /// committed length, never below the prime boundary, and the true-hidden refresh
1195    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1196    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1197    /// prefix-addressable; the prefix cache already refuses that class end to end).
1198    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1199        if self.scratch.kv.ring.is_some() {
1200            return None;
1201        }
1202        Some((
1203            &self.scratch.kv.k,
1204            &self.scratch.kv.v,
1205            self.scratch.kv.k_tok_bytes,
1206            self.scratch.kv.v_tok_bytes,
1207        ))
1208    }
1209    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1210    pub fn telemetry(&self) -> SpecTelemetry {
1211        self.telem.snapshot()
1212    }
1213    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1214    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1215    /// `spec_rewind_to_checkpoint`.
1216    pub fn rewind_pos(&self) -> Option<usize> {
1217        self.turn_ckpt.as_ref().map(|c| c.pos)
1218    }
1219    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1220    pub fn rewind_is_resident(&self) -> bool {
1221        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1222            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1223        })
1224    }
1225    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1226    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1227    /// session has never run a turn and has no prediction to hand over.
1228    pub fn demote_ready(&self) -> bool {
1229        self.pending_tok.is_none() && self.next_pred.is_some()
1230    }
1231    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1232    pub fn has_pending(&self) -> bool {
1233        self.pending_tok.is_some()
1234    }
1235    /// Committed row count == cache rows (the session invariant), for the caller's own
1236    /// `fed`-length cross-check at a handoff boundary.
1237    pub fn committed_len(&self) -> usize {
1238        self.committed.len()
1239    }
1240    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1241    /// cache + next-token prediction to the plain batched-decode path.
1242    ///
1243    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1244    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1245    /// tokenwise prime of the same `committed` sequence would have left it (that is the
1246    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1247    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1248    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1249    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1250    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1251    /// a state indistinguishable from one the batched path produced itself: the batched tick
1252    /// emits `next_pred`, feeds it into this same cache, and decodes on.
1253    ///
1254    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1255    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1256    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1257    /// path would silently skip a token.
1258    ///
1259    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1260    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1261    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1262    /// would mean an `mtp_kv_fill` over the whole committed history).
1263    pub fn into_demoted(self) -> Option<(Cache, u32)> {
1264        if self.pending_tok.is_some() {
1265            return None;
1266        }
1267        let np = self.next_pred?;
1268        debug_assert_eq!(
1269            self.cache.pos,
1270            self.committed.len(),
1271            "demotion handoff: cache rows != committed tokens"
1272        );
1273        Some((self.cache, np))
1274    }
1275    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1276    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1277    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1278    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1279    pub fn reset_graph_fallback_on_resume(&mut self) {
1280        if let Some(line) = self
1281            .draft_ctx
1282            .as_mut()
1283            .and_then(|c| c.failed.reset_on_resume())
1284        {
1285            eprintln!("{line}");
1286        }
1287    }
1288}
1289
1290/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1291///
1292/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1293/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1294/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1295/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1296/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1297/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1298///
1299/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1300/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1301/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1302/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1303/// below the boundary were written by this turn's fill and are never revisited (the per-round
1304/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1305/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1306/// predecessor-pairing anchor the next prime's fill reads for its first row.
1307///
1308/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1309pub(crate) struct SpecCheckpoint {
1310    snap: crate::cache::CacheSnapshot,
1311    /// Committed length at the boundary (== cache.pos there, the session invariant).
1312    pos: usize,
1313    /// Pre-output_norm hidden of row `pos - 1`.
1314    last_h: CudaSlice<f32>,
1315}
1316
1317/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1318/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1319/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1320/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1321/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1322/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1323/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1324/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1325pub struct SpecBoundaryCapture {
1326    pub snap: crate::cache::CacheSnapshot,
1327    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1328    pub pos: usize,
1329    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1330    pub logits: Vec<f32>,
1331    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1332    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1333    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1334    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1335    pub last_h: Vec<f32>,
1336}
1337
1338/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1339/// spec boundary capture carries for later restored-session fills. Failure is silent
1340/// (`turn_ckpt` convention): the capture publishes without an anchor.
1341fn capture_boundary_hidden(
1342    e: &Engine,
1343    h_rows: &CudaSlice<f32>,
1344    pos: usize,
1345    n_embd: usize,
1346) -> Vec<f32> {
1347    if pos == 0 || h_rows.len() < pos * n_embd {
1348        return Vec::new();
1349    }
1350    let Ok(mut row) = e.uninit(n_embd) else {
1351        return Vec::new();
1352    };
1353    if e.copy_view_into(
1354        &mut row,
1355        0,
1356        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1357        n_embd,
1358    )
1359    .is_err()
1360    {
1361        return Vec::new();
1362    }
1363    e.dtoh(&row).unwrap_or_default()
1364}
1365
1366/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1367/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1368/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1369/// every boundary) without touching greedy, which is byte-unaffected either way.
1370pub fn spec_sampled_boundary_on() -> bool {
1371    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1372    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1373}
1374
1375/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1376/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1377/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1378/// restores the pre-lane posture (each burst restarts the window from its own prompt
1379/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1380/// must keep refusing penalized sampled prefix-cache restores, because the restored
1381/// session's continuation burst is handed no prompt slice at all.
1382pub fn spec_pen_session_on() -> bool {
1383    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1384    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1385}
1386
1387/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1388/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1389/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1390/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1391/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1392/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1393pub fn spec_restore_republish_on() -> bool {
1394    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1395    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1396}
1397
1398/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1399/// the argmax the pre-lane code would have emitted from the same row. This is how the
1400/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1401fn spec_boundary_trace() -> bool {
1402    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1403    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1404}
1405
1406/// llama-parity floor for the penalty window when the request does not ask for a bigger
1407/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1408/// non-identity penalty, so this floor only matters to explicit small windows and to the
1409/// CLI env path.
1410const PEN_WINDOW_FLOOR: usize = 64;
1411
1412/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1413/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1414/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1415/// p column, the bonus column). The serve API uses this same bound for every non-identity
1416/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1417/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1418/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1419/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1420/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1421/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1422/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1423/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1424/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1425/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1426/// is a second thing to drift.
1427pub const PEN_WINDOW_MAX: usize = 8192;
1428
1429/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1430/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1431/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1432/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1433/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1434/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1435/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1436/// window through the SAME function (one definition of "the window" across both spec
1437/// routes and the gate binary's trunk-only reference arm).
1438pub fn pen_window_seed(
1439    session_committed: &[u32],
1440    burst_prompt: &[u32],
1441    penalty_last_n: usize,
1442) -> Vec<u32> {
1443    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1444    let take_prompt = burst_prompt.len().min(win);
1445    let take_sess = (win - take_prompt).min(session_committed.len());
1446    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1447    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1448    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1449    hist
1450}
1451
1452/// Draw a BOUNDARY token from the target distribution the request asked for
1453/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1454/// every burst boundary".
1455///
1456/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1457/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1458/// row after the last committed token on a continuation burst; the prefix-cache entry's
1459/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1460/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1461/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1462/// customer asked for a sampled token, so this draws one.
1463///
1464/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1465/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1466/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1467/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1468/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1469/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1470///
1471/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1472/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1473/// stream the accept walk uses — never a second, independently seeded stream (which would be
1474/// a new distributional bug: two streams from one seed correlate wherever their counters
1475/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1476/// to the cold session's own first draw from the same logits row, which is what preserves the
1477/// sampled-hit lane's per-seed hit==cold byte identity.
1478#[allow(clippy::too_many_arguments)]
1479pub fn sample_boundary_token_dev(
1480    e: &Engine,
1481    logits: &CudaSlice<f32>,
1482    n_vocab: usize,
1483    sp: &SpecSampling,
1484    pen_hist: &[u32],
1485    sctr: &mut u32,
1486    site: &str,
1487) -> Result<u32, Box<dyn std::error::Error>> {
1488    debug_assert!(
1489        sp.temp > 0.0,
1490        "boundary sampling is the sampled regime only"
1491    );
1492    // Own copy: penalize_logits mutates in place and the caller's row is live state
1493    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1494    let mut col = e.zeros(n_vocab)?;
1495    e.copy_into(&mut col, 0, logits, n_vocab)?;
1496    let pen_on = sp.penalty_last_n > 0
1497        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1498    if pen_on && !pen_hist.is_empty() {
1499        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1500        let w0 = pen_hist
1501            .len()
1502            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1503        let hist = &pen_hist[w0..];
1504        let hd = e.htod_u32_v(hist)?;
1505        e.penalize_logits(
1506            &mut col,
1507            &hd,
1508            hist.len(),
1509            sp.penalty_repeat,
1510            sp.penalty_freq,
1511            sp.penalty_present,
1512            n_vocab,
1513        )?;
1514    }
1515    let rows0 = e.htod_i32(&[0])?;
1516    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1517    e.filter_stats(
1518        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1519        sp.top_p, sp.min_p,
1520    )?;
1521    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1522    let mut perturb = e.zeros(n_vocab)?;
1523    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1524    *sctr = sctr.wrapping_add(1);
1525    let td = e.argmax_token_device(&perturb, n_vocab)?;
1526    let tok = guard_vocab_token(
1527        e.dtoh_u32_one(&td)?,
1528        n_vocab,
1529        &format!("sampled boundary token (site={site})"),
1530    )?;
1531    if spec_boundary_trace() {
1532        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1533        let raw = e.argmax_token_device(logits, n_vocab)?;
1534        let greedy = e.dtoh_u32_one(&raw)?;
1535        eprintln!(
1536            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1537             deviates={} temp={} sctr={}",
1538            (tok != greedy) as u8,
1539            sp.temp,
1540            sctr.wrapping_sub(1),
1541        );
1542    }
1543    Ok(tok)
1544}
1545
1546/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1547/// host `Vec<f32>`).
1548#[allow(clippy::too_many_arguments)]
1549pub fn sample_boundary_token(
1550    e: &Engine,
1551    logits: &[f32],
1552    sp: &SpecSampling,
1553    pen_hist: &[u32],
1554    sctr: &mut u32,
1555    site: &str,
1556) -> Result<u32, Box<dyn std::error::Error>> {
1557    let n_vocab = logits.len();
1558    let d = e.htod(logits)?;
1559    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1560}
1561
1562struct SpecPipeTraceClock {
1563    pair: usize,
1564    started: std::time::Instant,
1565}
1566
1567#[derive(Clone)]
1568struct SpecPipeTraceCtx {
1569    clock: std::sync::Arc<SpecPipeTraceClock>,
1570    round: usize,
1571    lane: usize,
1572}
1573
1574struct SpecPipeTraceMarker {
1575    trace: SpecPipeTraceCtx,
1576    phase: &'static str,
1577    edge: &'static str,
1578    slot: Option<usize>,
1579}
1580
1581unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1582    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1583    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1584    let slot = marker
1585        .slot
1586        .map(|v| v.to_string())
1587        .unwrap_or_else(|| "-".into());
1588    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1589    use std::io::Write as _;
1590    let stderr = std::io::stderr();
1591    let mut stderr = stderr.lock();
1592    let _ = writeln!(
1593        stderr,
1594        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1595         slot={slot} t_ms={t_ms:.3}",
1596        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1597    );
1598}
1599
1600fn enqueue_spec_pipe_trace_marker(
1601    stream: &cudarc::driver::CudaStream,
1602    trace: Option<&SpecPipeTraceCtx>,
1603    phase: &'static str,
1604    edge: &'static str,
1605    slot: Option<usize>,
1606) -> Result<(), Box<dyn std::error::Error>> {
1607    let Some(trace) = trace else {
1608        return Ok(());
1609    };
1610    let marker = Box::new(SpecPipeTraceMarker {
1611        trace: trace.clone(),
1612        phase,
1613        edge,
1614        slot,
1615    });
1616    let raw = Box::into_raw(marker);
1617    let result = unsafe {
1618        cudarc::driver::result::stream::launch_host_function(
1619            stream.cu_stream(),
1620            spec_pipe_trace_marker,
1621            raw.cast(),
1622        )
1623    };
1624    if let Err(err) = result {
1625        unsafe {
1626            drop(Box::from_raw(raw));
1627        }
1628        return Err(err.into());
1629    }
1630    Ok(())
1631}
1632
1633#[derive(Default)]
1634struct SpecPipeProgress {
1635    setup_done: [bool; 2],
1636    draft_done: [usize; 2],
1637    stage0_done: [usize; 2],
1638    verify_done: [usize; 2],
1639    accept_done: [usize; 2],
1640    finished: [bool; 2],
1641    aborted: bool,
1642}
1643
1644/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1645/// keeps its existing call stack and round locals; this object only orders phase entry. The
1646/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1647/// cannot be interleaved by the two host threads.
1648struct SpecPipeSync {
1649    progress: std::sync::Mutex<SpecPipeProgress>,
1650    changed: std::sync::Condvar,
1651    primary: std::sync::Mutex<()>,
1652    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1653}
1654
1655impl SpecPipeSync {
1656    fn new() -> Self {
1657        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1658        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1659            std::sync::Arc::new(SpecPipeTraceClock {
1660                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1661                started: std::time::Instant::now(),
1662            })
1663        });
1664        Self {
1665            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1666            changed: std::sync::Condvar::new(),
1667            primary: std::sync::Mutex::new(()),
1668            trace,
1669        }
1670    }
1671}
1672
1673#[derive(Clone)]
1674struct SpecPipeLane {
1675    sync: std::sync::Arc<SpecPipeSync>,
1676    lane: usize,
1677}
1678
1679impl SpecPipeLane {
1680    fn peer(&self) -> usize {
1681        1 - self.lane
1682    }
1683
1684    fn aborted() -> Box<dyn std::error::Error> {
1685        "paired speculative peer aborted".into()
1686    }
1687
1688    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1689        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1690            clock: clock.clone(),
1691            round,
1692            lane: self.lane,
1693        })
1694    }
1695
1696    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1697        let mut p = self.sync.progress.lock().unwrap();
1698        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1699            p = self.sync.changed.wait(p).unwrap();
1700        }
1701        if p.aborted {
1702            Err(Self::aborted())
1703        } else {
1704            Ok(())
1705        }
1706    }
1707
1708    fn setup_end(&self) {
1709        let mut p = self.sync.progress.lock().unwrap();
1710        p.setup_done[self.lane] = true;
1711        self.sync.changed.notify_all();
1712    }
1713
1714    fn draft_begin(
1715        &self,
1716        round: usize,
1717    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1718        let peer = self.peer();
1719        let mut p = self.sync.progress.lock().unwrap();
1720        loop {
1721            if p.aborted {
1722                return Err(Self::aborted());
1723            }
1724            let setup_ready =
1725                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1726            let prior_ready = p.accept_done[self.lane] >= round
1727                && (p.accept_done[peer] >= round || p.finished[peer]);
1728            let turn_ready = if self.lane == 0 {
1729                true
1730            } else {
1731                p.draft_done[0] > round || p.finished[0]
1732            };
1733            if setup_ready && prior_ready && turn_ready {
1734                break;
1735            }
1736            p = self.sync.changed.wait(p).unwrap();
1737        }
1738        drop(p);
1739        Ok(self.sync.primary.lock().unwrap())
1740    }
1741
1742    fn draft_end(&self, round: usize) {
1743        let mut p = self.sync.progress.lock().unwrap();
1744        p.draft_done[self.lane] = round + 1;
1745        self.sync.changed.notify_all();
1746    }
1747
1748    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1749    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1750    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1751        let peer = self.peer();
1752        let mut p = self.sync.progress.lock().unwrap();
1753        loop {
1754            if p.aborted {
1755                return Err(Self::aborted());
1756            }
1757            let ready = if self.lane == 0 {
1758                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1759            } else {
1760                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1761            };
1762            if ready {
1763                return Ok(self.lane == 0 || p.finished[peer]);
1764            }
1765            p = self.sync.changed.wait(p).unwrap();
1766        }
1767    }
1768
1769    fn stage0_end(&self, round: usize) {
1770        let mut p = self.sync.progress.lock().unwrap();
1771        p.stage0_done[self.lane] = round + 1;
1772        self.sync.changed.notify_all();
1773    }
1774
1775    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1776    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1777    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1778        let mut p = self.sync.progress.lock().unwrap();
1779        while !p.aborted
1780            && !(p.stage0_done[self.lane] > round
1781                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1782        {
1783            p = self.sync.changed.wait(p).unwrap();
1784        }
1785        if p.aborted {
1786            Err(Self::aborted())
1787        } else {
1788            Ok(())
1789        }
1790    }
1791
1792    fn verify_end(&self, round: usize) {
1793        let mut p = self.sync.progress.lock().unwrap();
1794        p.verify_done[self.lane] = round + 1;
1795        self.sync.changed.notify_all();
1796    }
1797
1798    fn accept_begin(
1799        &self,
1800        round: usize,
1801    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1802        let mut p = self.sync.progress.lock().unwrap();
1803        loop {
1804            if p.aborted {
1805                return Err(Self::aborted());
1806            }
1807            let ready = if self.lane == 0 {
1808                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1809            } else {
1810                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1811            };
1812            if ready {
1813                break;
1814            }
1815            p = self.sync.changed.wait(p).unwrap();
1816        }
1817        drop(p);
1818        Ok(self.sync.primary.lock().unwrap())
1819    }
1820
1821    fn accept_end(&self, round: usize) {
1822        let mut p = self.sync.progress.lock().unwrap();
1823        p.accept_done[self.lane] = round + 1;
1824        self.sync.changed.notify_all();
1825    }
1826
1827    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1828        self.sync.primary.lock().unwrap()
1829    }
1830
1831    fn finish(&self, failed: bool) {
1832        let mut p = self.sync.progress.lock().unwrap();
1833        p.finished[self.lane] = true;
1834        p.aborted |= failed;
1835        self.sync.changed.notify_all();
1836    }
1837}
1838
1839struct SpecPipeFinish<'a> {
1840    lane: &'a SpecPipeLane,
1841    closed: bool,
1842}
1843
1844impl<'a> SpecPipeFinish<'a> {
1845    fn new(lane: &'a SpecPipeLane) -> Self {
1846        Self {
1847            lane,
1848            closed: false,
1849        }
1850    }
1851
1852    fn close(&mut self, failed: bool) {
1853        self.lane.finish(failed);
1854        self.closed = true;
1855    }
1856}
1857
1858impl Drop for SpecPipeFinish<'_> {
1859    fn drop(&mut self) {
1860        if !self.closed {
1861            self.lane.finish(true);
1862        }
1863    }
1864}
1865
1866/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1867/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1868/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1869/// binds that context before touching the session, joins before returning, and never aliases the
1870/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1871/// session type Send.
1872struct SpecPipeSessionPtr(*mut SpecSession);
1873
1874unsafe impl Send for SpecPipeSessionPtr {}
1875
1876impl SpecPipeSessionPtr {
1877    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1878        unsafe { &mut *self.0 }
1879    }
1880}
1881
1882/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1883/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1884/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1885/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1886/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1887/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1888/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1889/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1890/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1891///
1892/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1893/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1894/// load-bearing:
1895///
1896/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1897///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1898///   This is all the key used to carry.
1899/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1900///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1901///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1902///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1903///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1904///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1905///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1906///
1907/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1908/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1909/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1910/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1911/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1912#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1913pub(crate) struct SampledGraphKey {
1914    seed: u64,
1915    temp_bits: u32,
1916    k: usize,
1917    top_k: i32,
1918    top_p_bits: u32,
1919    min_p_bits: u32,
1920    pen_on: bool,
1921}
1922
1923impl SampledGraphKey {
1924    pub(crate) fn new(
1925        seed: u64,
1926        temp: f32,
1927        k: usize,
1928        top_k: i32,
1929        top_p: f32,
1930        min_p: f32,
1931        pen_on: bool,
1932    ) -> Self {
1933        SampledGraphKey {
1934            seed,
1935            temp_bits: temp.to_bits(),
1936            k,
1937            top_k,
1938            top_p_bits: top_p.to_bits(),
1939            min_p_bits: min_p.to_bits(),
1940            pen_on,
1941        }
1942    }
1943
1944    /// The one regime the in-graph sampled chain may stand in for the eager one: nothing but
1945    /// temperature shapes `q`. Computed FROM THE KEY so the capture guard, the launch guard and
1946    /// the key can never drift apart (they were three separate expressions before this lane, and
1947    /// the launch site simply forgot to ask).
1948    pub(crate) fn pure_temp(&self) -> bool {
1949        self.top_k == 0
1950            && f32::from_bits(self.top_p_bits) >= 1.0
1951            && f32::from_bits(self.min_p_bits) <= 0.0
1952            && !self.pen_on
1953    }
1954}
1955
1956pub(crate) struct DraftGraphCtx {
1957    g_tok: CudaSlice<u32>,
1958    g_pos: CudaSlice<i32>,
1959    g_seed: CudaSlice<f32>,
1960    g_p: CudaSlice<f32>,
1961    g_ctr: CudaSlice<u32>,
1962    g_q: CudaSlice<f32>,
1963    g_perturb: CudaSlice<f32>,
1964    q_slots: Vec<CudaSlice<f32>>,
1965    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
1966    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
1967    /// per-position contents the host re-uploads before each replay (the graph-promote
1968    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
1969    g_dmask: CudaSlice<u32>,
1970    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
1971    graph_masked: bool,
1972    graph: Option<cudarc::driver::CudaGraph>,
1973    graph_s: Option<cudarc::driver::CudaGraph>,
1974    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
1975    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
1976    failed: DraftGraphFallback,
1977    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
1978    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
1979    s_key: Option<SampledGraphKey>,
1980    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
1981    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
1982    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
1983    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
1984    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
1985    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
1986    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
1987    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
1988    keeper: Vec<Box<dyn std::any::Any + Send>>,
1989    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
1990}
1991
1992/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
1993/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
1994///
1995/// Three contracts:
1996/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
1997///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
1998///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
1999///   an already-failed graph returns None (the per-burst memoization that keeps the eager
2000///   fallback from paying a doomed capture attempt every burst).
2001/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2002///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
2003///   failure for the pool's whole lifetime. Returns the note line only when a flag was
2004///   actually set (quiet on the common clean-resume path).
2005/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2006///   capture attempt whose own failure would re-flip loudly.
2007#[derive(Default)]
2008pub(crate) struct DraftGraphFallback {
2009    greedy: bool,
2010    sampled: bool,
2011}
2012impl DraftGraphFallback {
2013    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2014        if self.greedy {
2015            return None;
2016        }
2017        self.greedy = true;
2018        Some(format!(
2019            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2020        ))
2021    }
2022    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2023        if self.sampled {
2024            return None;
2025        }
2026        self.sampled = true;
2027        Some(format!(
2028            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2029        ))
2030    }
2031    fn greedy_failed(&self) -> bool {
2032        self.greedy
2033    }
2034    fn sampled_failed(&self) -> bool {
2035        self.sampled
2036    }
2037    fn clear_greedy(&mut self) {
2038        self.greedy = false;
2039    }
2040    fn clear_sampled(&mut self) {
2041        self.sampled = false;
2042    }
2043    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2044    /// was set (so clean resumes stay quiet).
2045    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2046        if !self.greedy && !self.sampled {
2047            return None;
2048        }
2049        let which = match (self.greedy, self.sampled) {
2050            (true, true) => "greedy+sampled",
2051            (true, false) => "greedy",
2052            _ => "sampled",
2053        };
2054        self.greedy = false;
2055        self.sampled = false;
2056        Some(format!(
2057            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2058        ))
2059    }
2060}
2061
2062impl DraftGraphCtx {
2063    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2064        Ok(DraftGraphCtx {
2065            g_tok: e.alloc_u32_zeroed(1)?,
2066            g_pos: e.htod_i32(&[0])?,
2067            g_seed: e.zeros(n_embd)?,
2068            g_p: e.zeros(1)?,
2069            g_ctr: e.alloc_u32_zeroed(1)?,
2070            g_q: e.zeros(qlen)?,
2071            g_perturb: e.zeros(qlen)?,
2072            q_slots: Vec::new(),
2073            g_dmask: e.alloc_u32_zeroed(1)?,
2074            graph_masked: false,
2075            graph: None,
2076            graph_s: None,
2077            failed: DraftGraphFallback::default(),
2078            s_key: None,
2079            keeper: Vec::new(),
2080            keeper_s: Vec::new(),
2081        })
2082    }
2083}
2084
2085pub(crate) struct MtpScratch {
2086    kv: KvLayer,
2087    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2088    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2089    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2090    /// smaller host-indexed SWA ring instead.
2091    cap: usize,
2092    extra: Vec<MtpScratchPlane>,
2093}
2094
2095struct MtpScratchPlane {
2096    kv: KvLayer,
2097    cap: usize,
2098}
2099
2100fn mtp_scratch_layout(
2101    cfg: &memra_gguf::config::ModelConfig,
2102    geom: Option<&crate::hybrid::DraftGeom>,
2103) -> (usize, usize, usize, usize) {
2104    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2105    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2106    let head_dim_k = cfg.head_dim_k as usize;
2107    let head_dim_v = cfg.head_dim_v as usize;
2108    assert!(
2109        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
2110        "KVQUANT requires head_dim%32==0 (MTP scratch)"
2111    );
2112    let kv_dim_k = head_dim_k * n_head_kv;
2113    let kv_dim_v = head_dim_v * n_head_kv;
2114    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2115    // policy shared with `MtpScratch::new` so admission scales the same allocation.
2116    let (kbb, vbb) = crate::kv_blk_bytes();
2117    let k_tok_bytes = (kv_dim_k / 32) * kbb;
2118    let v_tok_bytes = (kv_dim_v / 32) * vbb;
2119    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2120}
2121
2122fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2123    assert!(head_count > 0, "MTP chain requires at least one head");
2124    step % head_count
2125}
2126
2127impl MtpScratch {
2128    fn alloc_plane(
2129        e: &Engine,
2130        cfg: &memra_gguf::config::ModelConfig,
2131        plan: &memra_gguf::model_plan::ModelPlan,
2132        cap: usize,
2133        geom: Option<&crate::hybrid::DraftGeom>,
2134    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2135        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2136        let ring = if crate::cache::swa_ring_on()
2137            && crate::plan_backend::decode_batch_program(plan)
2138                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2139        {
2140            let window = plan
2141                .layers
2142                .iter()
2143                .find_map(|layer| match layer.attention {
2144                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2145                        Some(window as usize)
2146                    }
2147                    _ => None,
2148                })
2149                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2150            Some(crate::cache::KvRing::new(
2151                crate::cache::swa_ring_rows(window, cap),
2152                window,
2153            ))
2154        } else {
2155            None
2156        };
2157        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2158        // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2159        // KvLayer::base_d): the captured chain derives its physical rows from
2160        // (len_d, base_d, window) with zero per-token node updates.
2161        let base_d = match ring.as_ref() {
2162            Some(_) => Some(e.htod_i32(&[0])?),
2163            None => None,
2164        };
2165        Ok(MtpScratchPlane {
2166            kv: KvLayer {
2167                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2168                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2169                kv_dim_k,
2170                kv_dim_v,
2171                k_tok_bytes,
2172                v_tok_bytes,
2173                len: 0,
2174                ring,
2175                len_d: e.htod_i32(&[0])?,
2176                base_d,
2177            },
2178            cap,
2179        })
2180    }
2181
2182    fn new(
2183        e: &Engine,
2184        cfg: &memra_gguf::config::ModelConfig,
2185        plan: &memra_gguf::model_plan::ModelPlan,
2186        cap: usize,
2187        geom: Option<&crate::hybrid::DraftGeom>,
2188    ) -> Result<Self, Box<dyn std::error::Error>> {
2189        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2190        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2191        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2192        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2193        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2194        Ok(MtpScratch {
2195            kv: primary.kv,
2196            cap: primary.cap,
2197            extra: Vec::new(),
2198        })
2199    }
2200
2201    fn push_plane(
2202        &mut self,
2203        e: &Engine,
2204        cfg: &memra_gguf::config::ModelConfig,
2205        plan: &memra_gguf::model_plan::ModelPlan,
2206        geom: Option<&crate::hybrid::DraftGeom>,
2207    ) -> Result<(), Box<dyn std::error::Error>> {
2208        self.extra
2209            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2210        Ok(())
2211    }
2212
2213    fn plane_count(&self) -> usize {
2214        1 + self.extra.len()
2215    }
2216
2217    fn plane(&self, index: usize) -> (&KvLayer, usize) {
2218        if index == 0 {
2219            (&self.kv, self.cap)
2220        } else {
2221            let plane = &self.extra[index - 1];
2222            (&plane.kv, plane.cap)
2223        }
2224    }
2225
2226    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2227        if index == 0 {
2228            (&mut self.kv, self.cap)
2229        } else {
2230            let plane = &mut self.extra[index - 1];
2231            (&mut plane.kv, plane.cap)
2232        }
2233    }
2234
2235    // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2236    // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2237    // just that a rewind was refused.
2238    #[track_caller]
2239    fn set_plane_len(
2240        &mut self,
2241        e: &Engine,
2242        index: usize,
2243        n: usize,
2244    ) -> Result<(), Box<dyn std::error::Error>> {
2245        let caller = std::panic::Location::caller();
2246        let (kv, cap) = self.plane_mut(index);
2247        if let Some(ring) = kv.ring.as_ref() {
2248            if !ring.can_rewind_to(n) {
2249                // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2250                // vendor-default shape and it fires from more than one call path with more than
2251                // one trigger: a long generation walks the checkpoint out of the ring, but a
2252                // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2253                // explain. A bare message forced two rounds of guessing; the operands make each
2254                // trigger name itself.
2255                let raw = n.saturating_sub(ring.window().saturating_sub(1));
2256                return Err(format!(
2257                    "SWA ring MTP checkpoint has been lapped; full re-prime required (plane={index} rewind_to={n} window={} base={} rows={} cap={cap} needed_view_start={} < base, called from {caller})",
2258                    ring.window(),
2259                    ring.base(),
2260                    ring.rows(),
2261                    raw & !31usize,
2262                )
2263                .into());
2264            }
2265        }
2266        kv.len = n;
2267        e.set_i32_one(&mut kv.len_d, n as i32)
2268    }
2269
2270    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2271    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2272    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2273    #[track_caller]
2274    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2275        let caller = std::panic::Location::caller();
2276        if !self.can_rewind_to(n) {
2277            // set_plane_len re-checks and reports the operands; call it so the failure carries
2278            // which plane refused and why, instead of this bare aggregate.
2279            for index in 0..self.plane_count() {
2280                self.set_plane_len(e, index, n)?;
2281            }
2282            return Err(format!(
2283                "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2284            )
2285            .into());
2286        }
2287        for index in 0..self.plane_count() {
2288            self.set_plane_len(e, index, n)?;
2289        }
2290        Ok(())
2291    }
2292
2293    fn can_rewind_to(&self, n: usize) -> bool {
2294        (0..self.plane_count()).all(|index| {
2295            self.plane(index)
2296                .0
2297                .ring
2298                .as_ref()
2299                .is_none_or(|ring| ring.can_rewind_to(n))
2300        })
2301    }
2302
2303    /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2304    /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2305    /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2306    /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2307    /// flat planes and when the ring already has room; `len` is untouched either way.
2308    fn ensure_dcw_headroom(
2309        &mut self,
2310        e: &Engine,
2311        rows: usize,
2312    ) -> Result<(), Box<dyn std::error::Error>> {
2313        for index in 0..self.plane_count() {
2314            let (kv, _) = self.plane_mut(index);
2315            let Some(ring) = kv.ring.as_ref() else {
2316                continue;
2317            };
2318            let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2319            e.prepare_kv_append(kv, retain, rows)?;
2320        }
2321        Ok(())
2322    }
2323}
2324
2325/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2326/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2327/// full weight reads per round — recomputing columns the verify had already produced
2328/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2329/// to "after the first j verify columns" WITHOUT re-running the trunk:
2330/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2331///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2332///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2333///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2334///   pure-copy ring rebuild.
2335/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2336///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2337///   target: j <= t-1).
2338/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2339/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2340struct GdnStash {
2341    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2342    q_l2: CudaSlice<f32>,
2343    k_l2: CudaSlice<f32>,
2344    v_g: CudaSlice<f32>, // [t, num_v, d_state]
2345    g_log: CudaSlice<f32>,
2346    beta: CudaSlice<f32>, // [t, num_v]
2347}
2348pub(crate) struct VerifyCkpt {
2349    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2350    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2351}
2352/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2353pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2354
2355/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2356/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2357/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2358/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2359/// layers between full-attention layers are shape-static given vt — no positions, no
2360/// t_kv, state addressed through pointer tables — so runs of them capture per
2361/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2362/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2363///
2364/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2365/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2366/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2367/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2368/// before and restored after — the graph's first real launch starts from the exact
2369/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2370/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2371/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2372pub(crate) struct DsparkVerifyGraphs {
2373    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2374    lin: Vec<usize>,
2375    lin_pos: std::collections::HashMap<usize, usize>,
2376    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2377    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2378    table_all: CudaSlice<u64>,
2379    host_table: Vec<u64>,
2380    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2381    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2382    stash_conv: Vec<CudaSlice<f32>>,
2383    stash_ssm: Vec<CudaSlice<f32>>,
2384    conv_words: usize,
2385    ssm_words: usize,
2386    /// Per-vt input/output staging (stable addresses the graphs bake).
2387    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2388    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2389    /// so the sink buffer must live (and persist) with the graphs, not with the round.
2390    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2391    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2392    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2393    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2394    save_conv: CudaSlice<f32>,
2395    save_ssm: CudaSlice<f32>,
2396    max_run: usize,
2397    n_embd: usize,
2398    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2399    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2400    pub(crate) round_slab: bool,
2401    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2402    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2403    fa: Vec<usize>,
2404    fa_pos: std::collections::HashMap<usize, usize>,
2405    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2406    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2407    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2408    fa_table: CudaSlice<u64>,
2409    fa_host_table: Vec<u64>,
2410    t_cap: usize,
2411    /// Per-vt position staging for the captured bodies — contents refreshed per round
2412    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2413    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2414    /// Full-verify graphs keyed (vt, rung_end, hi).
2415    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2416    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2417    covered: usize,
2418    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2419    /// full-verify capture walks all of them.
2420    walk_uniform: bool,
2421    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2422    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2423    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2424    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2425    debt_obs: Option<(usize, usize)>,
2426}
2427
2428struct DsparkSegGraph {
2429    graph: cudarc::driver::CudaGraph,
2430    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2431}
2432
2433/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2434/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2435/// modes without a second copy of the math.
2436pub(crate) struct FaLayerArgs<'a> {
2437    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2438    /// them per-z (append slot = pos, T_kv = pos + 1).
2439    pub pos_d: &'a CudaSlice<i32>,
2440    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2441    /// arm builds/uses them (graph mode refuses that arm).
2442    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2443    pub pos0: usize,
2444    pub seqs_append: bool,
2445    pub batch_fa_on: bool,
2446    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2447    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2448    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2449    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2450    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2451    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2452    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2453    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2454    /// for FA layers that never touch it.
2455    pub ckpt: Option<&'a mut VerifyCkpt>,
2456}
2457
2458// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2459// no automatic trait; CUDA driver graph handles are context-scoped rather than
2460// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2461// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2462// single decode-stream thread.
2463unsafe impl Send for DsparkVerifyGraphs {}
2464
2465impl DsparkVerifyGraphs {
2466    /// Live capture count (segment + full graphs) — the denominator of
2467    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2468    pub(crate) fn captures(&self) -> usize {
2469        self.graphs.len() + self.full.len()
2470    }
2471
2472    /// Take the marginal-growth debt reading and record this observation for the next one.
2473    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2474    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2475        let captures = self.captures();
2476        let debt =
2477            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2478        if captures > 0 {
2479            match self.debt_obs {
2480                Some((c0, _)) if captures <= c0 => {}
2481                _ => self.debt_obs = Some((captures, reserved_bytes)),
2482            }
2483        }
2484        debt
2485    }
2486
2487    /// Build for this cache's shape. None when there are no linear layers, sizes are
2488    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2489    pub(crate) fn new(
2490        e: &Engine,
2491        cache: &Cache,
2492        t_max: usize,
2493        n_embd: usize,
2494    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2495        let lin: Vec<usize> = (0..cache.recur.len())
2496            .filter(|&il| cache.recur[il].is_some())
2497            .collect();
2498        if lin.is_empty() || t_max < 2 {
2499            return Ok(None);
2500        }
2501        let first = cache.recur[lin[0]].as_ref().unwrap();
2502        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2503        for &il in &lin {
2504            let rl = cache.recur[il].as_ref().unwrap();
2505            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2506                return Ok(None);
2507            }
2508        }
2509        let n = lin.len();
2510        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2511        for (k, &il) in lin.iter().enumerate() {
2512            lin_pos.insert(il, k);
2513        }
2514        // longest run of consecutive linear layers (save-scratch sizing)
2515        let mut max_run = 1usize;
2516        let mut run = 1usize;
2517        for w in lin.windows(2) {
2518            if w[1] == w[0] + 1 {
2519                run += 1;
2520                max_run = max_run.max(run);
2521            } else {
2522                run = 1;
2523            }
2524        }
2525        let rows = t_max - 1;
2526        let mut stash_conv = Vec::with_capacity(n);
2527        let mut stash_ssm = Vec::with_capacity(n);
2528        for _ in 0..n {
2529            stash_conv.push(e.uninit(rows * conv_words)?);
2530            stash_ssm.push(e.uninit(rows * ssm_words)?);
2531        }
2532        let host_table = vec![0u64; n * 6];
2533        let table_all = e.htod_u64(&host_table)?;
2534        // slice 4c: full-attention census for the full-verify graphs.
2535        let fa: Vec<usize> = (0..cache.kv.len())
2536            .filter(|&il| cache.kv[il].is_some())
2537            .collect();
2538        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2539        for (k, &il) in fa.iter().enumerate() {
2540            fa_pos.insert(il, k);
2541        }
2542        let n_layers = cache.kv.len().max(cache.recur.len());
2543        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2544        let walk_uniform = (0..n_layers).all(|il| {
2545            cache.recur.get(il).is_some_and(|r| r.is_some())
2546                != cache.kv.get(il).is_some_and(|k| k.is_some())
2547        });
2548        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2549        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2550        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2551        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2552        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2553        let covered = (0..n_layers)
2554            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2555            .count();
2556        let t_cap = t_max;
2557        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2558        let fa_table = e.htod_u64(&fa_host_table)?;
2559        Ok(Some(Self {
2560            lin,
2561            lin_pos,
2562            table_all,
2563            host_table,
2564            stash_conv,
2565            stash_ssm,
2566            conv_words,
2567            ssm_words,
2568            stage: std::collections::HashMap::new(),
2569            tap_bufs: std::collections::HashMap::new(),
2570            graphs: std::collections::HashMap::new(),
2571            save_conv: e.uninit(n * conv_words)?,
2572            save_ssm: e.uninit(n * ssm_words)?,
2573            max_run,
2574            n_embd,
2575            round_slab: false,
2576            fa,
2577            fa_pos,
2578            fa_table,
2579            fa_host_table,
2580            t_cap,
2581            pos_stage: std::collections::HashMap::new(),
2582            full: std::collections::HashMap::new(),
2583            covered,
2584            walk_uniform,
2585            debt_obs: None,
2586        }))
2587    }
2588
2589    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2590    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2591    /// cache buffers land at new addresses; a stale table would read the wrong state).
2592    pub(crate) fn refresh_tables(
2593        &mut self,
2594        e: &Engine,
2595        cache: &Cache,
2596    ) -> Result<(), Box<dyn std::error::Error>> {
2597        use cudarc::driver::DevicePtr;
2598        {
2599            let s = &e.gpu.stream();
2600            for (k, &il) in self.lin.iter().enumerate() {
2601                let rl = cache.recur[il].as_ref().unwrap();
2602                let (pc, _g0) = rl.conv_state.device_ptr(s);
2603                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2604                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2605                let o = k * 6;
2606                self.host_table[o] = pc as u64;
2607                self.host_table[o + 1] = p0 as u64;
2608                self.host_table[o + 2] = p1 as u64;
2609                self.host_table[o + 3] = pc as u64;
2610                self.host_table[o + 4] = p1 as u64;
2611                self.host_table[o + 5] = p0 as u64;
2612            }
2613            for (k, &il) in self.fa.iter().enumerate() {
2614                let kvl = cache.kv[il].as_ref().unwrap();
2615                let (pk, _g0) = kvl.k.device_ptr(s);
2616                let (pv, _g1) = kvl.v.device_ptr(s);
2617                let o = k * 2 * self.t_cap;
2618                for z in 0..self.t_cap {
2619                    self.fa_host_table[o + 2 * z] = pk as u64;
2620                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2621                }
2622            }
2623        }
2624        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2625        if !self.fa_host_table.is_empty() {
2626            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2627        }
2628        Ok(())
2629    }
2630
2631    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2632    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2633    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2634    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2635    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2636    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2637    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2638    /// captured graph is bit-identical for every round the rung covers.
2639    #[allow(clippy::too_many_arguments)]
2640    pub(crate) fn full_rung(
2641        &self,
2642        model: &crate::hybrid::HybridModel,
2643        cache: &Cache,
2644        lo: usize,
2645        hi: usize,
2646        t: usize,
2647        seqs_arms_on: bool,
2648    ) -> Option<usize> {
2649        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2650            static ONCE: std::sync::Once = std::sync::Once::new();
2651            let len0 = self
2652                .fa
2653                .first()
2654                .and_then(|&il| cache.kv[il].as_ref())
2655                .map(|k| k.len);
2656            ONCE.call_once(|| {
2657                eprintln!(
2658                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2659                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2660                    self.lin.len(), self.fa.len(), self.t_cap, len0
2661                );
2662            });
2663        }
2664        if !self.walk_uniform
2665            || !seqs_arms_on
2666            || !dspark_fa_rows_on()
2667            || t < 2
2668            || lo != 0
2669            || hi > self.covered
2670            || t > self.t_cap
2671            || self.fa.is_empty()
2672        {
2673            return None;
2674        }
2675        let cfg = &model.cfg;
2676        let head_dim_global = cfg.head_dim_k as usize;
2677        let nkv = cfg.n_head_kv as usize;
2678        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2679        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2680        // projection stride (the body's guard, hoisted so ineligible models fall back
2681        // instead of refusing mid-capture).
2682        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2683        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2684        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2685            return None;
2686        }
2687        let len0 = kvl0.len;
2688        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2689        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2690            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2691            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2692        {
2693            return None;
2694        }
2695        let rung = t_kv_last.next_power_of_two().max(256);
2696        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2697            return None;
2698        }
2699        Some(rung)
2700    }
2701
2702    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2703    /// the residual + refresh the per-vt position staging, capture on first encounter
2704    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2705    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2706    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2707    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2708    #[allow(clippy::too_many_arguments)]
2709    pub(crate) fn run_full(
2710        &mut self,
2711        model: &crate::hybrid::HybridModel,
2712        e: &Engine,
2713        lo: usize,
2714        hi: usize,
2715        x: &CudaSlice<f32>,
2716        t: usize,
2717        pos0: usize,
2718        rung: usize,
2719        cache: &mut Cache,
2720    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2721        let n_embd = self.n_embd;
2722        if !self.stage.contains_key(&t) {
2723            let xin = e.uninit(t * n_embd)?;
2724            let xout = e.uninit(t * n_embd)?;
2725            self.stage.insert(t, (xin, xout));
2726        }
2727        if !self.pos_stage.contains_key(&t) {
2728            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2729        }
2730        // Per-round refresh: position contents + input staging (both addresses are baked
2731        // by the captured bodies; only their CONTENTS change round to round).
2732        {
2733            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2734            let pb = self.pos_stage.get_mut(&t).unwrap();
2735            e.htod_i32_into(pb, &pos_host)?;
2736            let (xin, _) = self.stage.get_mut(&t).unwrap();
2737            e.copy_into(xin, 0, x, t * n_embd)?;
2738        }
2739        let key = (t, rung, hi);
2740        if !self.full.contains_key(&key) {
2741            // The warmups EXECUTE the whole walk on live state — save every linear
2742            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2743            // graph mode never bumps host lens and the appends write this round's own
2744            // slots).
2745            for (k, &il) in self.lin.iter().enumerate() {
2746                let rl = cache.recur[il].as_ref().unwrap();
2747                e.copy_into(
2748                    &mut self.save_conv,
2749                    k * self.conv_words,
2750                    &rl.conv_state,
2751                    self.conv_words,
2752                )?;
2753                e.copy_into(
2754                    &mut self.save_ssm,
2755                    k * self.ssm_words,
2756                    &rl.ssm_state,
2757                    self.ssm_words,
2758                )?;
2759            }
2760            let (graph, keeper) = {
2761                let table_all = &self.table_all;
2762                let lin_pos = &self.lin_pos;
2763                let fa_pos = &self.fa_pos;
2764                let fa_table = &self.fa_table;
2765                let t_cap = self.t_cap;
2766                let stash_conv = &mut self.stash_conv;
2767                let stash_ssm = &mut self.stash_ssm;
2768                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2769                let (xin, xout) = self
2770                    .stage
2771                    .get_mut(&t)
2772                    .map(|(a, b)| (&*a, b))
2773                    .expect("stage bucket created above");
2774                let cache_ref: &mut Cache = cache;
2775                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2776                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2777                } else {
2778                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2779                };
2780                e.capture_graph_retained_flags(iflag, move |e| {
2781                    let mut xc: Option<CudaSlice<f32>> = None;
2782                    for il in lo..hi {
2783                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2784                        let nx = if let Some(&k) = lin_pos.get(&il) {
2785                            model.qwen35_tparallel_linear_layer(
2786                                e,
2787                                il,
2788                                xr,
2789                                t,
2790                                cache_ref,
2791                                None,
2792                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2793                                Some((table_all, k * 6)),
2794                            )?
2795                        } else if let Some(&kf) = fa_pos.get(&il) {
2796                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2797                            model.qwen35_tparallel_fa_layer(
2798                                e,
2799                                il,
2800                                xr,
2801                                t,
2802                                cache_ref,
2803                                FaLayerArgs {
2804                                    pos_d,
2805                                    pos_rows: &mut no_rows,
2806                                    pos0,
2807                                    seqs_append: true,
2808                                    batch_fa_on: true,
2809                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2810                                    stream: None,
2811                                    ckpt: None,
2812                                },
2813                            )?
2814                        } else {
2815                            return Err(format!(
2816                                "run_full: layer {il} is neither linear nor full-attention"
2817                            )
2818                            .into());
2819                        };
2820                        xc = Some(nx);
2821                    }
2822                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2823                    Ok(())
2824                })?
2825            };
2826            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2827            // is odd -> 3 runs = net one swap), then restore the device state the
2828            // warmups consumed (walk scope only — layers past hi never executed). The
2829            // launch below then behaves exactly like one run.
2830            if t % 2 == 1 {
2831                for &il in &self.lin {
2832                    if il < lo || il >= hi {
2833                        continue;
2834                    }
2835                    let rl = cache.recur[il].as_mut().unwrap();
2836                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2837                }
2838            }
2839            for (k, &il) in self.lin.iter().enumerate() {
2840                if il < lo || il >= hi {
2841                    continue;
2842                }
2843                let rl = cache.recur[il].as_mut().unwrap();
2844                let (cw, sw) = (self.conv_words, self.ssm_words);
2845                {
2846                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2847                    let win = sv.slice(k * cw..(k + 1) * cw);
2848                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2849                }
2850                {
2851                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2852                    let win = sv.slice(k * sw..(k + 1) * sw);
2853                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2854                }
2855            }
2856            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2857                if let Ok(c) = crate::graph_update::node_census(&graph) {
2858                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2859                }
2860            }
2861            self.full.insert(
2862                key,
2863                DsparkSegGraph {
2864                    graph,
2865                    _keeper: keeper,
2866                },
2867            );
2868        }
2869        self.full[&key].graph.launch()?;
2870        // Host bookkeeping for the replayed body (captured host code does not re-run):
2871        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2872        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2873        // head layer's kv) that the walk never touches.
2874        if t % 2 == 1 {
2875            for &il in &self.lin {
2876                if il < lo || il >= hi {
2877                    continue;
2878                }
2879                let rl = cache.recur[il].as_mut().unwrap();
2880                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2881            }
2882        }
2883        for &il in &self.fa {
2884            if il < lo || il >= hi {
2885                continue;
2886            }
2887            cache.kv[il].as_mut().unwrap().len += t;
2888        }
2889        let (_, xout) = self.stage.get(&t).unwrap();
2890        let mut out = e.uninit(t * n_embd)?;
2891        e.copy_into(&mut out, 0, xout, t * n_embd)?;
2892        Ok(out)
2893    }
2894
2895    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
2896    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
2897    /// bracketed by a segment state save/restore), launch, then apply the host parity
2898    /// bookkeeping the captured body would have done. Returns the fresh residual.
2899    #[allow(clippy::too_many_arguments)]
2900    fn run_segment(
2901        &mut self,
2902        model: &crate::hybrid::HybridModel,
2903        e: &Engine,
2904        start: usize,
2905        end: usize,
2906        x: &CudaSlice<f32>,
2907        t: usize,
2908        cache: &mut Cache,
2909    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2910        let n_embd = self.n_embd;
2911        debug_assert!(end - start <= self.max_run);
2912        if !self.stage.contains_key(&t) {
2913            let xin = e.uninit(t * n_embd)?;
2914            let xout = e.uninit(t * n_embd)?;
2915            self.stage.insert(t, (xin, xout));
2916        }
2917        // Stage the residual at the bucket's baked input address.
2918        {
2919            let (xin, _) = self.stage.get_mut(&t).unwrap();
2920            e.copy_into(xin, 0, x, t * n_embd)?;
2921        }
2922        let key = (start, t);
2923        if !self.graphs.contains_key(&key) {
2924            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
2925            // ssm of every segment layer first, restore after, so the graph's first real
2926            // launch starts from the exact pre-round state (bytes gated e2e).
2927            for (k, il) in (start..end).enumerate() {
2928                let rl = cache.recur[il].as_ref().unwrap();
2929                e.copy_into(
2930                    &mut self.save_conv,
2931                    k * self.conv_words,
2932                    &rl.conv_state,
2933                    self.conv_words,
2934                )?;
2935                e.copy_into(
2936                    &mut self.save_ssm,
2937                    k * self.ssm_words,
2938                    &rl.ssm_state,
2939                    self.ssm_words,
2940                )?;
2941            }
2942            let (graph, keeper) = {
2943                let table_all = &self.table_all;
2944                let lin_pos = &self.lin_pos;
2945                let stash_conv = &mut self.stash_conv;
2946                let stash_ssm = &mut self.stash_ssm;
2947                let (xin, xout) = self
2948                    .stage
2949                    .get_mut(&t)
2950                    .map(|(a, b)| (&*a, b))
2951                    .expect("stage bucket created above");
2952                let cache_ref: &mut Cache = cache;
2953                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
2954                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
2955                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
2956                // = ~0.41 ms/round, most of the eager-launch savings. The captured
2957                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
2958                // (every transient drops inside the capture region — the generic
2959                // capture path's census precedent, 1589/1589), so AUTO_FREE has
2960                // nothing to reclaim and the graph is legal to instantiate without
2961                // it; PRIORITY is the flag the gemma slotted door ships for exactly
2962                // this reason (both alternatives drop the scan; UPLOAD via
2963                // cuGraphInstantiateWithFlags is WithParams-only and refused).
2964                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
2965                // the node census at capture (the ALLOC==FREE receipt).
2966                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2967                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2968                } else {
2969                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2970                };
2971                e.capture_graph_retained_flags(iflag, move |e| {
2972                    let mut xc: Option<CudaSlice<f32>> = None;
2973                    for il in start..end {
2974                        let k = lin_pos[&il];
2975                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2976                        let nx = model.qwen35_tparallel_linear_layer(
2977                            e,
2978                            il,
2979                            xr,
2980                            t,
2981                            cache_ref,
2982                            None,
2983                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
2984                            Some((table_all, k * 6)),
2985                        )?;
2986                        xc = Some(nx);
2987                    }
2988                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2989                    Ok(())
2990                })?
2991            };
2992            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2993            // is odd -> 3 runs = net one swap), then restore the device state the
2994            // warmups consumed. The launch below then behaves exactly like one run.
2995            if t % 2 == 1 {
2996                for il in start..end {
2997                    let rl = cache.recur[il].as_mut().unwrap();
2998                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2999                }
3000            }
3001            for (k, il) in (start..end).enumerate() {
3002                let rl = cache.recur[il].as_mut().unwrap();
3003                let (cw, sw) = (self.conv_words, self.ssm_words);
3004                {
3005                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3006                    let win = sv.slice(k * cw..(k + 1) * cw);
3007                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3008                }
3009                {
3010                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3011                    let win = sv.slice(k * sw..(k + 1) * sw);
3012                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3013                }
3014            }
3015            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3016                if let Ok(c) = crate::graph_update::node_census(&graph) {
3017                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3018                }
3019            }
3020            self.graphs.insert(
3021                key,
3022                DsparkSegGraph {
3023                    graph,
3024                    _keeper: keeper,
3025                },
3026            );
3027        }
3028        self.graphs[&key].graph.launch()?;
3029        // Host parity bookkeeping for the replayed body (the captured host swaps do not
3030        // re-run at replay).
3031        if t % 2 == 1 {
3032            for il in start..end {
3033                let rl = cache.recur[il].as_mut().unwrap();
3034                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3035            }
3036        }
3037        let (_, xout) = self.stage.get(&t).unwrap();
3038        let mut out = e.uninit(t * n_embd)?;
3039        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3040        Ok(out)
3041    }
3042
3043    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3044    fn can_capture(&self) -> bool {
3045        self.graphs.len() + self.full.len() < dspark_vg_cap()
3046    }
3047
3048    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3049    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3050    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3051    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3052    /// refusal would stash some layers in the ctx slabs and others in the round's cols
3053    /// while one commit reads only one of them.
3054    pub(crate) fn segments_ready(
3055        &self,
3056        model: &crate::hybrid::HybridModel,
3057        lo: usize,
3058        hi: usize,
3059        t: usize,
3060    ) -> bool {
3061        if self.can_capture() {
3062            return true;
3063        }
3064        let mut il = lo;
3065        while il < hi {
3066            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3067                let start = il;
3068                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3069                    il += 1;
3070                }
3071                if !self.graphs.contains_key(&(start, t)) {
3072                    return false;
3073                }
3074            } else {
3075                il += 1;
3076            }
3077        }
3078        true
3079    }
3080
3081    /// Widest verify window this pool was built for. A caller whose round exceeds it must
3082    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3083    /// past them is a panic rather than a refusal.
3084    pub(crate) fn t_capacity(&self) -> usize {
3085        self.t_cap
3086    }
3087
3088    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3089    /// `row` (0-based) of layer `il`. None for non-linear layers.
3090    pub(crate) fn slab_row(
3091        &self,
3092        e: &Engine,
3093        il: usize,
3094        row: usize,
3095    ) -> Option<(u64, u64, usize, usize)> {
3096        use cudarc::driver::DevicePtr;
3097        let k = *self.lin_pos.get(&il)?;
3098        let s = &e.gpu.stream();
3099        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3100        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3101        Some((
3102            pc as u64 + (row * self.conv_words * 4) as u64,
3103            ps as u64 + (row * self.ssm_words * 4) as u64,
3104            self.conv_words,
3105            self.ssm_words,
3106        ))
3107    }
3108}
3109
3110impl VerifyCkpt {
3111    fn new(n_layer: usize) -> Self {
3112        VerifyCkpt {
3113            gdn: (0..n_layer).map(|_| None).collect(),
3114            cols: (0..n_layer).map(|_| None).collect(),
3115        }
3116    }
3117}
3118
3119/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3120/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3121/// a logical round number.
3122struct VerifyBoundaryTicket {
3123    rt: &'static crate::pp::PpNRt,
3124    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3125    slot: usize,
3126    pos0: usize,
3127    t: usize,
3128    payload: usize,
3129    n_st: usize,
3130    pipelined: bool,
3131    pp_anatomy: bool,
3132    pp_started: std::time::Instant,
3133    reverse_ms: f64,
3134    stage0_ms: f64,
3135    tx_ms: f64,
3136    trace: Option<SpecPipeTraceCtx>,
3137}
3138
3139/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3140/// increment-2 controller can also be armed by the server's fresh-process research door.
3141#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3142pub enum OptiForkGateMode {
3143    Disabled,
3144    Hit,
3145    Miss,
3146    Alternate,
3147    Abort,
3148    Controller,
3149}
3150
3151static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3152static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3153    std::sync::atomic::AtomicU32::new(0);
3154static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3155static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3156static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3157static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3158static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3159static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3160static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3161static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3162static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3163static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3164    std::sync::atomic::AtomicU64::new(0);
3165static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3166    std::sync::atomic::AtomicU64::new(0);
3167static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3168
3169impl OptiForkGateMode {
3170    fn code(self) -> u8 {
3171        match self {
3172            Self::Disabled => 0,
3173            Self::Hit => 1,
3174            Self::Miss => 2,
3175            Self::Alternate => 3,
3176            Self::Abort => 4,
3177            Self::Controller => 5,
3178        }
3179    }
3180
3181    fn configured() -> Self {
3182        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3183            1 => Self::Hit,
3184            2 => Self::Miss,
3185            3 => Self::Alternate,
3186            4 => Self::Abort,
3187            5 => Self::Controller,
3188            _ => Self::Disabled,
3189        }
3190    }
3191
3192    fn action(self, generation: u64) -> OptiForkAction {
3193        match self {
3194            Self::Hit => OptiForkAction::Hit,
3195            Self::Miss => OptiForkAction::Miss,
3196            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3197            Self::Alternate => OptiForkAction::Miss,
3198            Self::Abort => OptiForkAction::Abort,
3199            Self::Disabled | Self::Controller => {
3200                unreachable!("non-forced mode cannot choose a forced fork action")
3201            }
3202        }
3203    }
3204
3205    fn is_forced(self) -> bool {
3206        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3207    }
3208}
3209
3210/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3211pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3212    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3213}
3214
3215/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3216/// two-token draft-probability product. Serving can call this only through its explicit
3217/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3218pub fn set_optipipe_controller_threshold(threshold: f32) {
3219    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3220    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3221    set_optipipe_gate_mode(OptiForkGateMode::Controller);
3222}
3223
3224#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3225pub struct OptiForkGateStats {
3226    pub attempts: u64,
3227    pub hits: u64,
3228    pub misses: u64,
3229    pub abort_drains: u64,
3230    pub refusals: u64,
3231    pub gate_checks: u64,
3232    pub gate_admits: u64,
3233    pub gate_rejects: u64,
3234    pub reconciles: u64,
3235    pub wasted_draft_tokens: u64,
3236    pub shadow_draft_tokens: u64,
3237    pub breaker_trips: u64,
3238}
3239
3240#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3241pub struct OptiForkStateIdentity {
3242    pub trunk_kv_bytes: usize,
3243    pub recurrent_bytes: usize,
3244    pub scratch_kv_bytes: usize,
3245    pub hidden_bytes: usize,
3246}
3247
3248pub fn reset_optipipe_gate_stats() {
3249    for counter in [
3250        &OPTI_FORK_ATTEMPTS,
3251        &OPTI_FORK_HITS,
3252        &OPTI_FORK_MISSES,
3253        &OPTI_FORK_ABORT_DRAINS,
3254        &OPTI_FORK_REFUSALS,
3255        &OPTI_GATE_CHECKS,
3256        &OPTI_GATE_ADMITS,
3257        &OPTI_GATE_REJECTS,
3258        &OPTI_RECONCILES,
3259        &OPTI_WASTED_DRAFT_TOKENS,
3260        &OPTI_SHADOW_DRAFT_TOKENS,
3261        &OPTI_BREAKER_TRIPS,
3262    ] {
3263        counter.store(0, std::sync::atomic::Ordering::Relaxed);
3264    }
3265}
3266
3267pub fn optipipe_gate_stats() -> OptiForkGateStats {
3268    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3269    OptiForkGateStats {
3270        attempts: load(&OPTI_FORK_ATTEMPTS),
3271        hits: load(&OPTI_FORK_HITS),
3272        misses: load(&OPTI_FORK_MISSES),
3273        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3274        refusals: load(&OPTI_FORK_REFUSALS),
3275        gate_checks: load(&OPTI_GATE_CHECKS),
3276        gate_admits: load(&OPTI_GATE_ADMITS),
3277        gate_rejects: load(&OPTI_GATE_REJECTS),
3278        reconciles: load(&OPTI_RECONCILES),
3279        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3280        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3281        breaker_trips: load(&OPTI_BREAKER_TRIPS),
3282    }
3283}
3284
3285#[derive(Clone, Copy, Debug)]
3286struct OptiControllerPolicy {
3287    threshold: f32,
3288    consecutive_misses: u8,
3289    breaker_tripped: bool,
3290}
3291
3292impl OptiControllerPolicy {
3293    fn configured() -> Self {
3294        Self {
3295            threshold: f32::from_bits(
3296                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3297            ),
3298            consecutive_misses: 0,
3299            breaker_tripped: false,
3300        }
3301    }
3302
3303    fn admit(&self, q_proxy: f32) -> bool {
3304        q_proxy.is_finite()
3305            && (0.0..=1.0).contains(&q_proxy)
3306            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3307    }
3308
3309    /// Returns true exactly when this resolution newly trips the three-miss breaker.
3310    fn resolve(&mut self, hit: bool) -> bool {
3311        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3312        // every optimistic opportunity, so the safety breaker is measured separately and must
3313        // not silently turn this arm into "three attempts then serial".
3314        if self.threshold == 0.0 {
3315            self.consecutive_misses = 0;
3316            return false;
3317        }
3318        if hit {
3319            self.consecutive_misses = 0;
3320            return false;
3321        }
3322        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3323        if !self.breaker_tripped && self.consecutive_misses >= 3 {
3324            self.breaker_tripped = true;
3325            return true;
3326        }
3327        false
3328    }
3329}
3330
3331#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3332enum OptiForkAction {
3333    Hit,
3334    Miss,
3335    Abort,
3336}
3337
3338#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3339struct OptiForkGeneration {
3340    id: u64,
3341    slot: usize,
3342}
3343
3344#[derive(Default)]
3345struct OptiForkGenerationTracker {
3346    next: u64,
3347    live: [Option<u64>; 2],
3348}
3349
3350impl OptiForkGenerationTracker {
3351    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3352        let generation = OptiForkGeneration {
3353            id: self.next,
3354            slot: (self.next & 1) as usize,
3355        };
3356        if let Some(live) = self.live[generation.slot] {
3357            return Err(format!(
3358                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3359                generation.slot,
3360            )
3361            .into());
3362        }
3363        self.next += 1;
3364        self.live[generation.slot] = Some(generation.id);
3365        Ok(generation)
3366    }
3367
3368    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3369        match self.live[generation.slot] {
3370            Some(id) if id == generation.id => {
3371                self.live[generation.slot] = None;
3372                Ok(())
3373            }
3374            other => Err(format!(
3375                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3376                generation.id, generation.slot,
3377            )
3378            .into()),
3379        }
3380    }
3381}
3382
3383struct OptiForkSeedGeneration {
3384    h_seed: CudaSlice<f32>,
3385    fill_prev: CudaSlice<f32>,
3386    scratch_len: usize,
3387}
3388
3389/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3390/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3391/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3392/// device ownership.
3393fn opti_snapshot_stage_owned(
3394    e: &Engine,
3395    cache: &Cache,
3396    rt: &'static crate::pp::PpNRt,
3397    fence: &[usize],
3398) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3399    let n = cache.kv.len();
3400    let mut snapshot = crate::cache::CacheSnapshot {
3401        kv_len: vec![None; n],
3402        tp_kv_len: vec![None; n],
3403        conv: (0..n).map(|_| None).collect(),
3404        ssm: (0..n).map(|_| None).collect(),
3405        pos: cache.pos,
3406    };
3407    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3408    Ok(snapshot)
3409}
3410
3411fn opti_snapshot_stage_owned_into(
3412    e: &Engine,
3413    cache: &Cache,
3414    rt: &'static crate::pp::PpNRt,
3415    fence: &[usize],
3416    snapshot: &mut crate::cache::CacheSnapshot,
3417) -> Result<(), Box<dyn std::error::Error>> {
3418    if fence.len() != rt.n_stages() + 1
3419        || snapshot.kv_len.len() != cache.kv.len()
3420        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3421    {
3422        return Err("optipipe stage-owned snapshot shape mismatch".into());
3423    }
3424    for stage in 0..rt.n_stages() {
3425        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3426    }
3427    snapshot.pos = cache.pos;
3428    Ok(())
3429}
3430
3431/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3432/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3433/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3434/// either point would capture one side of the fork at the wrong generation.
3435fn opti_snapshot_one_stage_owned_into(
3436    e: &Engine,
3437    cache: &Cache,
3438    rt: &'static crate::pp::PpNRt,
3439    fence: &[usize],
3440    stage: usize,
3441    snapshot: &mut crate::cache::CacheSnapshot,
3442) -> Result<(), Box<dyn std::error::Error>> {
3443    if fence.len() != rt.n_stages() + 1
3444        || snapshot.kv_len.len() != cache.kv.len()
3445        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3446        || stage >= rt.n_stages()
3447    {
3448        return Err("optipipe single-stage snapshot shape mismatch".into());
3449    }
3450    let _scope = rt.enter(stage);
3451    let owner = rt.engine(stage, e);
3452    for il in fence[stage]..fence[stage + 1] {
3453        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3454        snapshot.tp_kv_len[il] = cache.tp_kv[il]
3455            .as_ref()
3456            .map(crate::tp::ResidentTpKvCache::committed_len);
3457        match &cache.recur[il] {
3458            Some(recur) => {
3459                match snapshot.conv[il].as_mut() {
3460                    Some(dst) => {
3461                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3462                    }
3463                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3464                }
3465                match snapshot.ssm[il].as_mut() {
3466                    Some(dst) => {
3467                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3468                    }
3469                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3470                }
3471            }
3472            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3473                return Err(
3474                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3475                );
3476            }
3477            None => {}
3478        }
3479    }
3480    snapshot.pos = cache.pos;
3481    Ok(())
3482}
3483
3484/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3485/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3486/// resolve, so the reconcile tables and conditional restores are stage-local.
3487struct OptiForkState {
3488    mode: OptiForkGateMode,
3489    controller: Option<OptiControllerPolicy>,
3490    generations: OptiForkGenerationTracker,
3491    active_snapshot_slot: usize,
3492    alternate_snapshot: crate::cache::CacheSnapshot,
3493    seeds: [OptiForkSeedGeneration; 2],
3494    rt: &'static crate::pp::PpNRt,
3495    fence: [usize; 3],
3496    split: usize,
3497    len_ptrs: CudaSlice<u64>,
3498    saved_lens: CudaSlice<i32>,
3499    forced_acc: CudaSlice<u32>,
3500    valid: CudaSlice<u32>,
3501    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3502    logical_payload_bytes: [usize; 2],
3503}
3504
3505struct OptiForkTicket {
3506    generation: OptiForkGeneration,
3507    boundary: Option<VerifyBoundaryTicket>,
3508    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3509    settled: bool,
3510}
3511
3512struct OptiControllerTicket {
3513    generation: OptiForkGeneration,
3514    boundary: Option<VerifyBoundaryTicket>,
3515    ckpt: Option<VerifyCkpt>,
3516    verify_tokens: [u32; 2],
3517    draft_prob: f32,
3518    eager_seed: Option<CudaSlice<f32>>,
3519    q_proxy: f32,
3520    scratch_len: usize,
3521    issued_at: std::time::Instant,
3522    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3523    settled: bool,
3524}
3525
3526struct OptiControllerPrepared {
3527    verify_tokens: [u32; 2],
3528    draft_prob: f32,
3529    eager_seed: Option<CudaSlice<f32>>,
3530    q_proxy: f32,
3531    scratch_len: usize,
3532}
3533
3534impl OptiControllerTicket {
3535    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3536        self.boundary
3537            .take()
3538            .expect("controller boundary ticket already consumed")
3539    }
3540
3541    fn take_ckpt(&mut self) -> VerifyCkpt {
3542        self.ckpt
3543            .take()
3544            .expect("controller verify checkpoint already consumed")
3545    }
3546
3547    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3548        self.eager_seed.take()
3549    }
3550
3551    fn settle(&mut self) {
3552        self.settled = true;
3553    }
3554}
3555
3556impl Drop for OptiControllerTicket {
3557    fn drop(&mut self) {
3558        if !self.settled {
3559            let _ = self.drain.synchronize();
3560            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3561        }
3562    }
3563}
3564
3565impl OptiForkTicket {
3566    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3567        self.boundary
3568            .take()
3569            .expect("fork ticket boundary already consumed")
3570    }
3571
3572    fn settle(&mut self) {
3573        self.settled = true;
3574    }
3575}
3576
3577impl Drop for OptiForkTicket {
3578    fn drop(&mut self) {
3579        if !self.settled {
3580            let _ = self.drain.synchronize();
3581            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3582        }
3583    }
3584}
3585
3586impl OptiForkState {
3587    #[allow(clippy::too_many_arguments)]
3588    fn new(
3589        e: &Engine,
3590        cache: &Cache,
3591        mode: OptiForkGateMode,
3592        alternate_snapshot: crate::cache::CacheSnapshot,
3593        h_seed: &CudaSlice<f32>,
3594        fill_prev: &CudaSlice<f32>,
3595        rt: &'static crate::pp::PpNRt,
3596        split: usize,
3597        n_layer: usize,
3598    ) -> Result<Self, Box<dyn std::error::Error>> {
3599        let fence = [0, split, n_layer];
3600        let mut logical_payload_bytes = [0usize; 2];
3601        for stage in 0..2 {
3602            for il in fence[stage]..fence[stage + 1] {
3603                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3604                    .as_ref()
3605                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3606                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3607                    .as_ref()
3608                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3609            }
3610        }
3611        let seeds = [
3612            OptiForkSeedGeneration {
3613                h_seed: e.clone_dtod(h_seed)?,
3614                fill_prev: e.clone_dtod(fill_prev)?,
3615                scratch_len: 0,
3616            },
3617            OptiForkSeedGeneration {
3618                h_seed: e.clone_dtod(h_seed)?,
3619                fill_prev: e.clone_dtod(fill_prev)?,
3620                scratch_len: 0,
3621            },
3622        ];
3623        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3624            let _stage = rt.enter(0);
3625            let e0 = rt.engine(0, e);
3626            (
3627                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3628                e0.htod_i32(&vec![0; split])?,
3629                e0.alloc_u32_zeroed(2)?,
3630                e0.alloc_u32_zeroed(1)?,
3631                e0.stream(),
3632            )
3633        };
3634        logical_payload_bytes[0] += seeds
3635            .iter()
3636            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3637            .sum::<usize>();
3638        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3639            + saved_lens.len() * std::mem::size_of::<i32>()
3640            + forced_acc.len() * std::mem::size_of::<u32>()
3641            + valid.len() * std::mem::size_of::<u32>();
3642        Ok(Self {
3643            mode,
3644            controller: (mode == OptiForkGateMode::Controller)
3645                .then(OptiControllerPolicy::configured),
3646            generations: OptiForkGenerationTracker::default(),
3647            active_snapshot_slot: 0,
3648            alternate_snapshot,
3649            seeds,
3650            rt,
3651            fence,
3652            split,
3653            len_ptrs,
3654            saved_lens,
3655            forced_acc,
3656            valid,
3657            stage0_stream,
3658            logical_payload_bytes,
3659        })
3660    }
3661
3662    fn reserve(
3663        &mut self,
3664        current_snapshot: &mut crate::cache::CacheSnapshot,
3665    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3666        let generation = self.generations.reserve()?;
3667        if generation.slot != self.active_snapshot_slot {
3668            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3669            self.active_snapshot_slot = generation.slot;
3670        }
3671        Ok(generation)
3672    }
3673
3674    fn capture_seed(
3675        &mut self,
3676        e: &Engine,
3677        generation: OptiForkGeneration,
3678        h_seed: &CudaSlice<f32>,
3679        fill_prev: &CudaSlice<f32>,
3680        scratch_len: usize,
3681    ) -> Result<(), Box<dyn std::error::Error>> {
3682        let seed = &mut self.seeds[generation.slot];
3683        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3684        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3685        seed.scratch_len = scratch_len;
3686        Ok(())
3687    }
3688
3689    fn ticket(
3690        &self,
3691        generation: OptiForkGeneration,
3692        boundary: VerifyBoundaryTicket,
3693    ) -> OptiForkTicket {
3694        OptiForkTicket {
3695            generation,
3696            boundary: Some(boundary),
3697            drain: self.stage0_stream.clone(),
3698            settled: false,
3699        }
3700    }
3701
3702    #[allow(clippy::too_many_arguments)]
3703    fn controller_ticket(
3704        &self,
3705        generation: OptiForkGeneration,
3706        boundary: VerifyBoundaryTicket,
3707        ckpt: VerifyCkpt,
3708        verify_tokens: [u32; 2],
3709        draft_prob: f32,
3710        eager_seed: Option<CudaSlice<f32>>,
3711        q_proxy: f32,
3712        scratch_len: usize,
3713    ) -> OptiControllerTicket {
3714        OptiControllerTicket {
3715            generation,
3716            boundary: Some(boundary),
3717            ckpt: Some(ckpt),
3718            verify_tokens,
3719            draft_prob,
3720            eager_seed,
3721            q_proxy,
3722            scratch_len,
3723            issued_at: std::time::Instant::now(),
3724            drain: self.stage0_stream.clone(),
3725            settled: false,
3726        }
3727    }
3728
3729    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3730        self.generations.reserve()
3731    }
3732
3733    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3734        &mut self.alternate_snapshot
3735    }
3736
3737    fn promote_successor_snapshot(
3738        &mut self,
3739        current_snapshot: &mut crate::cache::CacheSnapshot,
3740        generation: OptiForkGeneration,
3741    ) {
3742        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3743        self.active_snapshot_slot = generation.slot;
3744    }
3745
3746    fn queue_actual_reconcile(
3747        &mut self,
3748        e: &Engine,
3749        snapshot: &crate::cache::CacheSnapshot,
3750        acc: &CudaSlice<u32>,
3751        optimistic_pending: u32,
3752        base: usize,
3753    ) -> Result<(), Box<dyn std::error::Error>> {
3754        let saved: Vec<i32> = (0..self.split)
3755            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3756            .collect();
3757        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3758        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3759        // the validity/reconcile kernels must never peer-read acc before it is written. The
3760        // increment-1 harness uses primary stage 0, where stream order already provides this.
3761        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3762            self.rt.fence_stages_behind(&e.stream())?;
3763        }
3764        let _stage = self.rt.enter(0);
3765        let e0 = self.rt.engine(0, e);
3766        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3767        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3768        e0.spec_fork_reconcile_kv(
3769            &self.len_ptrs,
3770            &self.saved_lens,
3771            acc,
3772            &self.valid,
3773            base,
3774            self.split,
3775        )
3776    }
3777
3778    fn finish_actual_reconcile(
3779        &mut self,
3780        e: &Engine,
3781        cache: &mut Cache,
3782        snapshot: &crate::cache::CacheSnapshot,
3783        n_acc: usize,
3784        base: usize,
3785        hit: bool,
3786    ) -> Result<(), Box<dyn std::error::Error>> {
3787        if hit {
3788            return Ok(());
3789        }
3790        let len_delta = base + n_acc;
3791        for il in 0..self.split {
3792            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3793                kv.len = saved + len_delta;
3794            }
3795        }
3796        {
3797            let _stage = self.rt.enter(1);
3798            let e1 = self.rt.engine(1, e);
3799            for il in self.split..self.fence[2] {
3800                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3801                    kv.len = saved + len_delta;
3802                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3803                }
3804            }
3805        }
3806        self.rt.publish_to(0, &e.stream())?;
3807        Ok(())
3808    }
3809
3810    fn cancel_controller_ticket(
3811        &mut self,
3812        e: &Engine,
3813        cache: &mut Cache,
3814        scratch: &mut MtpScratch,
3815        snapshot: &crate::cache::CacheSnapshot,
3816        ticket: &mut OptiControllerTicket,
3817    ) -> Result<(), Box<dyn std::error::Error>> {
3818        {
3819            let _stage = self.rt.enter(0);
3820            let e0 = self.rt.engine(0, e);
3821            for il in 0..self.split {
3822                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3823                    kv.len = saved;
3824                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3825                }
3826            }
3827        }
3828        scratch.set_len(e, snapshot.pos)?;
3829        ticket.settle();
3830        self.generations.retire(ticket.generation)?;
3831        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3832        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3833        eprintln!(
3834            "[opti-controller] tail-drain generation={} slot={}",
3835            ticket.generation.id, ticket.generation.slot,
3836        );
3837        Ok(())
3838    }
3839
3840    #[allow(clippy::too_many_arguments)]
3841    fn reconcile(
3842        &mut self,
3843        e: &Engine,
3844        cache: &mut Cache,
3845        scratch: &mut MtpScratch,
3846        snapshot: &crate::cache::CacheSnapshot,
3847        h_seed: &mut CudaSlice<f32>,
3848        fill_prev: &mut CudaSlice<f32>,
3849        generation: OptiForkGeneration,
3850        action: OptiForkAction,
3851        optimistic_pending: u32,
3852    ) -> Result<(), Box<dyn std::error::Error>> {
3853        debug_assert!(action != OptiForkAction::Abort);
3854        let miss_started = std::time::Instant::now();
3855        let keep = action == OptiForkAction::Hit;
3856        let saved: Vec<i32> = (0..self.split)
3857            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3858            .collect();
3859        let seed = &self.seeds[generation.slot];
3860        {
3861            let _stage = self.rt.enter(0);
3862            let e0 = self.rt.engine(0, e);
3863            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3864            let forced = if keep {
3865                [1u32, optimistic_pending]
3866            } else {
3867                [0u32, optimistic_pending]
3868            };
3869            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3870            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3871            e0.spec_fork_reconcile_kv(
3872                &self.len_ptrs,
3873                &self.saved_lens,
3874                &self.forced_acc,
3875                &self.valid,
3876                0,
3877                self.split,
3878            )?;
3879            for il in 0..self.split {
3880                if let Some(recur) = cache.recur[il].as_mut() {
3881                    let conv = snapshot.conv[il]
3882                        .as_ref()
3883                        .ok_or("optipipe stage0 snapshot missing conv state")?;
3884                    let ssm = snapshot.ssm[il]
3885                        .as_ref()
3886                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
3887                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
3888                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
3889                }
3890            }
3891            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
3892            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
3893        }
3894
3895        if keep {
3896            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3897            return Ok(());
3898        }
3899
3900        for il in 0..self.split {
3901            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3902                kv.len = saved;
3903            }
3904        }
3905        scratch.set_len(e, seed.scratch_len)?;
3906        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
3907        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
3908        let caller = e.stream();
3909        self.rt.publish_to(0, &caller)?;
3910        caller.synchronize()?;
3911        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
3912        eprintln!(
3913            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
3914            generation.id, generation.slot,
3915        );
3916        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3917        Ok(())
3918    }
3919
3920    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3921        self.generations.retire(generation)
3922    }
3923}
3924
3925fn rewind_tp_kv_verified_prefix(
3926    tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
3927    saved_lens: &[Option<usize>],
3928    accepted: usize,
3929) -> Result<(), Box<dyn std::error::Error>> {
3930    if tp_kv.len() != saved_lens.len() {
3931        return Err("spec TP KV snapshot shape mismatch".into());
3932    }
3933    for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
3934        match (cache.as_mut(), *saved) {
3935            (Some(cache), Some(saved)) => {
3936                let target = saved
3937                    .checked_add(accepted)
3938                    .ok_or("spec TP KV committed length overflow")?;
3939                cache.rewind_to(target)?;
3940            }
3941            (None, None) => {}
3942            _ => {
3943                return Err(
3944                    format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
3945                );
3946            }
3947        }
3948    }
3949    Ok(())
3950}
3951
3952/// MEMRA_WALK_SCRATCH=1: persistent per-layer temporaries for the verify walk. The walk allocated
3953/// `next`, `x1_t`, `z_t` and `x2_t` fresh EVERY LAYER — ~4 allocations x 45 layers x 2 ranks per
3954/// round. This session measured a device allocation at ~20 us (the sampled split head arrived
3955/// slower than the unsplit one purely on five allocations per token), so ~200 allocations is
3956/// ~4 ms of the 8.6 ms/round the [spec-phase] host-issue term reports. Stable addresses are also
3957/// the precondition for ever capturing this walk in a CUDA graph.
3958struct WalkScratch {
3959    dev: usize,
3960    cap: usize,
3961    x1: CudaSlice<f32>,
3962    z: CudaSlice<f32>,
3963}
3964static WALK_SCRATCH: std::sync::Mutex<Option<WalkScratch>> = std::sync::Mutex::new(None);
3965
3966fn walk_scratch_on() -> bool {
3967    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3968    *ON.get_or_init(|| std::env::var("MEMRA_WALK_SCRATCH").as_deref() == Ok("1"))
3969}
3970
3971/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
3972/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
3973static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3974static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3975static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3976
3977impl HybridModel {
3978    fn mtp_head_count(&self) -> usize {
3979        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
3980    }
3981
3982    fn mtp_head_at(&self, index: usize) -> &MtpHead {
3983        if index == 0 {
3984            self.mtp.as_ref().expect("MTP head 0 is unavailable")
3985        } else {
3986            &self.mtp_extra[index - 1]
3987        }
3988    }
3989
3990    fn new_mtp_scratch(
3991        &self,
3992        e: &Engine,
3993        cap: usize,
3994    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
3995        let mut scratch = MtpScratch::new(
3996            e,
3997            &self.cfg,
3998            &self.plan,
3999            cap,
4000            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4001        )?;
4002        for head in &self.mtp_extra {
4003            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4004        }
4005        Ok(scratch)
4006    }
4007
4008    fn opti_graph_draft_step(
4009        &self,
4010        e: &Engine,
4011        mtp: &MtpHead,
4012        dctx: &mut DraftGraphCtx,
4013        scratch: &mut MtpScratch,
4014        d_vocab: usize,
4015    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4016        // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4017        // host-side before launching (no-op on flat planes).
4018        if step35_draft_dcw_on() {
4019            scratch.ensure_dcw_headroom(e, 2)?;
4020        }
4021        dctx.graph
4022            .as_ref()
4023            .ok_or("optipipe controller requires the greedy draft graph")?
4024            .launch()?;
4025        scratch.kv.len += 1;
4026        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4027        if (idx as usize) >= d_vocab {
4028            return Err(
4029                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4030            );
4031        }
4032        let probability = e.dtoh(&dctx.g_p)?[0];
4033        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4034            return Err(format!("optipipe draft probability is invalid: {probability}").into());
4035        }
4036        let token = match &mtp.d2t {
4037            Some(map) => map[idx as usize],
4038            None => idx,
4039        };
4040        if token != idx {
4041            e.set_u32_one(&mut dctx.g_tok, token)?;
4042        }
4043        Ok((token, probability))
4044    }
4045
4046    #[allow(clippy::too_many_arguments)]
4047    fn opti_controller_draft_step(
4048        &self,
4049        e: &Engine,
4050        mtp: &MtpHead,
4051        dctx: &mut DraftGraphCtx,
4052        scratch: &mut MtpScratch,
4053        d_vocab: usize,
4054        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4055        eager_pos: usize,
4056        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4057    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4058        if dctx.graph.is_some() {
4059            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4060        }
4061        let (input_token, input_seed) = eager_state
4062            .take()
4063            .ok_or("optipipe eager continuation seed is unavailable")?;
4064        let (logits, next_seed) = self.mtp_head_forward_dev(
4065            e,
4066            mtp,
4067            input_token,
4068            &input_seed,
4069            scratch,
4070            eager_pos,
4071            embd_dev,
4072            None,
4073        )?;
4074        let token_d = e.argmax_token_device(&logits, d_vocab)?;
4075        let idx = e.dtoh_u32_one(&token_d)?;
4076        if (idx as usize) >= d_vocab {
4077            return Err(format!(
4078                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4079            )
4080            .into());
4081        }
4082        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4083        let probability = e.dtoh(&probability_d)?[0];
4084        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4085            return Err(
4086                format!("optipipe eager draft probability is invalid: {probability}").into(),
4087            );
4088        }
4089        let token = match &mtp.d2t {
4090            Some(map) => map[idx as usize],
4091            None => idx,
4092        };
4093        *eager_state = Some((token, next_seed));
4094        Ok((token, probability))
4095    }
4096
4097    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4098    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4099    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4100    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4101    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4102    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4103    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4104    /// transfer + host argmax per draft token from the K-token draft chain.
4105    #[allow(clippy::too_many_arguments)]
4106    fn mtp_head_forward_dev(
4107        &self,
4108        e: &Engine,
4109        mtp: &MtpHead,
4110        e_tok: u32,
4111        h_seed: &CudaSlice<f32>,
4112        scratch: &mut MtpScratch,
4113        mtp_pos: usize,
4114        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4115        mask: Option<(&CudaSlice<u32>, usize)>,
4116    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4117        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4118    }
4119
4120    #[allow(clippy::too_many_arguments)]
4121    fn mtp_head_forward_dev_at(
4122        &self,
4123        e: &Engine,
4124        mtp: &MtpHead,
4125        e_tok: u32,
4126        h_seed: &CudaSlice<f32>,
4127        scratch: &mut MtpScratch,
4128        scratch_index: usize,
4129        mtp_pos: usize,
4130        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4131        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4132        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4133        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4134        mask: Option<(&CudaSlice<u32>, usize)>,
4135    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4136        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4137        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4138        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4139        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4140        static ANAT_NS: [AtomicU64; 5] = [
4141            AtomicU64::new(0),
4142            AtomicU64::new(0),
4143            AtomicU64::new(0),
4144            AtomicU64::new(0),
4145            AtomicU64::new(0),
4146        ];
4147        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4148        let anat = {
4149            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4150            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4151        };
4152        if anat {
4153            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4154        }
4155        let t_all = std::time::Instant::now();
4156        let mut t_ph = std::time::Instant::now();
4157        let mut anat_mark = |i: usize,
4158                             e: &Engine,
4159                             t: &mut std::time::Instant|
4160         -> Result<(), Box<dyn std::error::Error>> {
4161            if anat {
4162                e.stream().synchronize()?;
4163                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4164                *t = std::time::Instant::now();
4165            }
4166            Ok(())
4167        };
4168        let cfg = &self.cfg;
4169        let n_embd = cfg.n_embd as usize;
4170        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4171        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4172        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4173        let eps = cfg.rms_eps;
4174        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4175
4176        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4177        // expands this one row on CPU and transfers n_embd f32 values instead.
4178        let e_emb = match embd_dev {
4179            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4180            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4181        };
4182
4183        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4184        let mut e_norm = e.zeros(n_embd)?;
4185        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4186        let mut h_norm = e.zeros(n_embd)?;
4187        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4188
4189        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4190        let mut concat = e.zeros(2 * n_embd)?;
4191        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4192        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4193
4194        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4195        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4196
4197        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4198        let mut a_norm = e.zeros(di)?;
4199        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4200        anat_mark(0, e, &mut t_ph)?;
4201
4202        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4203        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4204        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4205        // advances only the device counter).
4206        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4207            // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4208            // the captured chain (draft parity by construction). Per-step ring headroom runs
4209            // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4210            // plain dc arm below.
4211            (Mixer::Full(fa), Some(g)) if self.step35_dcw_eligible(g) => {
4212                {
4213                    let (kv, _) = scratch.plane_mut(scratch_index);
4214                    let retain = match kv.ring.as_ref() {
4215                        Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4216                        None => 0,
4217                    };
4218                    e.prepare_kv_append(kv, retain, 1)?;
4219                }
4220                let out =
4221                    self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4222                scratch.plane_mut(scratch_index).0.len += 1;
4223                out
4224            }
4225            // step35 MTP block, door off (the shipping default until receipts): PER-LAYER
4226            // geometry + a separate head-wise gate + an SWA window, none of which the plain dc
4227            // launcher can express (see `mtp_step35_attn`). Host-len arm. Advances BOTH the
4228            // host len and the device counter itself (unlike the dc arm, whose host-side
4229            // mirror the caller does).
4230            (Mixer::Full(fa), Some(g)) => {
4231                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4232            }
4233            (Mixer::Full(fa), None) => {
4234                let out = self.mtp_full_attn_dc(
4235                    e,
4236                    fa,
4237                    &a_norm,
4238                    &pos_d,
4239                    scratch,
4240                    scratch_index,
4241                    mtp.geom.as_ref(),
4242                )?;
4243                scratch.plane_mut(scratch_index).0.len += 1;
4244                out
4245            }
4246            (Mixer::Linear(_), _) => {
4247                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4248            }
4249            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
4250        };
4251        anat_mark(1, e, &mut t_ph)?;
4252
4253        // op 7: x1 = inpSA + attn_out
4254        let mut x1 = e.zeros(di)?;
4255        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4256
4257        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
4258        let mut z = e.zeros(di)?;
4259        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4260
4261        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4262        let ffn_out = match &mtp.ffn {
4263            crate::hybrid::Ffn::Dense {
4264                ffn_gate,
4265                ffn_up,
4266                ffn_down,
4267            } => {
4268                let n_ff = ffn_gate.out_features();
4269                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4270                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4271                    (
4272                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4273                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4274                    )
4275                } else {
4276                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4277                };
4278                let mut act = e.zeros(n_ff)?;
4279                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4280                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4281                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4282                // passes None, which is `ffn_act`'s dispatch verbatim.
4283                Self::ffn_act_lim(
4284                    e,
4285                    &self.cfg,
4286                    &gate,
4287                    &up,
4288                    1.0,
4289                    1.0,
4290                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
4291                    &mut act,
4292                    n_ff,
4293                )?;
4294                e.matmul(ffn_down, &act, 1)?
4295            }
4296            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4297            // so they never alias trunk layer 0's cache keys.
4298            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4299        };
4300        anat_mark(2, e, &mut t_ph)?;
4301
4302        // op 10: h_nextn = x1 + ffn_out (at di)
4303        let mut h_inner = e.zeros(di)?;
4304        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4305
4306        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4307        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4308        let h_nextn = match mtp.geom.as_ref() {
4309            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4310            None => h_inner,
4311        };
4312
4313        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4314        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4315        let mut final_h = e.zeros(n_embd)?;
4316        e.rms_norm(
4317            &h_nextn,
4318            final_norm.float_data(),
4319            &mut final_h,
4320            n_embd,
4321            1,
4322            eps,
4323        )?;
4324
4325        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4326        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4327        let mut logits = e.matmul(head, &final_h, 1)?;
4328        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4329        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4330        if let Some((mask_d, mw)) = mask {
4331            let d_vocab = head.out_features();
4332            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4333        }
4334        anat_mark(3, e, &mut t_ph)?;
4335        if anat {
4336            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4337            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4338            if n % 128 == 0 {
4339                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4340                eprintln!(
4341                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4342                    us(0),
4343                    us(1),
4344                    us(2),
4345                    us(3),
4346                    us(4)
4347                );
4348            }
4349        }
4350        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4351        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4352        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4353    }
4354
4355    #[allow(clippy::too_many_arguments)]
4356    fn mtp_chain_forward_dev(
4357        &self,
4358        e: &Engine,
4359        tokens: &[u32],
4360        seeds: &[CudaSlice<f32>],
4361        scratch: &mut MtpScratch,
4362        committed_scratch_len: usize,
4363        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4364        mask: Option<(&CudaSlice<u32>, usize)>,
4365    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4366        if tokens.is_empty() || tokens.len() != seeds.len() {
4367            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
4368        }
4369        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
4370        let head = self.mtp_head_at(index);
4371        scratch.set_plane_len(e, index, committed_scratch_len)?;
4372
4373        let mut last = None;
4374        for row in 0..tokens.len() {
4375            let is_last = row + 1 == tokens.len();
4376            last = Some(self.mtp_head_forward_dev_at(
4377                e,
4378                head,
4379                tokens[row],
4380                &seeds[row],
4381                scratch,
4382                index,
4383                committed_scratch_len + row + 1,
4384                embd_dev,
4385                if is_last { mask } else { None },
4386            )?);
4387        }
4388        Ok(last.expect("non-empty MTP prefix produced no row"))
4389    }
4390
4391    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
4392    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
4393    /// the dc path, and all three are properties of this arch's MTP block:
4394    ///
4395    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
4396    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
4397    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
4398    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
4399    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
4400    ///    starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
4401    ///    `mtp_step35_attn_dcw` (`fa_decode_dcw`, flag-doored via MEMRA_STEP35_DRAFT_DCW);
4402    ///    this host-len arm remains the default until the door's receipts land.
4403    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
4404    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
4405    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
4406    ///    resolved `Step35MtpGeom`, never from `cfg`.
4407    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
4408    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
4409    ///    fused-into-wq `q_gate_split` form the dc arm handles.
4410    ///
4411    /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW unset, `mtp_head_forward_cap` refuses step35
4412    /// heads explicitly (rather than silently capturing a window-less, wrong-past-`win` graph)
4413    /// and this eager chain IS the served path. With the door armed, BOTH draft modes run the
4414    /// `mtp_step35_attn_dcw` twin instead of this arm.
4415    ///
4416    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
4417    /// caller must not mirror.
4418    fn mtp_step35_attn(
4419        &self,
4420        e: &Engine,
4421        fa: &FullAttnLayer,
4422        g: &crate::hybrid::Step35MtpGeom,
4423        h: &CudaSlice<f32>,
4424        pos_d: &CudaSlice<i32>,
4425        scratch: &mut MtpScratch,
4426        scratch_index: usize,
4427    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4428        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4429        // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
4430        // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
4431        // the first three explanations for that gap were all wrong: head assignment (step-modulo
4432        // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
4433        // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
4434        // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
4435        // shows up only as acceptance — so it gets a standing receipt rather than another reading
4436        // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
4437        // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
4438        {
4439            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4440            ONCE.get_or_init(|| {
4441                eprintln!(
4442                    "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4443                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4444                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4445                );
4446            });
4447        }
4448        let eps = self.cfg.rms_eps;
4449        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4450        let n_embd = self.cfg.n_embd as usize;
4451        let gw = fa
4452            .attn_gate
4453            .as_ref()
4454            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4455
4456        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4457            && e.uses_q8_1_fast(&fa.wk)
4458            && e.uses_q8_1_fast(&fa.wv)
4459            && e.uses_q8_1_fast(gw)
4460        {
4461            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4462            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4463                Some(t3) => t3,
4464                None => (
4465                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4466                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4467                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4468                ),
4469            };
4470            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4471        } else {
4472            (
4473                e.matmul(&fa.wq, h, 1)?,
4474                e.matmul(&fa.wk, h, 1)?,
4475                e.matmul(&fa.wv, h, 1)?,
4476                e.matmul(gw, h, 1)?,
4477            )
4478        };
4479
4480        let mut q = e.uninit(nh * hd)?;
4481        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4482        let mut k = e.uninit(nkv * hd)?;
4483        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4484        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4485        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4486        // the resolved flag, not the constant, so an all-full sibling stays correct.
4487        let ff = if g.swa {
4488            None
4489        } else {
4490            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4491        };
4492        #[cfg(debug_assertions)]
4493        if let Some(ff) = ff {
4494            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4495        }
4496        e.rope_neox2(
4497            &mut q,
4498            &mut k,
4499            pos_d,
4500            hd,
4501            g.n_rot,
4502            nh,
4503            nkv,
4504            1,
4505            g.rope_base,
4506            1.0,
4507            ff,
4508        )?;
4509
4510        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4511        // length on the host anyway, and the windowed view below needs it there to compute the
4512        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4513        // dc-family consumer of this scratch still agree.
4514        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4515        assert!(
4516            kv.len < scratch_cap,
4517            "step35 MTP scratch overflow ({} >= {})",
4518            kv.len,
4519            scratch_cap
4520        );
4521        let next_len = kv.len + 1;
4522        let (off, t_kv) = if g.swa && next_len > g.window {
4523            (next_len - g.window, g.window)
4524        } else {
4525            (0, next_len)
4526        };
4527        // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
4528        // rewind that follows this append is still resident. THIS is the only site that rebases
4529        // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
4530        // that decides `base` for everyone.
4531        let retain_from = match kv.ring.as_ref() {
4532            Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4533            None => off & !31usize,
4534        };
4535        let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
4536        e.append_kv_quantized(
4537            &k,
4538            &v0,
4539            &mut kv.k,
4540            &mut kv.v,
4541            write_row,
4542            kv.kv_dim_k,
4543            kv.kv_dim_v,
4544            kv.k_tok_bytes,
4545            kv.v_tok_bytes,
4546            false,
4547        )?;
4548        kv.len = next_len;
4549        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4550        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4551        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4552        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4553        // therefore live, not theoretical.
4554        let physical = kv.physical_rows(off, off + t_kv)?;
4555        let k_view = e.view_u8_range(
4556            &kv.k,
4557            physical.start * kv.k_tok_bytes,
4558            physical.end * kv.k_tok_bytes,
4559        );
4560        let v_view = e.view_u8_range(
4561            &kv.v,
4562            physical.start * kv.v_tok_bytes,
4563            physical.end * kv.v_tok_bytes,
4564        );
4565        let mut attn = e.uninit(nh * hd)?;
4566        e.fa_decode_kvmod(
4567            &q,
4568            &k_view,
4569            &v_view,
4570            &mut attn,
4571            hd,
4572            nh,
4573            nkv,
4574            t_kv,
4575            scale,
4576            kv.k_tok_bytes,
4577            kv.v_tok_bytes,
4578            false,
4579        )?;
4580
4581        let mut ag = e.uninit(nh * hd)?;
4582        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4583        Ok(e.matmul(&fa.wo, &ag, 1)?)
4584    }
4585
4586    /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4587    /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4588    /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4589    /// fallback point) and the CAP site refuses with the named reason instead.
4590    fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom) -> bool {
4591        let hd = self.cfg.head_dim_k as usize;
4592        step35_draft_dcw_on()
4593            && g.swa
4594            && g.window >= crate::fa_vec_min_tkv()
4595            && std::env::var("MEMRA_NO_FA_VEC").is_err()
4596            && crate::fa_v3_active(hd)
4597            && hd <= 256
4598    }
4599
4600    /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4601    /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4602    /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4603    /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4604    /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4605    /// contract plus the view offset the plain `_dc` kernel could not express (the old
4606    /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4607    /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4608    /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4609    ///
4610    /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4611    /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4612    /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4613    /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4614    /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4615    /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4616    /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4617    /// arbitrates emitted bytes; acceptance is gated by the battery).
4618    ///
4619    /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4620    /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4621    /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4622    /// because a rebase is host work no captured chain may contain.
4623    fn mtp_step35_attn_dcw(
4624        &self,
4625        e: &Engine,
4626        fa: &FullAttnLayer,
4627        g: &crate::hybrid::Step35MtpGeom,
4628        h: &CudaSlice<f32>,
4629        pos_d: &CudaSlice<i32>,
4630        scratch: &mut MtpScratch,
4631        scratch_index: usize,
4632    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4633        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4634        // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4635        // naming the arm, so a serving log proves WHICH draft attention program ran (the
4636        // engagement receipt for the flag door, both directions).
4637        {
4638            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4639            ONCE.get_or_init(|| {
4640                eprintln!(
4641                    "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4642                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4643                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4644                );
4645            });
4646        }
4647        let eps = self.cfg.rms_eps;
4648        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4649        let n_embd = self.cfg.n_embd as usize;
4650        let gw = fa
4651            .attn_gate
4652            .as_ref()
4653            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4654
4655        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4656            && e.uses_q8_1_fast(&fa.wk)
4657            && e.uses_q8_1_fast(&fa.wv)
4658            && e.uses_q8_1_fast(gw)
4659        {
4660            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4661            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4662                Some(t3) => t3,
4663                None => (
4664                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4665                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4666                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4667                ),
4668            };
4669            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4670        } else {
4671            (
4672                e.matmul(&fa.wq, h, 1)?,
4673                e.matmul(&fa.wk, h, 1)?,
4674                e.matmul(&fa.wv, h, 1)?,
4675                e.matmul(gw, h, 1)?,
4676            )
4677        };
4678
4679        let mut q = e.zeros(nh * hd)?;
4680        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4681        let mut k = e.zeros(nkv * hd)?;
4682        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4683        // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4684        // (the eager twin's rule, resolved from the flag, not the constant).
4685        let ff = if g.swa {
4686            None
4687        } else {
4688            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4689        };
4690        #[cfg(debug_assertions)]
4691        if let Some(ff) = ff {
4692            crate::debug_assert_tensor_stream_device(
4693                ff,
4694                &e.stream(),
4695                "mtp_step35_attn_dcw.rope_freqs",
4696            );
4697        }
4698        e.rope_neox2(
4699            &mut q,
4700            &mut k,
4701            pos_d,
4702            hd,
4703            g.n_rot,
4704            nh,
4705            nkv,
4706            1,
4707            g.rope_base,
4708            1.0,
4709            ff,
4710        )?;
4711
4712        let (kv, cap) = scratch.plane_mut(scratch_index);
4713        // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4714        // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4715        e.append_kv_quantized_dcw(
4716            &k,
4717            &v0,
4718            &mut kv.k,
4719            &mut kv.v,
4720            &kv.len_d,
4721            kv.base_d.as_ref(),
4722            kv.kv_dim_k,
4723            kv.kv_dim_v,
4724            kv.k_tok_bytes,
4725            kv.v_tok_bytes,
4726        )?;
4727        e.inc_seqlen(&mut kv.len_d)?;
4728        // Full-buffer views (any in-round physical row stays in range under the headroom
4729        // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4730        let k_view = e.view_u8(&kv.k, kv.k.len());
4731        let v_view = e.view_u8(&kv.v, kv.v.len());
4732        let bucket = g.window.min(cap);
4733        let mut attn = e.zeros(nh * hd)?;
4734        e.fa_decode_dcw(
4735            &q,
4736            &k_view,
4737            &v_view,
4738            &mut attn,
4739            hd,
4740            nh,
4741            nkv,
4742            &kv.len_d,
4743            kv.base_d.as_ref(),
4744            if g.swa { g.window } else { 0 },
4745            bucket,
4746            scale,
4747            kv.k_tok_bytes,
4748            kv.v_tok_bytes,
4749            None,
4750        )?;
4751
4752        let mut ag = e.zeros(nh * hd)?;
4753        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4754        Ok(e.matmul(&fa.wo, &ag, 1)?)
4755    }
4756
4757    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4758    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4759    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4760    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4761    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4762    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4763    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4764    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4765    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4766    fn mtp_full_attn_dc(
4767        &self,
4768        e: &Engine,
4769        fa: &FullAttnLayer,
4770        h: &CudaSlice<f32>,
4771        pos_d: &CudaSlice<i32>,
4772        scratch: &mut MtpScratch,
4773        scratch_index: usize,
4774        geom: Option<&crate::hybrid::DraftGeom>,
4775    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4776        let cfg = &self.cfg;
4777        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4778        let geometry = cfg.full_attention_geometry_at(mtp_il);
4779        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4780        let n_head_kv = geom
4781            .map(|g| g.n_head_kv)
4782            .unwrap_or(geometry.n_head_kv as usize);
4783        let head_dim = geometry.head_dim_k as usize;
4784        let eps = cfg.rms_eps;
4785        let scale = geometry.attention_scale();
4786        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4787        let bucket_max = scratch.plane(scratch_index).1;
4788
4789        let (qf, mut k, v) =
4790            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4791                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4792                (
4793                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4794                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4795                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4796                )
4797            } else {
4798                (
4799                    e.matmul(&fa.wq, h, 1)?,
4800                    e.matmul(&fa.wk, h, 1)?,
4801                    e.matmul(&fa.wv, h, 1)?,
4802                )
4803            };
4804        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4805        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4806        let (mut q, gate) = if gated {
4807            let mut q = e.zeros(n_head * head_dim)?;
4808            let mut gate = e.zeros(n_head * head_dim)?;
4809            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4810            (q, Some(gate))
4811        } else {
4812            (qf, None)
4813        };
4814
4815        let mut qn = e.zeros(n_head * head_dim)?;
4816        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4817        q = qn;
4818        let mut kn = e.zeros(n_head_kv * head_dim)?;
4819        e.rms_norm(
4820            &k,
4821            fa.k_norm.float_data(),
4822            &mut kn,
4823            head_dim,
4824            n_head_kv,
4825            eps,
4826        )?;
4827        k = kn;
4828        let rope_dims = geometry.n_rot as usize;
4829        e.rope_neox(
4830            &mut q,
4831            pos_d,
4832            head_dim,
4833            rope_dims,
4834            n_head,
4835            1,
4836            geometry.rope_base,
4837            1.0,
4838        )?;
4839        e.rope_neox(
4840            &mut k,
4841            pos_d,
4842            head_dim,
4843            rope_dims,
4844            n_head_kv,
4845            1,
4846            geometry.rope_base,
4847            1.0,
4848        )?;
4849
4850        let kv = scratch.plane_mut(scratch_index).0;
4851        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4852        e.append_kv_quantized_dc(
4853            &k,
4854            &v,
4855            &mut kv.k,
4856            &mut kv.v,
4857            &kv.len_d,
4858            kv.kv_dim_k,
4859            kv.kv_dim_v,
4860            kv.k_tok_bytes,
4861            kv.v_tok_bytes,
4862            false,
4863        )?;
4864        e.inc_seqlen(&mut kv.len_d)?;
4865        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4866        // key range from the device counter.
4867        let k_view = e.view_u8(&kv.k, kv.k.len());
4868        let v_view = e.view_u8(&kv.v, kv.v.len());
4869        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4870        let mut attn = e.zeros(n_head * head_dim)?;
4871        e.fa_decode_dc(
4872            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4873            scale, ktb, vtb, false,
4874        )?;
4875
4876        let attn_g = match &gate {
4877            Some(gate) => {
4878                let mut gsig = e.zeros(n_head * head_dim)?;
4879                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4880                let mut ag = e.zeros(n_head * head_dim)?;
4881                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
4882                ag
4883            }
4884            None => attn,
4885        };
4886        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
4887    }
4888
4889    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
4890    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
4891    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
4892    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
4893    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
4894    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
4895    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
4896    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
4897    #[allow(clippy::too_many_arguments)]
4898    fn mtp_kv_fill_at(
4899        &self,
4900        e: &Engine,
4901        mtp: &MtpHead,
4902        tokens: &[u32],
4903        h: &CudaSlice<f32>,
4904        pos0: usize,
4905        scratch: &mut MtpScratch,
4906        scratch_index: usize,
4907        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4908    ) -> Result<(), Box<dyn std::error::Error>> {
4909        let cfg = &self.cfg;
4910        let n_embd = cfg.n_embd as usize;
4911        let eps = cfg.rms_eps;
4912        let t = tokens.len();
4913        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
4914        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
4915        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
4916        let Mixer::Full(fa) = &mtp.mixer else {
4917            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4918        };
4919        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
4920        let pos_d = e.htod_i32(&pos_vec)?;
4921
4922        // ops A/1/2: embed + the two input norms, T-wide.
4923        let e_emb = match embd_dev {
4924            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
4925            None => e.htod(&self.embd.gather(n_embd, tokens))?,
4926        };
4927        let mut e_norm = e.zeros(t * n_embd)?;
4928        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
4929        let mut h_norm = e.zeros(t * n_embd)?;
4930        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
4931
4932        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
4933        let mut concat = e.zeros(t * 2 * n_embd)?;
4934        for i in 0..t {
4935            e.copy_view_into(
4936                &mut concat,
4937                i * 2 * n_embd,
4938                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
4939                n_embd,
4940            )?;
4941            e.copy_view_into(
4942                &mut concat,
4943                i * 2 * n_embd + n_embd,
4944                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
4945                n_embd,
4946            )?;
4947        }
4948
4949        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
4950        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4951        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
4952        let mut a_norm = e.zeros(t * di)?;
4953        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
4954
4955        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
4956        // the fill only has to leave correct K/V rows behind for later chains to attend over.
4957        let n_head_kv = mtp
4958            .geom
4959            .as_ref()
4960            .map(|g| g.n_head_kv)
4961            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
4962            .unwrap_or_else(|| {
4963                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4964                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
4965            });
4966        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4967        let geometry = cfg.full_attention_geometry_at(mtp_il);
4968        let head_dim = geometry.head_dim_k as usize;
4969        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
4970        let v = e.matmul(&fa.wv, &a_norm, t)?;
4971        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
4972        e.rms_norm(
4973            &k,
4974            fa.k_norm.float_data(),
4975            &mut kn,
4976            head_dim,
4977            n_head_kv * t,
4978            eps,
4979        )?;
4980        k = kn;
4981        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
4982        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
4983        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
4984        // writes K rows the attention arm then re-derives at a different theta: correct-looking
4985        // output with dead acceptance, invisible to the exactness gates.
4986        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
4987            Some(s) => (
4988                s.n_rot,
4989                s.rope_base,
4990                if s.swa {
4991                    None
4992                } else {
4993                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4994                },
4995            ),
4996            None => (geometry.n_rot as usize, geometry.rope_base, None),
4997        };
4998        #[cfg(debug_assertions)]
4999        if let Some(ff) = ff {
5000            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5001        }
5002        match ff {
5003            Some(f) => e.rope_neox_ff(
5004                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5005            )?,
5006            None => e.rope_neox(
5007                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5008            )?,
5009        }
5010
5011        let kv = scratch.plane_mut(scratch_index).0;
5012        // Match the trunk prime contract: a chunk may need the aligned window immediately before
5013        // its first row, so preserve that prefix when the physical tail rebases at wrap.
5014        let retain_from = kv
5015            .ring
5016            .as_ref()
5017            .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5018            .unwrap_or(0);
5019        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5020        for i in 0..t {
5021            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5022            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5023            e.append_kv_quantized_view(
5024                &k_row,
5025                &v_row,
5026                &mut kv.k,
5027                &mut kv.v,
5028                write_row + i,
5029                kv.kv_dim_k,
5030                kv.kv_dim_v,
5031                kv.k_tok_bytes,
5032                kv.v_tok_bytes,
5033                false,
5034            )?;
5035        }
5036        kv.len = pos0 + t;
5037        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5038        Ok(())
5039    }
5040
5041    #[allow(clippy::too_many_arguments)]
5042    fn mtp_kv_fill_all(
5043        &self,
5044        e: &Engine,
5045        tokens: &[u32],
5046        h: &CudaSlice<f32>,
5047        pos0: usize,
5048        scratch: &mut MtpScratch,
5049        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5050    ) -> Result<(), Box<dyn std::error::Error>> {
5051        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5052        for index in 0..self.mtp_head_count() {
5053            self.mtp_kv_fill_at(
5054                e,
5055                self.mtp_head_at(index),
5056                tokens,
5057                h,
5058                pos0,
5059                scratch,
5060                index,
5061                embd_dev,
5062            )?;
5063        }
5064        Ok(())
5065    }
5066
5067    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5068    /// every varying input device-resident —
5069    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5070    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5071    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5072    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5073    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5074    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5075    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5076    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5077    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5078    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5079    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5080    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5081    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5082    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5083    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5084    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5085    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5086    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5087    #[allow(clippy::too_many_arguments)]
5088    fn mtp_head_forward_cap(
5089        &self,
5090        e: &Engine,
5091        mtp: &MtpHead,
5092        tok_d: &mut CudaSlice<u32>,
5093        pos_d: &mut CudaSlice<i32>,
5094        h_seed_d: &mut CudaSlice<f32>,
5095        p_d: &mut CudaSlice<f32>,
5096        scratch: &mut MtpScratch,
5097        with_prob: bool,
5098        with_head: bool,
5099        embd_gpu: &CudaSlice<u8>,
5100        embd_qt: i32,
5101        embd_rb: usize,
5102        d_vocab: usize,
5103        sampled_cap: Option<(
5104            &mut CudaSlice<u32>,
5105            &mut CudaSlice<f32>,
5106            &mut CudaSlice<f32>,
5107            u64,
5108            f32,
5109        )>,
5110        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5111        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5112        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5113        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5114        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5115        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5116        mask_cap: Option<(&CudaSlice<u32>, usize)>,
5117    ) -> Result<(), Box<dyn std::error::Error>> {
5118        let cfg = &self.cfg;
5119        let n_embd = cfg.n_embd as usize;
5120        // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5121        // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5122        // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5123        // row 0, cannot express this block's SWA view offset, and a captured chain would
5124        // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5125        // Returning Err (not a panic) is what the capture sites already handle by degrading to
5126        // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5127        // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5128        // step35_verify refusal), so a stream capture that succeeded here would only move the
5129        // failure from capture time (graceful stream-off) to serve time (a failed round).
5130        if let Some(g) = mtp.step35.as_ref() {
5131            if stream_pack.is_some() {
5132                return Err(
5133                    "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5134                     twin); stream off"
5135                        .into(),
5136                );
5137            }
5138            if !self.step35_dcw_eligible(g) {
5139                return Err(
5140                    "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5141                        block's SWA view offset; MEMRA_STEP35_DRAFT_DCW=1 arms the windowed dcw \
5142                        capture when the v3-vec class is live) - the eager draft chain serves \
5143                        this arch"
5144                        .into(),
5145                );
5146            }
5147        }
5148        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5149        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5150        let eps = cfg.rms_eps;
5151        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5152        let mut e_norm = e.zeros(n_embd)?;
5153        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5154        let mut h_norm = e.zeros(n_embd)?;
5155        e.rms_norm(
5156            &*h_seed_d,
5157            mtp.hnorm.float_data(),
5158            &mut h_norm,
5159            n_embd,
5160            1,
5161            eps,
5162        )?;
5163        let mut concat = e.zeros(2 * n_embd)?;
5164        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5165        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5166        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5167        let mut a_norm = e.zeros(di)?;
5168        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5169        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5170            // step35 (eligibility already enforced by the refusal above): the windowed dcw
5171            // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5172            // work here (this is the capture body); headroom is the callers' pre-arm.
5173            (Mixer::Full(fa), Some(g)) => {
5174                self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, 0)?
5175            }
5176            (Mixer::Full(fa), None) => {
5177                self.mtp_full_attn_dc(e, fa, &a_norm, pos_d, scratch, 0, mtp.geom.as_ref())?
5178            }
5179            (Mixer::Linear(_), _) => {
5180                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5181            }
5182            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
5183        };
5184        let mut x1 = e.zeros(di)?;
5185        e.add(&inp_sa, &attn_out, &mut x1, di)?;
5186        let mut z = e.zeros(di)?;
5187        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5188        let ffn_out = match &mtp.ffn {
5189            crate::hybrid::Ffn::Dense {
5190                ffn_gate,
5191                ffn_up,
5192                ffn_down,
5193            } => {
5194                let n_ff = ffn_gate.out_features();
5195                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5196                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5197                    (
5198                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5199                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5200                    )
5201                } else {
5202                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5203                };
5204                let mut act = e.zeros(n_ff)?;
5205                // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5206                // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5207                // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5208                // run the ONE activation program.
5209                Self::ffn_act_lim(
5210                    e,
5211                    &self.cfg,
5212                    &gate,
5213                    &up,
5214                    1.0,
5215                    1.0,
5216                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
5217                    &mut act,
5218                    n_ff,
5219                )?;
5220                e.matmul(ffn_down, &act, 1)?
5221            }
5222            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
5223            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
5224            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
5225            // error arm degrades the caller to eager/stream-off.
5226            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
5227                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5228            }
5229            crate::hybrid::Ffn::Moe(_) => {
5230                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
5231            }
5232        };
5233        let mut h_inner = e.zeros(di)?;
5234        e.add(&x1, &ffn_out, &mut h_inner, di)?;
5235        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5236        let h_nextn = match mtp.geom.as_ref() {
5237            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5238            None => h_inner,
5239        };
5240        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5241        let final_h = if with_head || spec_hpost() {
5242            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5243            let mut fh = e.zeros(n_embd)?;
5244            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5245            Some(fh)
5246        } else {
5247            None
5248        };
5249        if with_head {
5250            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5251            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
5252            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
5253            // before the argmax — proposals become legal by construction. Contents-only
5254            // per-replay upload keeps the capture valid.
5255            if let Some((mask_d, mw)) = mask_cap {
5256                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
5257            }
5258            if let Some((ctr_d, perturb_d, q_out_d, seed, temp)) = sampled_cap {
5259                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
5260                // own buffer is pool-recycled after the capture body returns, so it can't be the
5261                // retention target), bump the device event counter, gumbel-perturb reading it,
5262                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
5263                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
5264                e.sctr_inc(ctr_d)?;
5265                e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?;
5266                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
5267                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
5268                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
5269                if with_prob {
5270                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5271                }
5272            } else {
5273                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
5274                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
5275                // p-min under a draft mask reads the MASKED row: confidence relative to the
5276                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
5277                // is the right semantics for "does the drafter know what comes next here" and
5278                // the same row the pick came from. Draft-quality only — verify arbitrates.
5279                if with_prob {
5280                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5281                }
5282            }
5283        }
5284        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
5285        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
5286        if let Some((out, slot, d2t)) = stream_pack {
5287            e.pack_tok_p(tok_d, p_d, out, slot)?;
5288            if let Some(map) = d2t {
5289                e.tok_map_u32(tok_d, map)?;
5290            }
5291        }
5292        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
5293        if spec_hpost() {
5294            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
5295        } else {
5296            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
5297        }
5298        // advance the draft rope position in-graph.
5299        e.inc_seqlen(pos_d)?;
5300        Ok(())
5301    }
5302
5303    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
5304    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
5305    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
5306    /// Advances `cache.pos` by T.
5307    pub fn decode_step_t(
5308        &self,
5309        e: &Engine,
5310        tokens: &[u32],
5311        pos0: usize,
5312        cache: &mut Cache,
5313    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5314        if self.is_gemma4_e4b() {
5315            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
5316        }
5317        if self.gemma_batch_program() {
5318            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
5319        }
5320        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
5321    }
5322
5323    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
5324    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
5325    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
5326    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
5327    pub fn decode_step_t_h(
5328        &self,
5329        e: &Engine,
5330        tokens: &[u32],
5331        pos0: usize,
5332        cache: &mut Cache,
5333    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5334        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
5335    }
5336
5337    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
5338    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
5339    pub fn decode_step_t_h_emb(
5340        &self,
5341        e: &Engine,
5342        tokens: &[u32],
5343        pos0: usize,
5344        cache: &mut Cache,
5345        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5346    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5347        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
5348        Ok((e.dtoh(&logits_d)?, h_seed))
5349    }
5350
5351    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
5352    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
5353    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
5354    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
5355    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
5356    pub fn decode_step_t_h_emb_dev(
5357        &self,
5358        e: &Engine,
5359        tokens: &[u32],
5360        pos0: usize,
5361        cache: &mut Cache,
5362        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5363    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5364        let n_embd = self.cfg.n_embd as usize;
5365        let t = tokens.len();
5366        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
5367        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
5368        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
5369        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5370        Ok((logits, hs))
5371    }
5372
5373    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
5374    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
5375    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
5376    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
5377    /// retains/copies — they never change what any kernel computes).
5378    fn decode_step_t_core(
5379        &self,
5380        e: &Engine,
5381        tokens: &[u32],
5382        pos0: usize,
5383        cache: &mut Cache,
5384        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5385        mut ckpt: Option<&mut VerifyCkpt>,
5386    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5387        self.decode_step_t_core_stream(
5388            e,
5389            tokens,
5390            pos0,
5391            cache,
5392            embd_dev,
5393            ckpt.take(),
5394            None,
5395            None,
5396            None,
5397            None,
5398        )
5399    }
5400
5401    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
5402    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
5403    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
5404    fn decode_step_t_core_vg(
5405        &self,
5406        e: &Engine,
5407        tokens: &[u32],
5408        pos0: usize,
5409        cache: &mut Cache,
5410        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5411        mut ckpt: Option<&mut VerifyCkpt>,
5412        graphs: Option<&mut DsparkVerifyGraphs>,
5413    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5414        self.decode_step_t_core_stream(
5415            e,
5416            tokens,
5417            pos0,
5418            cache,
5419            embd_dev,
5420            ckpt.take(),
5421            None,
5422            None,
5423            None,
5424            graphs,
5425        )
5426    }
5427
5428    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
5429    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
5430    fn decode_step_t_core_pipelined(
5431        &self,
5432        e: &Engine,
5433        tokens: &[u32],
5434        pos0: usize,
5435        cache: &mut Cache,
5436        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5437        mut ckpt: Option<&mut VerifyCkpt>,
5438        pipe: &SpecPipeLane,
5439        round: usize,
5440    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5441        let fence = crate::pp::pp_cuts(self.layers.len())
5442            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
5443        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5444            return Err("two-session speculative pipeline requires the PP verify split".into());
5445        }
5446        let interval_fence = pipe.stage0_begin(round)?;
5447        let ticket = self.verify_stage0_issue(
5448            e,
5449            tokens,
5450            pos0,
5451            cache,
5452            embd_dev,
5453            ckpt.as_deref_mut(),
5454            None,
5455            &fence,
5456            Some(interval_fence),
5457            pipe.trace(round),
5458        )?;
5459        pipe.stage0_end(round);
5460        pipe.stage1_begin(round)?;
5461        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
5462        pipe.verify_end(round);
5463        Ok(result)
5464    }
5465
5466    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
5467    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
5468    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
5469    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
5470    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
5471    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
5472    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
5473    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
5474    #[allow(clippy::too_many_arguments)]
5475    fn decode_step_t_core_stream(
5476        &self,
5477        e: &Engine,
5478        tokens: &[u32],
5479        pos0: usize,
5480        cache: &mut Cache,
5481        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5482        mut ckpt: Option<&mut VerifyCkpt>,
5483        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5484        pp_pipe: Option<bool>,
5485        vtok_dev: Option<&CudaSlice<u32>>,
5486        graphs: Option<&mut DsparkVerifyGraphs>,
5487    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5488        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
5489        // exactly as the eager and batched steps do. This is the single funnel every verify
5490        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
5491        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
5492        // is untouched.
5493        //
5494        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
5495        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
5496        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
5497        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
5498        // or a placement whose PpNRt fails to build — so a config that would still walk the
5499        // whole trunk on one stream refuses instead of regressing 28x.
5500        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5501            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
5502                if vtok_dev.is_some() {
5503                    return Err(
5504                        "device-token dspark verify (slice-2 deferred readback) has no PP \
5505                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
5506                         route on one device"
5507                            .into(),
5508                    );
5509                }
5510                return self.decode_step_t_core_ppn(
5511                    e,
5512                    tokens,
5513                    pos0,
5514                    cache,
5515                    embd_dev,
5516                    ckpt.take(),
5517                    stream,
5518                    &fence,
5519                    pp_pipe,
5520                );
5521            }
5522        }
5523        crate::pp::refuse_unsplit_if_remote(
5524            "decode_step_t (spec verify)",
5525            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
5526             split (decode_step_t_core_ppn); or run spec on one device",
5527        )?;
5528        let cfg = &self.cfg;
5529        let n_embd = cfg.n_embd as usize;
5530        let eps = cfg.rms_eps;
5531        let t = tokens.len();
5532        let pos_d = match stream {
5533            Some((_, ctr)) => {
5534                let mut p = e.alloc_uninit::<i32>(t)?;
5535                e.pos_iota(ctr, &mut p, t)?;
5536                p
5537            }
5538            None => {
5539                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5540                e.htod_i32(&pos_vec)?
5541            }
5542        };
5543
5544        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5545        let x = match (stream, embd_dev) {
5546            (Some((vtok, _)), Some((g, qt, rb))) => {
5547                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5548            }
5549            (None, Some((g, qt, rb))) => match vtok_dev {
5550                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5551                // bit-identical rows to the host-token arm (same per-dtype deq).
5552                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5553                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5554            },
5555            _ => {
5556                assert!(
5557                    vtok_dev.is_none(),
5558                    "device-token verify requires the resident embed table (embd_dev)"
5559                );
5560                e.htod(&self.embd.gather(n_embd, tokens))?
5561            }
5562        };
5563
5564        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5565        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5566        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5567        let x = self.verify_layers(
5568            e,
5569            x,
5570            0,
5571            self.layers.len(),
5572            &pos_d,
5573            pos0,
5574            t,
5575            cache,
5576            ckpt.take(),
5577            stream,
5578            graphs,
5579        )?;
5580        if spec_nan_scan() {
5581            nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5582        }
5583
5584        let mut hn = vbuf(e, t * n_embd)?;
5585        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5586        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5587        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5588        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5589        let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5590        if eager_tail {
5591            let n_vocab = self.cfg.n_vocab as usize;
5592            // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5593            //
5594            // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5595            // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5596            // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5597            // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5598            //
5599            // The loop's justification is the comment above: the batched cuBLASLt head is a
5600            // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5601            // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5602            // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5603            // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5604            // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5605            // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5606            // documented "bit-identical to t single-row calls". So the batched form is the SAME
5607            // arithmetic per row on both paths, with one weight read instead of t.
5608            //
5609            // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5610            //
5611            // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5612            // claims" is still an argument. The greedy byte tape decides, and the door flips only
5613            // once the tape is a receipt.
5614            if head_rows_on() {
5615                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5616                let logits = e.matmul(&self.output, &hn, t)?;
5617                if stream.is_none() {
5618                    cache.pos += t;
5619                }
5620                return Ok((logits, if spec_hpost() { hn } else { x }));
5621            }
5622            let mut logits = vbuf(e, t * n_vocab)?;
5623            for r in 0..t {
5624                let mut row = e.uninit(n_embd)?;
5625                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5626                let mut hr = e.uninit(n_embd)?;
5627                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5628                let lr = e.matmul(&self.output, &hr, 1)?;
5629                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5630                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5631            }
5632            if stream.is_none() {
5633                cache.pos += t;
5634            }
5635            return Ok((logits, if spec_hpost() { hn } else { x }));
5636        }
5637        let serving_head =
5638            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5639        let logits = if serving_head {
5640            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5641            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5642            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5643            // serve one batched numeric class at every live width, including B=1. Keep the
5644            // verify head in that same class; other generic families retain the decode-exact
5645            // head that their run-spec contract pins.
5646            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5647            e.matmul(&self.output, &hn, t)?
5648        } else {
5649            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5650            e.matmul_decode_exact(&self.output, &hn, t)?
5651        };
5652        // stream: the device pos counter owns position; host mirror reconciles at drain.
5653        if stream.is_none() {
5654            cache.pos += t;
5655        }
5656        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5657        Ok((logits, if spec_hpost() { hn } else { x }))
5658    }
5659
5660    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5661    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5662    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5663    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5664    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5665    /// the payload).
5666    ///
5667    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5668    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5669    /// receipts):
5670    ///
5671    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5672    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5673    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5674    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5675    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
5676    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5677    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5678    ///
5679    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5680    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5681    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5682    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5683    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
5684    ///
5685    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5686    ///    sharded loader leaves the table with stage 0 by construction).
5687    ///
5688    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5689    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5690    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5691    ///    model, every round.
5692    ///
5693    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5694    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5695    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5696    /// through the primary context by UVA — the same read the batched serving epilogue's
5697    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5698    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5699    ///
5700    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5701    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5702    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5703    ///
5704    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5705    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5706    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5707    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5708    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5709    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5710    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5711    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5712    #[allow(clippy::too_many_arguments)]
5713    fn decode_step_t_core_ppn(
5714        &self,
5715        e: &Engine,
5716        tokens: &[u32],
5717        pos0: usize,
5718        cache: &mut Cache,
5719        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5720        mut ckpt: Option<&mut VerifyCkpt>,
5721        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5722        fence: &[usize],
5723        pp_pipe: Option<bool>,
5724    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5725        let ticket = self.verify_stage0_issue(
5726            e,
5727            tokens,
5728            pos0,
5729            cache,
5730            embd_dev,
5731            ckpt.as_deref_mut(),
5732            stream,
5733            fence,
5734            pp_pipe,
5735            None,
5736        )?;
5737        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5738    }
5739
5740    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5741    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5742    #[allow(clippy::too_many_arguments)]
5743    fn verify_stage0_issue(
5744        &self,
5745        e: &Engine,
5746        tokens: &[u32],
5747        pos0: usize,
5748        cache: &mut Cache,
5749        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5750        mut ckpt: Option<&mut VerifyCkpt>,
5751        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5752        fence: &[usize],
5753        pp_pipe: Option<bool>,
5754        trace: Option<SpecPipeTraceCtx>,
5755    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5756        assert!(
5757            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5758            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5759             (the gemma4 arms have their own decode_step_t twins)"
5760        );
5761        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5762            return Err(
5763                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5764                 boundary itself is host-staged, but device-resident verify still peer-reads \
5765                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5766                 serving on this host class; spec requires local per-stage inputs first."
5767                    .into(),
5768            );
5769        }
5770        let rt = crate::pp::PpNRt::get(e)?;
5771        let n_st = fence.len() - 1;
5772        assert_eq!(
5773            rt.n_stages(),
5774            n_st,
5775            "PpNRt stage count {} != fence stages {n_st}",
5776            rt.n_stages()
5777        );
5778        let n_embd = self.cfg.n_embd as usize;
5779        let t = tokens.len();
5780        let payload = t * n_embd;
5781        if pp_pipe.is_some() {
5782            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5783        }
5784        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5785        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5786        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5787        // the report below names exactly two stages and must never imply it measured middle ones.
5788        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5789        let pp_started = std::time::Instant::now();
5790        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5791        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5792        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5793        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5794        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5795        // stage stream and the wait would self-order into a no-op.
5796        let caller_stream = e.stream();
5797        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5798        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5799        // the primary stream still holds queued reads of them — with event tracking elided,
5800        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5801        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5802        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5803        // stage stream behind the caller before enqueueing new stage work.
5804        let reverse_started = std::time::Instant::now();
5805        if pp_pipe != Some(false) {
5806            rt.fence_stages_behind(&caller_stream)?;
5807        }
5808        if pp_pipe == Some(true) {
5809            // Both session verifies must alternate boundary slots even when the ordinary
5810            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5811            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5812            rt.prepare_overlap_slots(0, payload)?;
5813        }
5814        if pp_anatomy {
5815            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5816            // prices any primary-stream rollback/refresh tail inherited from the prior round.
5817            for s in 0..n_st {
5818                let _st = rt.enter(s);
5819                rt.engine(s, e).stream().synchronize()?;
5820            }
5821            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5822        }
5823
5824        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5825        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5826        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5827            match stream {
5828                Some((_, ctr)) => {
5829                    let mut p = es.alloc_uninit::<i32>(t)?;
5830                    es.pos_iota(ctr, &mut p, t)?;
5831                    Ok(p)
5832                }
5833                None => {
5834                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5835                    es.htod_i32(&pos_vec)
5836                }
5837            }
5838        };
5839
5840        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5841        let slot = {
5842            let _st0 = rt.enter(0);
5843            let e0 = rt.engine(0, e);
5844            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5845            let stage0_started = std::time::Instant::now();
5846            let pos_d = stage_pos(e0)?;
5847            let x = match (stream, embd_dev) {
5848                (Some((vtok, _)), Some((g, qt, rb))) => {
5849                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5850                }
5851                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5852                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
5853            };
5854            let x = self.verify_layers(
5855                e0,
5856                x,
5857                fence[0],
5858                fence[1],
5859                &pos_d,
5860                pos0,
5861                t,
5862                cache,
5863                ckpt.as_deref_mut(),
5864                stream,
5865                None,
5866            )?;
5867            if pp_anatomy {
5868                e0.stream().synchronize()?;
5869                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
5870            }
5871            let tx_started = std::time::Instant::now();
5872            let slot = if pp_pipe.is_some() {
5873                rt.tx_pipelined(0, &x, payload)?
5874            } else {
5875                rt.tx(0, &x, payload)?
5876            };
5877            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
5878            if pp_anatomy {
5879                e0.stream().synchronize()?;
5880                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
5881            }
5882            slot
5883            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
5884        };
5885
5886        Ok(VerifyBoundaryTicket {
5887            rt,
5888            caller_stream,
5889            slot,
5890            pos0,
5891            t,
5892            payload,
5893            n_st,
5894            pipelined: pp_pipe.is_some(),
5895            pp_anatomy,
5896            pp_started,
5897            reverse_ms,
5898            stage0_ms,
5899            tx_ms,
5900            trace,
5901        })
5902    }
5903
5904    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
5905    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
5906    #[allow(clippy::too_many_arguments)]
5907    fn verify_stage1_finish(
5908        &self,
5909        e: &Engine,
5910        ticket: VerifyBoundaryTicket,
5911        cache: &mut Cache,
5912        mut ckpt: Option<&mut VerifyCkpt>,
5913        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5914        fence: &[usize],
5915        publish_to_caller: bool,
5916    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5917        let VerifyBoundaryTicket {
5918            rt,
5919            caller_stream,
5920            slot,
5921            pos0,
5922            t,
5923            payload,
5924            n_st,
5925            pipelined,
5926            pp_anatomy,
5927            pp_started,
5928            reverse_ms,
5929            stage0_ms,
5930            tx_ms,
5931            trace,
5932        } = ticket;
5933        let n_embd = self.cfg.n_embd as usize;
5934        let eps = self.cfg.rms_eps;
5935        let mut slot = slot;
5936        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
5937        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5938            match stream {
5939                Some((_, ctr)) => {
5940                    let mut p = es.alloc_uninit::<i32>(t)?;
5941                    es.pos_iota(ctr, &mut p, t)?;
5942                    Ok(p)
5943                }
5944                None => {
5945                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5946                    es.htod_i32(&pos_vec)
5947                }
5948            }
5949        };
5950
5951        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
5952        for s in 1..n_st - 1 {
5953            let _st = rt.enter(s);
5954            let es = rt.engine(s, e);
5955            let pos_d = stage_pos(es)?;
5956            let x = rt.rx(s - 1, slot, payload)?;
5957            let x = self.verify_layers(
5958                es,
5959                x,
5960                fence[s],
5961                fence[s + 1],
5962                &pos_d,
5963                pos0,
5964                t,
5965                cache,
5966                ckpt.as_deref_mut(),
5967                stream,
5968                None,
5969            )?;
5970            slot = if pipelined {
5971                rt.tx_pipelined(s, &x, payload)?
5972            } else {
5973                rt.tx(s, &x, payload)?
5974            };
5975        }
5976
5977        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
5978        let _stl = rt.enter(n_st - 1);
5979        let el = rt.engine(n_st - 1, e);
5980        let pos_d = stage_pos(el)?;
5981        let rx_started = std::time::Instant::now();
5982        let x = rt.rx(n_st - 2, slot, payload)?;
5983        if pp_anatomy {
5984            el.stream().synchronize()?;
5985            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
5986        }
5987        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
5988        let stage1_started = std::time::Instant::now();
5989        let x = self.verify_layers(
5990            el,
5991            x,
5992            fence[n_st - 1],
5993            fence[n_st],
5994            &pos_d,
5995            pos0,
5996            t,
5997            cache,
5998            ckpt.as_deref_mut(),
5999            stream,
6000            None,
6001        )?;
6002
6003        let mut hn = vbuf(el, payload)?;
6004        let logits = if self.sliding_gated_moe_batch_program() {
6005            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6006            // Verify must not switch numeric class merely because the same session speculates.
6007            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6008            el.matmul(&self.output, &hn, t)?
6009        } else {
6010            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6011            el.matmul_decode_exact(&self.output, &hn, t)?
6012        };
6013        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6014        if pp_anatomy {
6015            el.stream().synchronize()?;
6016            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6017        }
6018        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6019        // stream. Order the caller's stream behind that work before the buffers escape this
6020        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6021        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6022        // the following arm's KV in the same process).
6023        if publish_to_caller {
6024            rt.publish_to(n_st - 1, &caller_stream)?;
6025        }
6026        if pp_anatomy {
6027            if publish_to_caller {
6028                caller_stream.synchronize()?;
6029            }
6030            eprintln!(
6031                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6032                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6033                pp_started.elapsed().as_secs_f64() * 1e3,
6034            );
6035        }
6036        // stream: the device pos counter owns position; host mirror reconciles at drain.
6037        if stream.is_none() {
6038            cache.pos += t;
6039        }
6040        Ok((logits, if spec_hpost() { hn } else { x }))
6041    }
6042
6043    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6044    ///
6045    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6046    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6047    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6048    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6049    /// bytes when a request moves from batched plain serving into speculative verify. Run the
6050    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6051    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6052    /// every norm/projection/FFN uses exactly the live serving dispatch.
6053    #[allow(clippy::too_many_arguments)]
6054    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6055    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6056    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6057    /// reference while replacing the host-canonical per-token prime. Requires the walk
6058    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6059    #[allow(clippy::type_complexity)]
6060    pub(crate) fn step35_prime_trows(
6061        &self,
6062        e: &Engine,
6063        tokens: &[u32],
6064        cache: &mut Cache,
6065    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6066    {
6067        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6068        if !prime_trows_on() {
6069            return Ok(None);
6070        }
6071        if !self.uses_sliding_gated_moe_program()
6072            || cache.pos != 0
6073            || cache.dflash_taps.is_some()
6074            || !spec_verify_eager_on()
6075            || !spec_verify_tcol_on()
6076        {
6077            if dbg {
6078                eprintln!(
6079                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6080                    self.uses_sliding_gated_moe_program(),
6081                    cache.pos,
6082                    cache.dflash_taps.is_some(),
6083                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6084                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6085                );
6086            }
6087            return Ok(None);
6088        }
6089        let n_embd = self.cfg.n_embd as usize;
6090        let n_layers = self.layers.len();
6091        let t_total = tokens.len();
6092        let Some(embd_gpu) = self.embd_gpu_try(e) else {
6093            if dbg {
6094                eprintln!("[prime-trows] refuse: no device embed table");
6095            }
6096            return Ok(None);
6097        };
6098        let embd_qtype = match self.embd.ggml_type {
6099            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6100            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6101            other => {
6102                if dbg {
6103                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6104                }
6105                return Ok(None);
6106            }
6107        };
6108        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6109        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6110        // (the walk floor is t >= 2).
6111        let mut bounds = Vec::new();
6112        let mut start = 0usize;
6113        while start < t_total {
6114            let mut end = (start + 32).min(t_total);
6115            if t_total - end == 1 {
6116                end -= 1;
6117            }
6118            bounds.push((start, end));
6119            start = end;
6120        }
6121        if bounds.iter().any(|(a, b)| b - a < 2) {
6122            return Ok(None); // degenerate short prompt keeps the ordinary prime
6123        }
6124        let mut hiddens = e.uninit(t_total * n_embd)?;
6125        let mut last: Option<CudaSlice<f32>> = None;
6126        for &(a, b) in &bounds {
6127            let tc = b - a;
6128            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6129            let x =
6130                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6131            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6132            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6133            if b == t_total {
6134                let mut h = e.uninit(n_embd)?;
6135                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6136                last = Some(h);
6137            }
6138        }
6139        let h_seed = last.expect("last chunk produced the seed row");
6140        let mut hn = e.uninit(n_embd)?;
6141        e.rms_norm_decode(
6142            &h_seed,
6143            self.output_norm.float_data(),
6144            &mut hn,
6145            n_embd,
6146            1,
6147            self.cfg.rms_eps,
6148        )?;
6149        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6150        let logits = e.dtoh(&logits_d)?;
6151        cache.pos = t_total;
6152        Ok(Some((logits, h_seed, hiddens)))
6153    }
6154
6155    fn step35_verify_batch_layers(
6156        &self,
6157        e: &Engine,
6158        mut x: CudaSlice<f32>,
6159        lo: usize,
6160        hi: usize,
6161        pos0: usize,
6162        t: usize,
6163        cache: &mut Cache,
6164    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6165        let n_embd = self.cfg.n_embd as usize;
6166        if !self.uses_sliding_gated_moe_program() {
6167            return Err(
6168                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6169            );
6170        }
6171        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6172        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6173        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6174        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6175        // and the tap path keep the batch-layer class.
6176        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6177        let eager_verify =
6178            *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6179        if eager_verify {
6180            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6181            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6182            // column runs the UNMODIFIED t=1 attention program via the col-select door and
6183            // the ordinary residual/FFN body. Values per column are bit-equal to the
6184            // row-outer walk: rms over the materialized residual == the fused add+norm
6185            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6186            // kernel, and every downstream op IS the t=1 program.
6187            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6188            let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6189            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6190            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6191            // so a chunked call is value-identical to the row-outer loop it replaces.
6192            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6193            // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6194            // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6195            // flag precedence between two existing doors, not a new flag. Without this, both
6196            // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6197            let trows_prefill =
6198                *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6199            // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6200            // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6201            // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6202            // its accumulators to local memory), so a wider chunk fails the request with
6203            // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6204            // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6205            static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6206            let trows_w = match TROWS_W.get_or_init(|| {
6207                let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6208                parse_prime_trows_width(value.as_deref())
6209            }) {
6210                Ok(width) => *width,
6211                Err(err) => return Err(err.clone().into()),
6212            };
6213            if tcol && trows_prefill && t > trows_w {
6214                // One-time engagement receipt: without it a prefill gate cannot tell a
6215                // chunked walk from the row-outer fallback it is supposed to replace
6216                // (the first PRIME_TROWS gate passed vacuously on exactly that).
6217                static SEEN: std::sync::atomic::AtomicBool =
6218                    std::sync::atomic::AtomicBool::new(false);
6219                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6220                    eprintln!(
6221                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6222                        t.div_ceil(trows_w),
6223                        lo,
6224                        hi
6225                    );
6226                }
6227                let mut out = e.uninit(t * n_embd)?;
6228                let mut start = 0usize;
6229                while start < t {
6230                    let mut end = (start + trows_w).min(t);
6231                    if t - end == 1 {
6232                        end -= 1;
6233                    }
6234                    let tc = end - start;
6235                    let mut xc = e.uninit(tc * n_embd)?;
6236                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
6237                    let oc =
6238                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
6239                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
6240                    start = end;
6241                }
6242                return Ok(out);
6243            }
6244            if tcol && t >= 2 && t <= 32 {
6245                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
6246                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
6247                // syncs serialize the stream, so the split is for TARGETING amortization
6248                // work only — never a perf claim.
6249                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6250                let prof =
6251                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
6252                let mut prof_ms = [0f64; 3];
6253                let eps = self.cfg.rms_eps;
6254                let mut x_t = x;
6255                let mut h_t = e.uninit(t * n_embd)?;
6256                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
6257                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
6258                // pageable htod was an in-stream engine turnaround x t x 45).
6259                let mut pos_rows = Vec::with_capacity(t);
6260                for r in 0..t {
6261                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
6262                }
6263                let mut ok = true;
6264                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
6265                // stashes `gated` instead of joining per column; one b4_tcol per rank +
6266                // one slab join produce every column's `mixed` after the attention pass.
6267                // Bit-exact per column (t=1 b4 program per column; elementwise join).
6268                // MEMRA_TCOL_FFN=1 (implies the o_proj defer): when every column of a
6269                // MoE layer deferred, the residual norm runs as one t-grid launch
6270                // (per-row program == t=1) and the FFN as ONE two-column device-routed
6271                // sweep + per-column shexp — the two columns' expert weights dedup
6272                // through L2 instead of reading HBM twice.
6273                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6274                let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
6275                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
6276                // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
6277                // the per-column pass norms/ropes/appends and stashes q+gate, then one
6278                // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
6279                // [2, o_out] mixed slab. The precheck runs before arming (stashing is
6280                // unrecoverable); ineligible/boundary layers run the ordinary program.
6281                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
6282                let mut mixed_row = e.uninit(n_embd)?;
6283                let mut pos_staged = false;
6284                for il in lo..hi {
6285                    let layer = &self.layers[il];
6286                    // BEFORE this layer touches its planes: is the history it is about to
6287                    // attend already poisoned? Global (non-ring) layers only, which are the
6288                    // ones the level-2 bitmap implicates.
6289                    if kv_plane_scan_on() && self.step35_geom(il).window.is_none() {
6290                        if let Some(distributed) = cache.tp_kv[il].as_ref() {
6291                            scan_kv_plane(e, distributed, il, pos0)?;
6292                        }
6293                    }
6294                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
6295                    let mut seg = std::time::Instant::now();
6296                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
6297                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
6298                        ok = false;
6299                        break;
6300                    }
6301                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
6302                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
6303                    // advance by t. Host cache bookkeeping mirrors the per-column tail.
6304                    if fa2_layer {
6305                        if let Some(mixed_t) =
6306                            self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
6307                        {
6308                            pos_staged = true;
6309                            {
6310                                let tp_kv = cache.tp_kv[il]
6311                                    .as_mut()
6312                                    .expect("precheck verified the distributed cache");
6313                                let transaction = tp_kv.begin_transaction()?;
6314                                let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
6315                                    return Err("verify rope pass expects full attention".into());
6316                                };
6317                                let tp = fa
6318                                    .step_tp_qkv
6319                                    .as_ref()
6320                                    .ok_or("verify rope pass lost its TP state")?;
6321                                let empty: [CudaSlice<f32>; 0] = [];
6322                                tp.runtime.append_tp_kv_transaction_inner(
6323                                    tp_kv,
6324                                    transaction,
6325                                    &empty,
6326                                    &empty,
6327                                    t,
6328                                    true,
6329                                )?;
6330                                tp.runtime.commit_tp_kv_transaction_external(
6331                                    tp_kv,
6332                                    transaction,
6333                                    t,
6334                                )?;
6335                                if let Some(local) = cache.kv[il].as_mut() {
6336                                    local.len = pos0 + t;
6337                                    if !crate::tp::len_mirror_lazy_on() {
6338                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
6339                                    }
6340                                }
6341                            }
6342                            if prof {
6343                                e.stream().synchronize()?;
6344                                prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6345                                seg = std::time::Instant::now();
6346                            }
6347                            let o_out = mixed_t.len() / t;
6348                            let mut next = e.uninit(t * n_embd)?;
6349                            let mut batched = false;
6350                            if ffn_batch && o_out == n_embd {
6351                                // MEMRA_WALK_SCRATCH=1: reuse persistent slabs instead of four
6352                                // fresh allocations per layer (see WalkScratch). Same kernels,
6353                                // same order, same values — only the buffers' provenance changes.
6354                                if walk_scratch_on() {
6355                                    let mut guard = WALK_SCRATCH
6356                                        .lock()
6357                                        .map_err(|_| "walk scratch lock poisoned")?;
6358                                    let dev = e.ctx().ordinal();
6359                                    if guard
6360                                        .as_ref()
6361                                        .is_none_or(|w| w.dev != dev || w.cap < t * n_embd)
6362                                    {
6363                                        *guard = Some(WalkScratch {
6364                                            dev,
6365                                            cap: 32 * n_embd,
6366                                            x1: e.uninit(32 * n_embd)?,
6367                                            z: e.uninit(32 * n_embd)?,
6368                                        });
6369                                    }
6370                                    let w = guard.as_mut().expect("armed above");
6371                                    // `add_rms_norm` and `step35_verify_moe_tn` both take an
6372                                    // explicit element count and tolerate a longer slab, so the
6373                                    // persistent buffers drop straight in: same kernels, same
6374                                    // order, same values, two fewer allocations per layer.
6375                                    e.add_rms_norm(
6376                                        &x_t,
6377                                        &mixed_t,
6378                                        layer.post_attn_norm.float_data(),
6379                                        &mut w.x1,
6380                                        &mut w.z,
6381                                        n_embd,
6382                                        t,
6383                                        eps,
6384                                    )?;
6385                                    if spec_nan_scan_level() >= 2 {
6386                                        nan_scan_rows(
6387                                            e,
6388                                            &w.z,
6389                                            t,
6390                                            n_embd,
6391                                            &format!("tcol layer {il} post-attn norm z"),
6392                                        )?;
6393                                    }
6394                                    if let Some(ffn_t) =
6395                                        self.step35_verify_moe_tn(e, il, &w.z, t)?
6396                                    {
6397                                        if spec_nan_scan_level() >= 2 {
6398                                            nan_scan_rows(
6399                                                e,
6400                                                &ffn_t,
6401                                                t,
6402                                                n_embd,
6403                                                &format!("tcol layer {il} batched routed-MoE out"),
6404                                            )?;
6405                                        }
6406                                        let mut x2_t = e.uninit(t * n_embd)?;
6407                                        e.add(&w.x1, &ffn_t, &mut x2_t, t * n_embd)?;
6408                                        next = x2_t;
6409                                        batched = true;
6410                                    }
6411                                } else {
6412                                    let mut x1_t = e.uninit(t * n_embd)?;
6413                                    let mut z_t = e.uninit(t * n_embd)?;
6414                                    e.add_rms_norm(
6415                                        &x_t,
6416                                        &mixed_t,
6417                                        layer.post_attn_norm.float_data(),
6418                                        &mut x1_t,
6419                                        &mut z_t,
6420                                        n_embd,
6421                                        t,
6422                                        eps,
6423                                    )?;
6424                                    if spec_nan_scan_level() >= 2 {
6425                                        nan_scan_rows(
6426                                            e,
6427                                            &z_t,
6428                                            t,
6429                                            n_embd,
6430                                            &format!("tcol layer {il} post-attn norm z"),
6431                                        )?;
6432                                    }
6433                                    if let Some(ffn_t) =
6434                                        self.step35_verify_moe_tn(e, il, &z_t, t)?
6435                                    {
6436                                        if spec_nan_scan_level() >= 2 {
6437                                            nan_scan_rows(
6438                                                e,
6439                                                &ffn_t,
6440                                                t,
6441                                                n_embd,
6442                                                &format!("tcol layer {il} batched routed-MoE out"),
6443                                            )?;
6444                                        }
6445                                        let mut x2_t = e.uninit(t * n_embd)?;
6446                                        e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
6447                                        next = x2_t;
6448                                        batched = true;
6449                                    }
6450                                }
6451                            }
6452                            if !batched {
6453                                for r in 0..t {
6454                                    e.dtod_copy_view(
6455                                        &mixed_t.slice(r * o_out..(r + 1) * o_out),
6456                                        &mut mixed_row,
6457                                    )?;
6458                                    let mut x_row = e.uninit(n_embd)?;
6459                                    e.dtod_copy_view(
6460                                        &x_t.slice(r * n_embd..(r + 1) * n_embd),
6461                                        &mut x_row,
6462                                    )?;
6463                                    let (x1, ffn_out) = self.residual_norm_ffn(
6464                                        e, layer, &x_row, &mixed_row, n_embd, il, eps,
6465                                    )?;
6466                                    let mut x2 = e.uninit(n_embd)?;
6467                                    e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6468                                    e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
6469                                }
6470                            }
6471                            if prof {
6472                                e.stream().synchronize()?;
6473                                prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6474                            }
6475                            x_t = next;
6476                            if spec_nan_scan() {
6477                                // The scan MUST sit on this arm too. It used to live only on
6478                                // the non-fused tail, so a fused layer's poison was first
6479                                // reported by the next non-fused layer.
6480                                verify_arm_receipt(
6481                                    "fused",
6482                                    il,
6483                                    pos0,
6484                                    t,
6485                                    cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6486                                );
6487                                nan_scan_rows(
6488                                    e,
6489                                    &x_t,
6490                                    t,
6491                                    n_embd,
6492                                    &format!("tcol layer {il} pos0={pos0} arm=fused"),
6493                                )?;
6494                            }
6495                            continue;
6496                        }
6497                    }
6498                    if prof {
6499                        e.stream().synchronize()?;
6500                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
6501                        seg = std::time::Instant::now();
6502                    }
6503                    let mut next = e.uninit(t * n_embd)?;
6504                    // Columns whose o_proj was deferred (their FFN runs after the join).
6505                    // A NON-deferred column's FFN must run INSIDE the column loop: the
6506                    // oproj-tail handoff is a single cell that the same column's
6507                    // residual_norm_ffn consumes before the next column's finish.
6508                    let mut deferred: Vec<usize> = Vec::new();
6509                    let mut fa2_deferred: Vec<usize> = Vec::new();
6510                    let mut ffn_col =
6511                        |r: usize,
6512                         mixed: &CudaSlice<f32>,
6513                         next: &mut CudaSlice<f32>|
6514                         -> Result<(), Box<dyn std::error::Error>> {
6515                            let mut x_row = e.uninit(n_embd)?;
6516                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
6517                            let (x1, ffn_out) =
6518                                self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
6519                            if spec_nan_scan_level() >= 2 {
6520                                nan_scan_rows(
6521                                    e,
6522                                    &ffn_out,
6523                                    1,
6524                                    n_embd,
6525                                    &format!("tcol layer {il} col {r} per-column FFN out"),
6526                                )?;
6527                            }
6528                            let mut x2 = e.uninit(n_embd)?;
6529                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6530                            e.dtod_copy_into(&x2, next, r * n_embd)?;
6531                            Ok(())
6532                        };
6533                    for r in 0..t {
6534                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
6535                        let row_pos = &pos_rows[r];
6536                        crate::tp::set_verify_tcol(Some(r));
6537                        if fa2_layer {
6538                            crate::tp::set_spec_fa2_defer(Some(r));
6539                        } else if oproj_batch {
6540                            crate::tp::set_tcol_oproj_defer(Some(r));
6541                        }
6542                        let mixed = match &layer.mixer {
6543                            crate::hybrid::Mixer::Full(fa) => {
6544                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
6545                            }
6546                            _ => Err("step35 verify expects full attention".into()),
6547                        };
6548                        crate::tp::set_verify_tcol(None);
6549                        crate::tp::set_spec_fa2_defer(None);
6550                        crate::tp::set_tcol_oproj_defer(None);
6551                        let mixed = mixed?;
6552                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
6553                            fa2_deferred.push(r);
6554                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
6555                            deferred.push(r);
6556                        } else {
6557                            if spec_nan_scan_level() >= 2 {
6558                                let cols = mixed.len();
6559                                nan_scan_rows(
6560                                    e,
6561                                    &mixed,
6562                                    1,
6563                                    cols,
6564                                    &format!("tcol layer {il} col {r} per-column ATTN out"),
6565                                )?;
6566                            }
6567                            ffn_col(r, &mixed, &mut next)?;
6568                        }
6569                    }
6570                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
6571                        // The precheck guarantees both columns stash or neither; a strict
6572                        // subset means a column's output was never produced anywhere.
6573                        return Err("spec fa2 stash engaged for a subset of columns".into());
6574                    }
6575                    if prof {
6576                        e.stream().synchronize()?;
6577                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6578                        seg = std::time::Instant::now();
6579                    }
6580                    if !fa2_deferred.is_empty() {
6581                        deferred = fa2_deferred;
6582                    }
6583                    if !deferred.is_empty() {
6584                        let mixed_t = if fa2_layer {
6585                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
6586                        } else {
6587                            self.step35_verify_oproj_tcol(e, il, t)?
6588                        };
6589                        let o_out = mixed_t.len() / t;
6590                        if spec_nan_scan_level() >= 2 {
6591                            nan_scan_rows(
6592                                e,
6593                                &mixed_t,
6594                                t,
6595                                o_out,
6596                                &format!("tcol layer {il} JOINED attn over deferred cols"),
6597                            )?;
6598                        }
6599                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
6600                        // program == t=1; bit-identical to the oproj-tail join per the
6601                        // M2 verbatim-program contract) feeding the two-column routed
6602                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
6603                        // to the per-column body.
6604                        let mut batched = false;
6605                        if ffn_batch && deferred.len() == t && o_out == n_embd {
6606                            let mut x1_t = e.uninit(t * n_embd)?;
6607                            let mut z_t = e.uninit(t * n_embd)?;
6608                            e.add_rms_norm(
6609                                &x_t,
6610                                &mixed_t,
6611                                layer.post_attn_norm.float_data(),
6612                                &mut x1_t,
6613                                &mut z_t,
6614                                n_embd,
6615                                t,
6616                                eps,
6617                            )?;
6618                            if let Some(ffn_t) = self.step35_verify_moe_tn(e, il, &z_t, t)? {
6619                                let mut x2_t = e.uninit(t * n_embd)?;
6620                                e.add(&x1_t, &ffn_t, &mut x2_t, t * n_embd)?;
6621                                next = x2_t;
6622                                batched = true;
6623                            }
6624                        }
6625                        if !batched {
6626                            for &r in &deferred {
6627                                e.dtod_copy_view(
6628                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
6629                                    &mut mixed_row,
6630                                )?;
6631                                ffn_col(r, &mixed_row, &mut next)?;
6632                            }
6633                        }
6634                    }
6635                    if prof {
6636                        e.stream().synchronize()?;
6637                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6638                    }
6639                    drop(ffn_col);
6640                    x_t = next;
6641                    if spec_nan_scan() {
6642                        verify_arm_receipt(
6643                            if fa2_layer { "join" } else { "percol" },
6644                            il,
6645                            pos0,
6646                            t,
6647                            cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6648                        );
6649                        nan_scan_rows(
6650                            e,
6651                            &x_t,
6652                            t,
6653                            n_embd,
6654                            &format!(
6655                                "tcol layer {il} pos0={pos0} arm={}",
6656                                if fa2_layer { "join" } else { "percol" }
6657                            ),
6658                        )?;
6659                    }
6660                }
6661                if prof {
6662                    eprintln!(
6663                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6664                        prof_ms[0], prof_ms[1], prof_ms[2]
6665                    );
6666                }
6667                if ok {
6668                    return Ok(x_t);
6669                }
6670                // fall through to the row-outer walk on ineligible layers
6671                x = x_t;
6672            }
6673            let mut next = e.uninit(t * n_embd)?;
6674            let scan = spec_nan_scan();
6675            for r in 0..t {
6676                let mut row = e.uninit(n_embd)?;
6677                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6678                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6679                let out = if scan {
6680                    // Diagnostic arm: the same range walked one layer at a time so the first
6681                    // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6682                    // and executes its trailing residual add, so a per-layer chain is the same
6683                    // program with the cross-layer add+norm fusion unrolled.
6684                    nan_scan_rows(
6685                        e,
6686                        &row,
6687                        1,
6688                        n_embd,
6689                        &format!("embed row r={r} pos={}", pos0 + r),
6690                    )?;
6691                    let mut acc = row;
6692                    for il in lo..hi {
6693                        acc = self.decode_layers_eager(
6694                            e,
6695                            acc,
6696                            il,
6697                            il + 1,
6698                            &row_pos,
6699                            pos0 + r,
6700                            cache,
6701                        )?;
6702                        nan_scan_rows(
6703                            e,
6704                            &acc,
6705                            1,
6706                            n_embd,
6707                            &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6708                        )?;
6709                    }
6710                    acc
6711                } else {
6712                    self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6713                };
6714                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6715            }
6716            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6717            // row-outer walk does not materialize); the door is a step37 MTP bring-up
6718            // surface where taps are unused.
6719            return Ok(next);
6720        }
6721        let mut ph_last = std::time::Instant::now();
6722        for il in lo..hi {
6723            let mut next = e.uninit(t * n_embd)?;
6724            for r in 0..t {
6725                let mut row = e.uninit(n_embd)?;
6726                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6727                // The caller owns this verify's position. During controller overlap, cache.pos
6728                // still describes generation N while this stage-0 walk belongs to N+1.
6729                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6730                let mut one = [&mut *cache];
6731                let out = self.step35_decode_batch_layers(
6732                    e,
6733                    row,
6734                    &mut one,
6735                    &[(pos0 + r) as i32],
6736                    &row_pos,
6737                    il,
6738                    il + 1,
6739                    &mut ph_last,
6740                )?;
6741                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6742            }
6743            self.dflash_tap(e, cache, il, &next, t)?;
6744            x = next;
6745            if spec_nan_scan() {
6746                nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6747            }
6748        }
6749        Ok(x)
6750    }
6751
6752    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6753    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6754    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6755    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6756    /// prefix-keep, not all-or-nothing).
6757    pub(crate) fn dspark_verify_t_am(
6758        &self,
6759        e: &Engine,
6760        tokens: &[u32],
6761        pos0: usize,
6762        cache: &mut Cache,
6763    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6764        let (logits, _hn) = self.decode_step_t_core_stream(
6765            e, tokens, pos0, cache, None, None, None, None, None, None,
6766        )?;
6767        let t = tokens.len();
6768        let v = self.output.out_features();
6769        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6770        for r in 0..t {
6771            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6772        }
6773        Ok(e.dtoh_u32(&am_d)?)
6774    }
6775
6776    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6777    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6778    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6779    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6780    pub(crate) fn dspark_verify_t_logits(
6781        &self,
6782        e: &Engine,
6783        tokens: &[u32],
6784        pos0: usize,
6785        cache: &mut Cache,
6786    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6787        let (logits, _hn) = self.decode_step_t_core_stream(
6788            e, tokens, pos0, cache, None, None, None, None, None, None,
6789        )?;
6790        Ok(logits)
6791    }
6792
6793    /// DSpark verify with the MTP column-stash armed: identical forward to
6794    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6795    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6796    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6797    pub(crate) fn dspark_verify_t_am_ckpt(
6798        &self,
6799        e: &Engine,
6800        tokens: &[u32],
6801        pos0: usize,
6802        cache: &mut Cache,
6803    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6804        let mut ck = VerifyCkpt::new(self.layers.len());
6805        let (logits, _hn) = self.decode_step_t_core_stream(
6806            e,
6807            tokens,
6808            pos0,
6809            cache,
6810            None,
6811            Some(&mut ck),
6812            None,
6813            None,
6814            None,
6815            None,
6816        )?;
6817        let t = tokens.len();
6818        let v = self.output.out_features();
6819        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6820        for r in 0..t {
6821            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6822        }
6823        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6824    }
6825
6826    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6827    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6828    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6829    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6830    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6831    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6832    pub(crate) fn dspark_verify_t_am_ckpt_dev(
6833        &self,
6834        e: &Engine,
6835        vtok: &CudaSlice<u32>,
6836        t: usize,
6837        pos0: usize,
6838        cache: &mut Cache,
6839        embd_dev: (&CudaSlice<u8>, i32, usize),
6840        graphs: Option<&mut DsparkVerifyGraphs>,
6841    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6842        debug_assert!(
6843            vtok.len() >= t,
6844            "verify window exceeds the device token buffer"
6845        );
6846        // The slab flag is a per-round statement: clear it here so a verify that never
6847        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6848        // stale `true` steering the commit at slabs the round never wrote.
6849        let mut graphs = graphs;
6850        if let Some(g) = graphs.as_deref_mut() {
6851            g.round_slab = false;
6852        }
6853        let mut ck = VerifyCkpt::new(self.layers.len());
6854        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6855        // arm's established pattern — spec.rs stream-mode verify does the same).
6856        let dummy = vec![0u32; t];
6857        let (logits, _hn) = self.decode_step_t_core_stream(
6858            e,
6859            &dummy,
6860            pos0,
6861            cache,
6862            Some(embd_dev),
6863            Some(&mut ck),
6864            None,
6865            None,
6866            Some(vtok),
6867            graphs,
6868        )?;
6869        let v = self.output.out_features();
6870        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6871        for r in 0..t {
6872            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6873        }
6874        Ok((am_d, DsparkVerifyCkpt(ck)))
6875    }
6876
6877    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6878    pub(crate) fn dspark_verify_t_logits_ckpt(
6879        &self,
6880        e: &Engine,
6881        tokens: &[u32],
6882        pos0: usize,
6883        cache: &mut Cache,
6884    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6885        let mut ck = VerifyCkpt::new(self.layers.len());
6886        let (logits, _hn) = self.decode_step_t_core_stream(
6887            e,
6888            tokens,
6889            pos0,
6890            cache,
6891            None,
6892            Some(&mut ck),
6893            None,
6894            None,
6895            None,
6896            None,
6897        )?;
6898        Ok((logits, DsparkVerifyCkpt(ck)))
6899    }
6900
6901    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6902    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6903    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6904    pub(crate) fn dspark_commit_prefix(
6905        &self,
6906        e: &Engine,
6907        cache: &mut Cache,
6908        snap: &crate::cache::CacheSnapshot,
6909        ckpt: &DsparkVerifyCkpt,
6910        keep: usize,
6911    ) -> Result<(), Box<dyn std::error::Error>> {
6912        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6913    }
6914
6915    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6916    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6917    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6918    /// from the stash of column keep-1), slab-addressed and batched into two copy
6919    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6920    pub(crate) fn dspark_commit_prefix_slab(
6921        &self,
6922        e: &Engine,
6923        cache: &mut Cache,
6924        snap: &crate::cache::CacheSnapshot,
6925        ctx: &DsparkVerifyGraphs,
6926        keep: usize,
6927    ) -> Result<(), Box<dyn std::error::Error>> {
6928        use cudarc::driver::DevicePtr;
6929        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6930        let mut conv_src: Vec<u64> = Vec::new();
6931        let mut ssm_src: Vec<u64> = Vec::new();
6932        let mut conv_dst: Vec<u64> = Vec::new();
6933        let mut ssm_dst: Vec<u64> = Vec::new();
6934        for il in 0..self.layers.len() {
6935            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6936                kvl.len = saved + keep;
6937                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6938            }
6939            if let Some(rl) = cache.recur[il].as_ref() {
6940                let (pc, ps, _cw, _sw) = ctx
6941                    .slab_row(e, il, keep - 1)
6942                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6943                conv_src.push(pc);
6944                ssm_src.push(ps);
6945                let st = &e.gpu.stream();
6946                let (dc, _g0) = rl.conv_state.device_ptr(st);
6947                let (ds, _g1) = rl.ssm_state.device_ptr(st);
6948                conv_dst.push(dc as u64);
6949                ssm_dst.push(ds as u64);
6950            }
6951        }
6952        let n = conv_src.len();
6953        if n > 0 {
6954            if state_copy_batch_on() {
6955                let mut tt = vec![0u64; 2 * n];
6956                tt[..n].copy_from_slice(&conv_src);
6957                tt[n..].copy_from_slice(&conv_dst);
6958                let ct = e.htod_u64(&tt)?;
6959                tt[..n].copy_from_slice(&ssm_src);
6960                tt[n..].copy_from_slice(&ssm_dst);
6961                let st = e.htod_u64(&tt)?;
6962                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6963                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6964            } else {
6965                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6966                let row = keep - 1;
6967                for il in 0..self.layers.len() {
6968                    let Some(rl) = cache.recur[il].as_mut() else {
6969                        continue;
6970                    };
6971                    let k = ctx.lin_pos[&il];
6972                    {
6973                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
6974                        let win = sv.slice(row * cw..(row + 1) * cw);
6975                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
6976                    }
6977                    {
6978                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
6979                        let win = sv.slice(row * sw..(row + 1) * sw);
6980                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
6981                    }
6982                }
6983            }
6984        }
6985        cache.pos = snap.pos + keep;
6986        Ok(())
6987    }
6988
6989    /// Qwen35-family verify trunk in the live serving numeric class.
6990    ///
6991    /// Serving intentionally keeps this architecture in the generic batched program even at
6992    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
6993    ///
6994    /// Two arms, one numeric class:
6995    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
6996    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
6997    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
6998    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
6999    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7000    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7001    ///   program its isolated serving step would). One weight read per layer per round
7002    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
7003    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7004    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7005    ///   serving layer body, preserving single-session autoregressive cache order (the
7006    ///   correctness reference; also the rollback seam for the t-parallel arm).
7007    ///
7008    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7009    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7010    #[allow(clippy::too_many_arguments)]
7011    fn qwen35_verify_batch_layers(
7012        &self,
7013        e: &Engine,
7014        x: CudaSlice<f32>,
7015        lo: usize,
7016        hi: usize,
7017        pos0: usize,
7018        t: usize,
7019        cache: &mut Cache,
7020        ckpt: Option<&mut VerifyCkpt>,
7021        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7022        graphs: Option<&mut DsparkVerifyGraphs>,
7023    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7024        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7025        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7026        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7027        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7028        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7029        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7030        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7031            || !self.batched_serving_numeric_class()
7032            || t > 16;
7033        if rowwise {
7034            if stream.is_some() {
7035                // rowwise replays per row with host cache.pos — irreconcilable with a
7036                // device position counter. Burst callers must keep t <= 16 and the
7037                // ROWWISE env unset; refusing beats silently mispositioned rows.
7038                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7039                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7040                    .into());
7041            }
7042            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7043        } else {
7044            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7045        }
7046    }
7047
7048    /// The per-row correctness reference: replay each verify row through the authoritative
7049    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7050    #[allow(clippy::too_many_arguments)]
7051    fn qwen35_verify_rowwise(
7052        &self,
7053        e: &Engine,
7054        mut x: CudaSlice<f32>,
7055        lo: usize,
7056        hi: usize,
7057        pos0: usize,
7058        t: usize,
7059        cache: &mut Cache,
7060        mut ckpt: Option<&mut VerifyCkpt>,
7061    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7062        let n_embd = self.cfg.n_embd as usize;
7063        let saved_pos = cache.pos;
7064        let mut ph_last = std::time::Instant::now();
7065        for il in lo..hi {
7066            let mut next = e.uninit(t * n_embd)?;
7067            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7068                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7069                    Some(Vec::with_capacity(t - 1))
7070                } else {
7071                    None
7072                };
7073            for r in 0..t {
7074                cache.pos = pos0 + r;
7075                let mut row = e.uninit(n_embd)?;
7076                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7077                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7078                let mut one = [&mut *cache];
7079                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7080                let out = match self.decode_batch_layers(
7081                    e,
7082                    row,
7083                    &mut one,
7084                    &ctx,
7085                    &row_pos,
7086                    &mut ph_last,
7087                ) {
7088                    Ok(out) => out,
7089                    Err(error) => {
7090                        cache.pos = saved_pos;
7091                        return Err(error);
7092                    }
7093                };
7094                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7095                if r + 1 < t {
7096                    if let Some(states) = col_states.as_mut() {
7097                        let recur = cache.recur[il]
7098                            .as_ref()
7099                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7100                        states.push((
7101                            e.clone_dtod(&recur.conv_state)?,
7102                            e.clone_dtod(&recur.ssm_state)?,
7103                        ));
7104                    }
7105                }
7106            }
7107            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7108                checkpoint.cols[il] = Some(states);
7109            }
7110            x = next;
7111        }
7112        cache.pos = saved_pos;
7113        Ok(x)
7114    }
7115
7116    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7117    ///
7118    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7119    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7120    /// pins the serving batch tier already carries:
7121    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7122    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7123    ///     alone;
7124    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7125    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7126    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
7127    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7128    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7129    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7130    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7131    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7132    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7133    /// program its isolated B=1 serving step would.
7134    ///
7135    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7136    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7137    #[allow(clippy::too_many_arguments)]
7138    fn qwen35_verify_tparallel(
7139        &self,
7140        e: &Engine,
7141        mut x: CudaSlice<f32>,
7142        lo: usize,
7143        hi: usize,
7144        pos0: usize,
7145        t: usize,
7146        cache: &mut Cache,
7147        mut ckpt: Option<&mut VerifyCkpt>,
7148        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7149        mut graphs: Option<&mut DsparkVerifyGraphs>,
7150    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7151        let seqs_append =
7152            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7153        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7154
7155        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7156        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7157        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7158        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7159        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7160        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7161        // full-verify bodies).
7162        if stream.is_some() && graphs.is_some() {
7163            return Err(
7164                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7165                        cannot arm together"
7166                    .into(),
7167            );
7168        }
7169        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7170        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7171        // moves the kv caches). Then:
7172        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7173        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7174        //    full-verify graph per (vt, rung) — linear layers through the shared
7175        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
7176        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
7177        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
7178        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7179        //    the full-attention layers run eager (batched rows when eligible).
7180        if let Some(g) = graphs.as_deref_mut() {
7181            g.refresh_tables(e, cache)?;
7182            g.round_slab = false;
7183            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7184                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7185                // full capture past the ceiling falls through to the segment/eager arms.
7186                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7187                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7188                    g.round_slab = true;
7189                    return Ok(out);
7190                }
7191            }
7192            // Round-atomic ceiling check for the segment door: if any linear run in this
7193            // walk would need a NEW capture past the ceiling, the whole round runs the
7194            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7195            // would corrupt the commit).
7196            if !g.segments_ready(self, lo, hi, t) {
7197                graphs = None;
7198            }
7199        }
7200        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7201        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7202        let pos_d = match stream {
7203            Some((_, ctr)) => {
7204                let mut p = e.alloc_uninit::<i32>(t)?;
7205                e.pos_iota(ctr, &mut p, t)?;
7206                p
7207            }
7208            None => {
7209                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7210                e.htod_i32(&pos_host)?
7211            }
7212        };
7213        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7214        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7215        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7216        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7217        // rides the dc rows kernels and never reaches the fallback).
7218        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7219        let mut il = lo;
7220        while il < hi {
7221            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7222                let mut end = il;
7223                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7224                    end += 1;
7225                }
7226                let g = graphs.as_deref_mut().expect("checked above");
7227                x = g.run_segment(self, e, il, end, &x, t, cache)?;
7228                g.round_slab = true;
7229                il = end;
7230                continue;
7231            }
7232            let layer = &self.layers[il];
7233            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7234                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7235                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7236                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7237                x = self.qwen35_tparallel_linear_layer(
7238                    e,
7239                    il,
7240                    &x,
7241                    t,
7242                    cache,
7243                    ckpt.as_deref_mut(),
7244                    None,
7245                    None,
7246                )?;
7247                il += 1;
7248                continue;
7249            }
7250            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7251            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7252            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7253            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7254            // run (lane/draftcost-moe).
7255            x = self.qwen35_tparallel_fa_layer(
7256                e,
7257                il,
7258                &x,
7259                t,
7260                cache,
7261                FaLayerArgs {
7262                    pos_d: &pos_d,
7263                    pos_rows: &mut pos_rows,
7264                    pos0,
7265                    seqs_append,
7266                    batch_fa_on,
7267                    graph_cap: None,
7268                    stream,
7269                    ckpt: ckpt.as_deref_mut(),
7270                },
7271            )?;
7272            il += 1;
7273        }
7274        Ok(x)
7275    }
7276
7277    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7278    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7279    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7280    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7281    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7282    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7283    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7284    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7285    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7286    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7287    /// original singles chain, byte-for-byte.
7288    #[allow(clippy::too_many_arguments)]
7289    fn qwen35_tparallel_dense_ffn(
7290        &self,
7291        e: &Engine,
7292        ffn_gate: &crate::model::GpuTensor,
7293        ffn_up: &crate::model::GpuTensor,
7294        ffn_down: &crate::model::GpuTensor,
7295        zn: &CudaSlice<f32>,
7296        t: usize,
7297        n_embd: usize,
7298    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7299        let n_ff = ffn_gate.out_features();
7300        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7301        if Engine::tk_ffn_dual_on() {
7302            if let Some(((g, gs), (u, us))) =
7303                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7304            {
7305                if e.uses_q8_1_fast(ffn_down) {
7306                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7307                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7308                }
7309                let mut act = e.uninit(t * n_ff)?;
7310                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7311                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7312                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7313            }
7314        }
7315        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7316        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7317        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7318        let mut act = e.uninit(t * n_ff)?;
7319        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7320        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7321        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7322    }
7323
7324    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7325    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7326    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7327    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7328    ///
7329    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7330    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7331    ///   generation's cache lands at new addresses that only the per-verify table refresh
7332    ///   knows — the slice-3 baked-address lesson);
7333    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
7334    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
7335    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
7336    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
7337    ///   round whose rows all sit inside the rung;
7338    /// - the host len bump moves to the replay caller (captured host code does not
7339    ///   re-run at replay).
7340    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
7341    /// host-branches on t_kv and must never be captured.
7342    #[allow(clippy::too_many_arguments)]
7343    fn qwen35_tparallel_fa_layer(
7344        &self,
7345        e: &Engine,
7346        il: usize,
7347        x: &CudaSlice<f32>,
7348        t: usize,
7349        cache: &mut Cache,
7350        args: FaLayerArgs<'_>,
7351    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7352        use cudarc::driver::DevicePtr;
7353        let cfg = &self.cfg;
7354        let n_embd = cfg.n_embd as usize;
7355        let eps = cfg.rms_eps;
7356        let head_dim_global = cfg.head_dim_k as usize;
7357        let layer = &self.layers[il];
7358        let FaLayerArgs {
7359            pos_d,
7360            pos_rows,
7361            pos0,
7362            seqs_append,
7363            batch_fa_on,
7364            graph_cap,
7365            stream,
7366            mut ckpt,
7367        } = args;
7368
7369        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7370        let anorm = layer.attn_norm.float_data();
7371        let mut xn = e.uninit(t * n_embd)?;
7372        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7373        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7374
7375        let mixed: CudaSlice<f32> = match &layer.mixer {
7376            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7377            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
7378            // per-row serving-kernel chain cannot run (host state swaps keyed on host
7379            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
7380            // rebuild — the per-row chain only produces per-column clones). GDN rides
7381            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
7382            // and its one-scan recurrence is pinned bit-identical to T chained T=1
7383            // steps (its header + kernel-check). Position-independent, so no counter
7384            // plumbing is needed. Guards mirror the generic call site exactly.
7385            Mixer::Linear(la) if stream.is_some() => {
7386                if !(t >= 3 || (t == 2 && spec_m2()))
7387                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
7388                    || !e.uses_q8_1_fast(&la.ssm_out)
7389                {
7390                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
7391                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
7392                        .into());
7393                }
7394                let want = ckpt.is_some();
7395                let (out, stash) =
7396                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
7397                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7398                    ck.gdn[il] = Some(st);
7399                }
7400                out
7401            }
7402            Mixer::Linear(_) => {
7403                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
7404            }
7405            Mixer::Full(fa) => {
7406                let geometry = cfg.full_attention_geometry_at(il as u32);
7407                let n_head = geometry.n_head as usize;
7408                let n_head_kv = geometry.n_head_kv as usize;
7409                let head_dim = geometry.head_dim_k as usize;
7410                let rope_dims = geometry.n_rot as usize;
7411                let rope_base = geometry.rope_base;
7412                let scale = geometry.attention_scale();
7413                // Batched projections: one weight read serves all T rows.
7414                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
7415                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
7416                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
7417                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
7418                    [&fa.wq, &fa.wk, &fa.wv],
7419                    &hq,
7420                    &hd,
7421                    t,
7422                )? {
7423                    Some(mut g3) => {
7424                        let v = g3.pop().unwrap();
7425                        let k = g3.pop().unwrap();
7426                        let qf = g3.pop().unwrap();
7427                        (qf, k, v)
7428                    }
7429                    None => (
7430                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
7431                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
7432                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
7433                    ),
7434                };
7435                let gated =
7436                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7437                let (mut q, gate) = if gated {
7438                    let mut qs = e.uninit(t * n_head * head_dim)?;
7439                    let mut gs = e.uninit(t * n_head * head_dim)?;
7440                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
7441                    (qs, Some(gs))
7442                } else {
7443                    (qf, None)
7444                };
7445                let mut qn = e.uninit(t * n_head * head_dim)?;
7446                e.rms_norm(
7447                    &q,
7448                    fa.q_norm.float_data(),
7449                    &mut qn,
7450                    head_dim,
7451                    t * n_head,
7452                    eps,
7453                )?;
7454                q = qn;
7455                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7456                e.rms_norm(
7457                    &k,
7458                    fa.k_norm.float_data(),
7459                    &mut kn,
7460                    head_dim,
7461                    t * n_head_kv,
7462                    eps,
7463                )?;
7464                k = kn;
7465                e.rope_neox(
7466                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
7467                )?;
7468                e.rope_neox(
7469                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
7470                )?;
7471
7472                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
7473                // draft), each through the b_n=1 serving kernels at its own t_kv.
7474                let q_dim = n_head * head_dim;
7475                let kv_dim = n_head_kv * head_dim;
7476                let mut attn = e.uninit(t * q_dim)?;
7477                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
7478                    let kvl = cache.kv[il].as_ref().unwrap();
7479                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
7480                    // the batched twins; the per-row fallback reads pair 0 (same cache
7481                    // for every row of one layer). Graph mode reads the ctx table.
7482                    let local: Option<CudaSlice<u64>> = match graph_cap {
7483                        Some(_) => None,
7484                        None => {
7485                            let s = &e.gpu.stream();
7486                            let (pk, _g) = kvl.k.device_ptr(s);
7487                            let (pv, _g2) = kvl.v.device_ptr(s);
7488                            let mut tbl = Vec::with_capacity(2 * t);
7489                            for _ in 0..t {
7490                                tbl.push(pk as u64);
7491                                tbl.push(pv as u64);
7492                            }
7493                            Some(e.htod_u64(&tbl)?)
7494                        }
7495                    };
7496                    (
7497                        kvl.kv_dim_k,
7498                        kvl.kv_dim_v,
7499                        kvl.k_tok_bytes,
7500                        kvl.v_tok_bytes,
7501                        kvl.len,
7502                        local,
7503                    )
7504                };
7505                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
7506                    Some((tb, off, _)) => (tb, off),
7507                    None => (kv_local.as_ref().expect("built above"), 0),
7508                };
7509                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
7510                // section batches into the z-batched serving twins when every row of
7511                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
7512                // guards are evaluated at the round's FIRST and LAST t_kv — the
7513                // eligibility window (vec floor .. v4 max) and each split-ladder rung
7514                // are intervals in t_kv, so ends-inside means all-inside (the straddle
7515                // law). Appending all T rows before any attend is read-equivalent to
7516                // the interleaved order: row r's walk reads keys 0..len0+r only, and
7517                // rows > r land at slots it never touches; every written cache row is
7518                // the per-token appender's exact warp program (kernel-check pinned).
7519                let t_kv_first = len0 + 1;
7520                let t_kv_last = len0 + t;
7521                let rows_batched = t >= 2
7522                    && seqs_append
7523                    && batch_fa_on
7524                    && dspark_fa_rows_on()
7525                    // the z-batched twins read stacked rows at the CACHE's kv dims;
7526                    // the projection stack is [T, n_head_kv*head_dim] — they must be
7527                    // the same stride or row z misaligns (true for this family; the
7528                    // guard keeps any asymmetric-kv model on the per-row loop).
7529                    && kdk == kv_dim
7530                    && kdv == kv_dim
7531                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
7532                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
7533                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
7534                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
7535                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
7536                // grid only — bytes proven equal above). Capture-time invariants refuse
7537                // loudly rather than bake a divergent body.
7538                let (size_kv_max, sp) = match graph_cap {
7539                    Some((_, _, rung)) => {
7540                        if !rows_batched {
7541                            return Err(format!(
7542                                "fa graph capture: layer {il} round is not batchable \
7543                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
7544                                 must never be captured"
7545                            )
7546                            .into());
7547                        }
7548                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
7549                        if t_kv_last > rung
7550                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
7551                        {
7552                            return Err(format!(
7553                                "fa graph capture: rung {rung} does not cover round \
7554                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
7555                            )
7556                            .into());
7557                        }
7558                        (rung, sp_r)
7559                    }
7560                    None => (
7561                        t_kv_last,
7562                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
7563                    ),
7564                };
7565                if let Some((_, ctr)) = stream {
7566                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
7567                    // — the generic stream arm's exact shape (rows kernels are pinned
7568                    // byte-identical to the per-row programs by kernel-check). Host len
7569                    // stays a stale lower bound; the burst drain reconciles it.
7570                    let kvl = cache.kv[il].as_mut().unwrap();
7571                    e.append_kv_quantized_rows_dc(
7572                        &k,
7573                        &v,
7574                        &mut kvl.k,
7575                        &mut kvl.v,
7576                        ctr,
7577                        t,
7578                        kdk,
7579                        kdv,
7580                        ktb,
7581                        vtb,
7582                        Engine::kv_fp8_on(),
7583                    )?;
7584                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
7585                    let k_view = e.view_u8(&kvl.k, upper * ktb);
7586                    let v_view = e.view_u8(&kvl.v, upper * vtb);
7587                    e.fa_decode_rows_dc(
7588                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
7589                        t, scale, ktb, vtb, 0, false,
7590                    )?;
7591                } else if rows_batched {
7592                    e.append_kv_quantized_seqs(
7593                        &k,
7594                        &v,
7595                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7596                        pos_d,
7597                        t,
7598                        kdk,
7599                        kdv,
7600                        ktb,
7601                        vtb,
7602                    )?;
7603                    if graph_cap.is_none() {
7604                        cache.kv[il].as_mut().unwrap().len += t;
7605                    }
7606                    e.fa_decode_batch_seqs_v4(
7607                        &q,
7608                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7609                        pos_d,
7610                        &mut attn,
7611                        head_dim,
7612                        n_head,
7613                        n_head_kv,
7614                        t,
7615                        size_kv_max,
7616                        scale,
7617                        sp,
7618                        ktb,
7619                        vtb,
7620                    )?;
7621                } else {
7622                    if pos_rows.is_none() {
7623                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
7624                        // the dc rows kernels above and never reaches this fallback).
7625                        *pos_rows = Some(match stream {
7626                            Some((_, ctr)) => (0..t)
7627                                .map(|r| {
7628                                    let mut b = e.alloc_uninit::<i32>(1)?;
7629                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
7630                                    Ok(b)
7631                                })
7632                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
7633                            None => (0..t)
7634                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
7635                                .collect::<Result<_, _>>()?,
7636                        });
7637                    }
7638                    let pos_rows = pos_rows.as_ref().unwrap();
7639                    for r in 0..t {
7640                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
7641                        // whose row 0 is this row (arithmetic-free materialization copies,
7642                        // same as decode's per-seq fallback arm).
7643                        let mut k_row = e.uninit(kv_dim)?;
7644                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
7645                        let mut v_row = e.uninit(kv_dim)?;
7646                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
7647                        let pos_row = &pos_rows[r];
7648                        let kvl = cache.kv[il].as_mut().unwrap();
7649                        if seqs_append {
7650                            e.append_kv_quantized_seqs(
7651                                &k_row,
7652                                &v_row,
7653                                &kv_tbl.slice(kv_off..kv_off + 2),
7654                                pos_row,
7655                                1,
7656                                kdk,
7657                                kdv,
7658                                ktb,
7659                                vtb,
7660                            )?;
7661                            kvl.len += 1;
7662                        } else {
7663                            e.append_kv_quantized_view(
7664                                &k_row.slice(0..kv_dim),
7665                                &v_row.slice(0..kv_dim),
7666                                &mut kvl.k,
7667                                &mut kvl.v,
7668                                kvl.len,
7669                                kvl.kv_dim_k,
7670                                kvl.kv_dim_v,
7671                                kvl.k_tok_bytes,
7672                                kvl.v_tok_bytes,
7673                                Engine::kv_fp8_on(),
7674                            )?;
7675                            kvl.len += 1;
7676                        }
7677                        let t_kv = kvl.len;
7678                        let mut q_row = e.uninit(q_dim)?;
7679                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
7680                        let mut a_row = e.uninit(q_dim)?;
7681                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
7682                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
7683                            e.fa_decode_batch_seqs_v4(
7684                                &q_row,
7685                                &kv_tbl.slice(kv_off..kv_off + 2),
7686                                pos_row,
7687                                &mut a_row,
7688                                head_dim,
7689                                n_head,
7690                                n_head_kv,
7691                                1,
7692                                t_kv,
7693                                scale,
7694                                sp0_r,
7695                                ktb,
7696                                vtb,
7697                            )?;
7698                        } else {
7699                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7700                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7701                            let mut a_view = a_row.slice_mut(0..q_dim);
7702                            e.fa_decode_kvmod_view(
7703                                &q_row.slice(0..q_dim),
7704                                &k_view,
7705                                &v_view,
7706                                &mut a_view,
7707                                head_dim,
7708                                n_head,
7709                                n_head_kv,
7710                                t_kv,
7711                                scale,
7712                                kvl.k_tok_bytes,
7713                                kvl.v_tok_bytes,
7714                                Engine::kv_fp8_on(),
7715                            )?;
7716                        }
7717                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7718                    }
7719                }
7720
7721                // Output gate (element-wise) + o-proj at m=T.
7722                let attn_g = match &gate {
7723                    Some(g) => {
7724                        let n = t * q_dim;
7725                        let mut gsig = e.uninit(n)?;
7726                        e.sigmoid(g, &mut gsig, n)?;
7727                        let mut ag = e.uninit(n)?;
7728                        e.mul(&attn, &gsig, &mut ag, n)?;
7729                        ag
7730                    }
7731                    None => attn,
7732                };
7733                e.matmul(&fa.wo, &attn_g, t)?
7734            }
7735        };
7736
7737        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7738        let pnorm = layer.post_attn_norm.float_data();
7739        let mut x1 = e.uninit(t * n_embd)?;
7740        let mut zn = e.uninit(t * n_embd)?;
7741        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7742        let ffn_out = match &layer.ffn {
7743            crate::hybrid::Ffn::Dense {
7744                ffn_gate,
7745                ffn_up,
7746                ffn_down,
7747            } => {
7748                assert!(
7749                    self.cfg.m3.is_none(),
7750                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7751                );
7752                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7753            }
7754            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7755        };
7756        let mut x2 = e.uninit(t * n_embd)?;
7757        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7758        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7759        self.dflash_tap(e, cache, il, &x2, t)?;
7760        Ok(x2)
7761    }
7762
7763    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7764    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7765    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7766    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7767    /// bit-identical by construction:
7768    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7769    ///   the device sequence is driven entirely by the 6-entry pointer table, which
7770    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
7771    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7772    ///   legacy post-swap clone read.
7773    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7774    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7775    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7776    /// None builds the per-verify table exactly as before.
7777    #[allow(clippy::too_many_arguments)]
7778    fn qwen35_tparallel_linear_layer(
7779        &self,
7780        e: &Engine,
7781        il: usize,
7782        x: &CudaSlice<f32>,
7783        t: usize,
7784        cache: &mut Cache,
7785        mut ckpt: Option<&mut VerifyCkpt>,
7786        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7787        table_src: Option<(&CudaSlice<u64>, usize)>,
7788    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7789        use cudarc::driver::DevicePtr;
7790        let cfg = &self.cfg;
7791        let n_embd = cfg.n_embd as usize;
7792        let eps = cfg.rms_eps;
7793        let layer = &self.layers[il];
7794        let Mixer::Linear(la) = &layer.mixer else {
7795            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7796        };
7797        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7798        let anorm = layer.attn_norm.float_data();
7799        let mut xn = e.uninit(t * n_embd)?;
7800        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7801        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7802
7803        let geometry = la.geometry;
7804        let d_state = geometry.key_head_dim as usize;
7805        let num_k = geometry.key_heads as usize;
7806        let num_v = geometry.value_heads as usize;
7807        let d_conv = geometry.conv_kernel as usize;
7808        let key_dim = d_state * num_k;
7809        let value_dim = geometry.value_head_dim as usize * num_v;
7810        let conv_dim = key_dim * 2 + value_dim;
7811        let gdn_scale = 1.0 / (d_state as f32).sqrt();
7812
7813        // ---- batched projections: one weight read for all T rows ----
7814        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7815        // per (tensor, token, row) to the four singles; refused (layout/tier) or
7816        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7817        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7818            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7819            &hq,
7820            &hd,
7821            t,
7822        )? {
7823            Some(mut g4) => {
7824                let alpha = g4.pop().unwrap();
7825                let beta_raw = g4.pop().unwrap();
7826                let z = g4.pop().unwrap();
7827                let qkv_mixed = g4.pop().unwrap();
7828                (qkv_mixed, z, beta_raw, alpha)
7829            }
7830            None => (
7831                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7832                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7833                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7834                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7835            ),
7836        };
7837        let beta_w = la.ssm_beta.out_features();
7838        let alpha_w = la.ssm_alpha.out_features();
7839        let qkv_w = la.wqkv.out_features();
7840
7841        // ---- per-row state chain through the b_n=1 serving kernels ----
7842        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7843        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7844        let table_local: Option<CudaSlice<u64>> = match table_src {
7845            Some(_) => None,
7846            None => {
7847                let rl = cache.recur[il].as_ref().unwrap();
7848                let s = &e.gpu.stream();
7849                let (pc, _g0) = rl.conv_state.device_ptr(s);
7850                let (p0, _g1) = rl.ssm_state.device_ptr(s);
7851                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7852                Some(e.htod_u64(&[
7853                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
7854                ])?)
7855            }
7856        };
7857        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7858            Some((tb, off)) => (tb, off),
7859            None => (table_local.as_ref().unwrap(), 0),
7860        };
7861        let mut o_all = e.uninit(t * value_dim)?;
7862        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7863            if ckpt.is_some() && stash.is_none() && t >= 2 {
7864                Some(Vec::with_capacity(t - 1))
7865            } else {
7866                None
7867            };
7868        let mut stash = stash;
7869        // Per-row scratch reused across rows (uninit is cheap but not free at
7870        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7871        // [T, ...] buffers — zero arithmetic-free copies in this loop.
7872        let mut conv_out = e.uninit(conv_dim)?;
7873        let mut q_l2 = e.uninit(value_dim)?;
7874        let mut k_l2 = e.uninit(value_dim)?;
7875        let mut v_gd = e.uninit(value_dim)?;
7876        let mut beta_b = e.uninit(num_v)?;
7877        let mut g_log = e.uninit(num_v)?;
7878        for r in 0..t {
7879            let base = toff + if r % 2 == 0 { 0 } else { 3 };
7880            let conv_view = table.slice(base..base + 1);
7881            let in_view = table.slice(base + 1..base + 2);
7882            let out_view = table.slice(base + 2..base + 3);
7883            e.ssm_conv1d_fused_decode_b_view(
7884                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7885                &conv_view,
7886                la.ssm_conv1d.float_data(),
7887                &mut conv_out,
7888                conv_dim,
7889                d_conv,
7890                1,
7891            )?;
7892            e.gdn_prep_decode_b_view(
7893                &conv_out,
7894                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7895                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7896                la.ssm_dt.float_data(),
7897                la.ssm_a.float_data(),
7898                &mut q_l2,
7899                &mut k_l2,
7900                &mut v_gd,
7901                &mut beta_b,
7902                &mut g_log,
7903                d_state,
7904                num_v,
7905                num_k,
7906                key_dim,
7907                eps,
7908                conv_dim,
7909                1,
7910            )?;
7911            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7912            e.gdn_scan_s128_batched_view(
7913                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7914                gdn_scale,
7915            )?;
7916            if r + 1 < t {
7917                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7918                // odd rows write s0 — the same physical state the legacy post-swap
7919                // canonical clone read.
7920                let rl = cache.recur[il]
7921                    .as_ref()
7922                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
7923                let ssm_src = if r % 2 == 0 {
7924                    &rl.ssm_state_alt
7925                } else {
7926                    &rl.ssm_state
7927                };
7928                match stash.as_mut() {
7929                    Some((conv_slab, ssm_slab)) => {
7930                        // BOTH stash reads go through the pointer table at run time: the
7931                        // ssm handles ping-pong between rounds, and the ctx (with its
7932                        // captured graphs) outlives the Cache — a fresh generation's
7933                        // conv/ssm buffers land at new addresses that only the per-round
7934                        // table refresh knows. A baked direct copy would read freed
7935                        // memory (parity was the slice-3 smoke divergence; cache
7936                        // lifetime is the cross-generation twin).
7937                        e.copy_indirect_src_f32(
7938                            &conv_view,
7939                            conv_slab,
7940                            r * conv_dim * (d_conv - 1),
7941                            conv_dim * (d_conv - 1),
7942                        )?;
7943                        // The ssm handles PING-PONG between rounds: a captured direct
7944                        // copy would bake the capture-time physical buffer and read the
7945                        // wrong parity after any odd-vt round (the slice-3 smoke
7946                        // divergence). Read the src address from row r's OUT table
7947                        // entry at run time — the same entry the scan just wrote.
7948                        e.copy_indirect_src_f32(
7949                            &out_view,
7950                            ssm_slab,
7951                            r * d_state * d_state * num_v,
7952                            d_state * d_state * num_v,
7953                        )?;
7954                    }
7955                    None => {
7956                        if let Some(states) = col_states.as_mut() {
7957                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7958                        }
7959                    }
7960                }
7961            }
7962        }
7963        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7964        // handle motion is identical and the device sequence never read the handles.
7965        if t % 2 == 1 {
7966            let rl = cache.recur[il].as_mut().unwrap();
7967            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
7968        }
7969        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7970            checkpoint.cols[il] = Some(states);
7971        }
7972
7973        // ---- batched gated norm + out-projection at m=T ----
7974        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
7975            let (gq, gd) = e.gated_rmsnorm_q8_1(
7976                &o_all,
7977                la.ssm_norm.float_data(),
7978                &z,
7979                d_state,
7980                t * num_v,
7981                eps,
7982            )?;
7983            let g0 = e.zeros(0)?;
7984            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
7985        } else {
7986            let mut gn = e.uninit(t * value_dim)?;
7987            e.gated_rmsnorm(
7988                &o_all,
7989                la.ssm_norm.float_data(),
7990                &z,
7991                &mut gn,
7992                d_state,
7993                t * num_v,
7994                eps,
7995            )?;
7996            e.matmul(&la.ssm_out, &gn, t)?
7997        };
7998
7999        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8000        let pnorm = layer.post_attn_norm.float_data();
8001        let mut x1 = e.uninit(t * n_embd)?;
8002        let mut zn = e.uninit(t * n_embd)?;
8003        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8004        let ffn_out = match &layer.ffn {
8005            crate::hybrid::Ffn::Dense {
8006                ffn_gate,
8007                ffn_up,
8008                ffn_down,
8009            } => {
8010                assert!(
8011                    self.cfg.m3.is_none(),
8012                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8013                );
8014                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8015            }
8016            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8017        };
8018        let mut x2 = e.uninit(t * n_embd)?;
8019        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8020        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8021        self.dflash_tap(e, cache, il, &x2, t)?;
8022        Ok(x2)
8023    }
8024
8025    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8026    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8027    /// carried in from outside the range) and exits with the range's final residual materialized
8028    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8029    /// instead of one.
8030    ///
8031    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8032    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8033    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8034    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8035    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8036    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8037    /// code — there is no "split version" of the verify math.
8038    ///
8039    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8040    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8041    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8042    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8043    #[allow(clippy::too_many_arguments)]
8044    fn verify_layers(
8045        &self,
8046        e: &Engine,
8047        mut x: CudaSlice<f32>,
8048        lo: usize,
8049        hi: usize,
8050        pos_d: &CudaSlice<i32>,
8051        pos0: usize,
8052        t: usize,
8053        cache: &mut Cache,
8054        mut ckpt: Option<&mut VerifyCkpt>,
8055        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8056        graphs: Option<&mut DsparkVerifyGraphs>,
8057    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8058        if self.sliding_gated_moe_batch_program() {
8059            if stream.is_some() {
8060                return Err(
8061                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8062                            cannot express the SWA offset KV view)"
8063                        .into(),
8064                );
8065            }
8066            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8067        }
8068        if self.batched_serving_numeric_class() {
8069            return self.qwen35_verify_batch_layers(
8070                e,
8071                x,
8072                lo,
8073                hi,
8074                pos0,
8075                t,
8076                cache,
8077                ckpt.take(),
8078                stream,
8079                graphs,
8080            );
8081        }
8082        let n_embd = self.cfg.n_embd as usize;
8083        let eps = self.cfg.rms_eps;
8084        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8085        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8086        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8087        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8088        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8089        // residual the next layer needs) as its `res` output. Falls back to the separate add
8090        // when the next layer is off the fused-q8 path.
8091        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8092        for il in lo..hi {
8093            let layer = &self.layers[il];
8094            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8095            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8096            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8097            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8098            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8099            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8100            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8101            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8102            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8103            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8104            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8105            // projections only; Linear mixer: the batched arm — the per-column fallback needs
8106            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8107            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8108            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8109            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8110            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8111            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8112            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8113            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8114            let lin_q8_only = match &layer.mixer {
8115                Mixer::Linear(la) => {
8116                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8117                }
8118                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8119                _ => true,
8120            };
8121            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8122            // a non-fused layer still performs the residual add.
8123            let taken = pending.take();
8124            let (h, h_q8) = if norm_fused && lin_q8_only {
8125                let pair = match taken {
8126                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8127                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8128                    Some((x1p, f1p)) => {
8129                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8130                        let p = e.add_rms_norm_q8_1(
8131                            &x1p,
8132                            &f1p,
8133                            layer.attn_norm.float_data(),
8134                            &mut x2,
8135                            n_embd,
8136                            t,
8137                            eps,
8138                        )?;
8139                        x = x2;
8140                        p
8141                    }
8142                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8143                };
8144                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8145            } else {
8146                if let Some((x1p, f1p)) = taken {
8147                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8148                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8149                    x = x2;
8150                }
8151                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8152                if norm_fused {
8153                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8154                } else {
8155                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8156                }
8157                (h, None)
8158            };
8159            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8160
8161            let mixed = match &layer.mixer {
8162                Mixer::Full(fa) => self.full_attn_verify(
8163                    e,
8164                    fa,
8165                    &h,
8166                    h_q8_ref,
8167                    pos_d,
8168                    t,
8169                    cache,
8170                    il,
8171                    stream.map(|(_, c)| c),
8172                )?,
8173                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8174                Mixer::Linear(la) => {
8175                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8176                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8177                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8178                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8179                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8180                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8181                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8182                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8183                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8184                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8185                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8186                    if (t >= 3 || (t == 2 && spec_m2()))
8187                        && mixer_fast
8188                        && e.uses_q8_1_fast(&la.ssm_out)
8189                    {
8190                        let want = ckpt.is_some();
8191                        let (out, stash) =
8192                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8193                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8194                            ck.gdn[il] = Some(st);
8195                        }
8196                        out
8197                    } else {
8198                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8199                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8200                            if ckpt.is_some() && t >= 2 {
8201                                Some(Vec::with_capacity(t - 1))
8202                            } else {
8203                                None
8204                            };
8205                        for col in 0..t {
8206                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8207                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
8208                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8209                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8210                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8211                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8212                            // (pure dtod — cannot change any computed value). Last column skipped:
8213                            // rebuild targets are j <= t-1 columns.
8214                            if let Some(cs) = col_states.as_mut() {
8215                                if col + 1 < t {
8216                                    let rl = cache.recur[il].as_ref().unwrap();
8217                                    cs.push((
8218                                        e.clone_dtod(&rl.conv_state)?,
8219                                        e.clone_dtod(&rl.ssm_state)?,
8220                                    ));
8221                                }
8222                            }
8223                        }
8224                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8225                            // ReplaySSM-assessment instrumentation (2026-07-30): the
8226                            // per-column clones are the only true state snapshots left in
8227                            // the verify (the batched path stashes INPUTS and replays).
8228                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8229                                static ONCE: std::sync::Once = std::sync::Once::new();
8230                                let bytes: usize =
8231                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8232                                ONCE.call_once(|| eprintln!(
8233                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8234                                    cs.len(), bytes as f64 / 1e6));
8235                            }
8236                            ck.cols[il] = Some(cs);
8237                        }
8238                        out
8239                    }
8240                }
8241            };
8242
8243            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8244            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8245            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8246            let ffn_fuse = match &layer.ffn {
8247                crate::hybrid::Ffn::Dense {
8248                    ffn_gate, ffn_up, ..
8249                } => {
8250                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8251                        && e.uses_q8_1_fast(ffn_gate)
8252                        && e.uses_q8_1_fast(ffn_up)
8253                }
8254                crate::hybrid::Ffn::Moe(_) => false,
8255            };
8256            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8257            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8258            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8259            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8260            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8261            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8262            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8263            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8264            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8265            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8266            // mirror decode's dispatch or spec self-consistency fails.
8267            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8268            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8269            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8270            let mut z = e.zeros(0)?; // replaced below on the unfused arms
8271            let z_q8 = if fuse_q8 {
8272                Some(e.add_rms_norm_q8_1(
8273                    &x,
8274                    &mixed,
8275                    layer.post_attn_norm.float_data(),
8276                    &mut x1,
8277                    n_embd,
8278                    t,
8279                    eps,
8280                )?)
8281            } else {
8282                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8283                if ffn_fuse {
8284                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
8285                    e.rms_norm_decode(
8286                        &x1,
8287                        layer.post_attn_norm.float_data(),
8288                        &mut zf,
8289                        n_embd,
8290                        t,
8291                        eps,
8292                    )?;
8293                } else {
8294                    e.add_rms_norm(
8295                        &x,
8296                        &mixed,
8297                        layer.post_attn_norm.float_data(),
8298                        &mut x1,
8299                        &mut zf,
8300                        n_embd,
8301                        t,
8302                        eps,
8303                    )?;
8304                }
8305                z = zf;
8306                None
8307            };
8308            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8309            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8310            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
8311            let ffn_out = match &layer.ffn {
8312                crate::hybrid::Ffn::Dense {
8313                    ffn_gate,
8314                    ffn_up,
8315                    ffn_down,
8316                } => {
8317                    let n_ff = ffn_gate.out_features();
8318                    if let Some((zq, zd)) = z_q8.as_ref() {
8319                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
8320                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
8321                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
8322                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
8323                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
8324                        // structure at nrows=t.
8325                        let pair =
8326                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
8327                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
8328                                None => None,
8329                            };
8330                        let (gate, gs, up, us) = match pair {
8331                            Some(x4) => x4,
8332                            None => (
8333                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
8334                                1.0, // scale already applied inside _pre
8335                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
8336                                1.0,
8337                            ),
8338                        };
8339                        if e.uses_q8_1_fast(ffn_down) {
8340                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
8341                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
8342                        } else {
8343                            let mut act = vbuf(e, t * n_ff)?;
8344                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
8345                            e.matmul_decode_exact(ffn_down, &act, t)?
8346                        }
8347                    } else {
8348                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
8349                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
8350                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
8351                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
8352                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
8353                        let (gate, up) =
8354                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
8355                                Some(pair) => pair,
8356                                None => (
8357                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
8358                                    e.matmul_decode_exact(ffn_up, &z, t)?,
8359                                ),
8360                            };
8361                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8362                        Self::ffn_act_lim(
8363                            e,
8364                            &self.cfg,
8365                            &gate,
8366                            &up,
8367                            1.0,
8368                            1.0,
8369                            dense_lim,
8370                            &mut act,
8371                            t * n_ff,
8372                        )?;
8373                        e.matmul_decode_exact(ffn_down, &act, t)?
8374                    }
8375                }
8376                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8377            };
8378            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
8379            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
8380            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
8381            pending = Some((x1, ffn_out));
8382        }
8383        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
8384        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
8385        if let Some((x1p, f1p)) = pending.take() {
8386            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8387            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8388            x = x2;
8389        }
8390        Ok(x)
8391    }
8392    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
8393    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
8394    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
8395    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
8396    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
8397    /// ssm state exactly like T sequential decode steps.
8398    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
8399    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
8400    #[allow(clippy::too_many_arguments)]
8401    fn linear_attn_verify_t(
8402        &self,
8403        e: &Engine,
8404        la: &LinearAttnLayer,
8405        h: &CudaSlice<f32>,
8406        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8407        t: usize,
8408        cache: &mut Cache,
8409        il: usize,
8410        want_stash: bool,
8411    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
8412        let cfg = &self.cfg;
8413        let geometry = la.geometry;
8414        let d_state = geometry.key_head_dim as usize;
8415        let num_k = geometry.key_heads as usize;
8416        let num_v = geometry.value_heads as usize;
8417        let d_conv = geometry.conv_kernel as usize;
8418        let key_dim = d_state * num_k;
8419        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
8420        let eps = cfg.rms_eps;
8421        let scale = 1.0 / (d_state as f32).sqrt();
8422
8423        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
8424        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
8425        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
8426        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
8427        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
8428        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
8429        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
8430        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
8431        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
8432        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
8433        // Bit-identical per (tensor,token,row) — see spec_fused_t().
8434        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
8435        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
8436        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
8437        // and feeds every projection; the caller guaranteed all four input projections are
8438        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
8439        let h_q8_t = if h_q8.is_none()
8440            && spec_fused_t()
8441            && (2..=4).contains(&t)
8442            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
8443                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
8444        {
8445            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
8446        } else {
8447            None
8448        };
8449        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
8450        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
8451            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
8452        let (qkv_mixed, z) = {
8453            let mut fused = None;
8454            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
8455                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8456                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
8457            } else if let Some((hq, hd)) = hq8_any {
8458                if spec_fused_t() && (2..=4).contains(&t) {
8459                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
8460                }
8461            }
8462            match (fused, hq8_any) {
8463                (Some(pair), _) => pair,
8464                (None, Some((hq, hd))) if h_q8.is_some() => (
8465                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
8466                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
8467                ),
8468                (None, _) => (
8469                    e.matmul_decode_exact(&la.wqkv, h, t)?,
8470                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
8471                ),
8472            }
8473        };
8474        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
8475        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
8476        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
8477        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
8478        let (beta_raw, alpha) = if t == 1 {
8479            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8480            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
8481                Some(((mut b, bs), (mut a, as_))) => {
8482                    if bs != 1.0 {
8483                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
8484                    }
8485                    if as_ != 1.0 {
8486                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
8487                    }
8488                    (b, a)
8489                }
8490                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
8491                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
8492                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
8493                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
8494                    Some((b, a)) => (b, a),
8495                    None => (
8496                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
8497                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
8498                    ),
8499                },
8500            }
8501        } else {
8502            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
8503            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
8504            let mut nvfp4_fused = None;
8505            let mut q8_fused = None;
8506            if let Some((hq, hd)) = hq8_any {
8507                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
8508                    nvfp4_fused =
8509                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8510                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
8511                        static ONCE: std::sync::Once = std::sync::Once::new();
8512                        ONCE.call_once(|| {
8513                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
8514                        });
8515                    }
8516                }
8517                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
8518                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8519                }
8520            }
8521            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
8522                if bs != 1.0 {
8523                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
8524                }
8525                if as_ != 1.0 {
8526                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
8527                }
8528                (b, a)
8529            } else if let Some(pair) = q8_fused {
8530                pair
8531            } else {
8532                match hq8_any {
8533                    Some((hq, hd)) if h_q8.is_some() => (
8534                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
8535                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
8536                    ),
8537                    _ => (
8538                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
8539                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
8540                    ),
8541                }
8542            }
8543        };
8544
8545        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
8546        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
8547        let rl = cache.recur[il].as_mut().unwrap();
8548        let mut conv_out = e.uninit(conv_dim * t)?;
8549        e.ssm_conv1d_tm_state(
8550            &qkv_mixed,
8551            &mut rl.conv_state,
8552            la.ssm_conv1d.float_data(),
8553            &mut conv_out,
8554            conv_dim,
8555            t,
8556            d_conv,
8557        )?;
8558
8559        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
8560        let mut q_g = e.uninit(d_state * num_v * t)?;
8561        let mut k_g = e.uninit(d_state * num_v * t)?;
8562        let mut v_g = e.uninit(d_state * num_v * t)?;
8563        e.qkv_to_gdn_repack(
8564            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
8565        )?;
8566        let mut q_l2 = e.uninit(d_state * num_v * t)?;
8567        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8568        let mut k_l2 = e.uninit(d_state * num_v * t)?;
8569        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8570        let mut beta = e.uninit(t * num_v)?;
8571        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8572        let mut g_log = e.uninit(t * num_v)?;
8573        e.gdn_glog(
8574            &alpha,
8575            la.ssm_dt.float_data(),
8576            la.ssm_a.float_data(),
8577            &mut g_log,
8578            num_v,
8579            t,
8580        )?;
8581
8582        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
8583        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
8584        let mut o = e.uninit(d_state * num_v * t)?;
8585        {
8586            let crate::cache::RecurLayer {
8587                ssm_state,
8588                ssm_state_alt,
8589                ..
8590            } = rl;
8591            e.gdn_scan_s128(
8592                &q_l2,
8593                &k_l2,
8594                &v_g,
8595                &g_log,
8596                &beta,
8597                ssm_state,
8598                ssm_state_alt,
8599                &mut o,
8600                num_v,
8601                t,
8602                scale,
8603            )?;
8604        }
8605        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8606
8607        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
8608        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
8609        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
8610        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
8611        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
8612        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
8613        let out = if e.uses_q8_1_fast(&la.ssm_out) {
8614            let (gq, gd) =
8615                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
8616            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
8617        } else {
8618            let mut gn = e.uninit(d_state * num_v * t)?;
8619            e.gated_rmsnorm(
8620                &o,
8621                la.ssm_norm.float_data(),
8622                &z,
8623                &mut gn,
8624                d_state,
8625                num_v * t,
8626                eps,
8627            )?;
8628            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
8629            // would fall to dp4a with a different FP reduction order — same class of bug as
8630            // the input projs).
8631            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
8632        };
8633        let stash = if want_stash {
8634            Some(GdnStash {
8635                qkv_mixed,
8636                q_l2,
8637                k_l2,
8638                v_g,
8639                g_log,
8640                beta,
8641            })
8642        } else {
8643            None
8644        };
8645        Ok((out, stash))
8646    }
8647
8648    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
8649    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
8650    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
8651    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
8652    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
8653    ///   replaying them.
8654    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
8655    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
8656    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
8657    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
8658    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
8659    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
8660    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
8661    fn commit_verified_prefix(
8662        &self,
8663        e: &Engine,
8664        cache: &mut Cache,
8665        snap: &crate::cache::CacheSnapshot,
8666        ckpt: &VerifyCkpt,
8667        j: usize,
8668        kv_lens_done: bool,
8669        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
8670    ) -> Result<(), Box<dyn std::error::Error>> {
8671        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8672        // recurrent state and must never be forced through a synthetic SSM geometry.
8673        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8674        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8675        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8676        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8677        // buffers and stream order are identical to the per-layer memcpy sequence; the
8678        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
8679        let mut batched_cols = false;
8680        if state_copy_batch_on() && dev_j.is_none() {
8681            use cudarc::driver::DevicePtr;
8682            let s = &e.gpu.stream();
8683            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8684            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8685            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8686            let mut uniform = true;
8687            for il in 0..self.layers.len() {
8688                let Some(rl) = cache.recur[il].as_ref() else {
8689                    continue;
8690                };
8691                if ckpt.gdn[il].is_some() {
8692                    continue; // kernel-rebuild arm restores below, per layer
8693                }
8694                let Some(cols) = &ckpt.cols[il] else {
8695                    continue; // missing-ckpt error surfaces in the main loop
8696                };
8697                let (c, st) = &cols[j - 1];
8698                if conv_pairs.is_empty() {
8699                    conv_words = c.len();
8700                    ssm_words = st.len();
8701                } else if c.len() != conv_words || st.len() != ssm_words {
8702                    uniform = false;
8703                    break;
8704                }
8705                let (pc, _g0) = c.device_ptr(s);
8706                let (dc, _g1) = rl.conv_state.device_ptr(s);
8707                let (ps, _g2) = st.device_ptr(s);
8708                let (ds, _g3) = rl.ssm_state.device_ptr(s);
8709                conv_pairs.push((pc as u64, dc as u64));
8710                ssm_pairs.push((ps as u64, ds as u64));
8711            }
8712            if uniform && !conv_pairs.is_empty() {
8713                let n = conv_pairs.len();
8714                let mut t = vec![0u64; 2 * n];
8715                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8716                    t[k] = src;
8717                    t[n + k] = dst;
8718                }
8719                let conv_t = e.htod_u64(&t)?;
8720                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8721                    t[k] = src;
8722                    t[n + k] = dst;
8723                }
8724                let ssm_t = e.htod_u64(&t)?;
8725                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8726                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8727                batched_cols = true;
8728            }
8729        }
8730        rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
8731        for il in 0..self.layers.len() {
8732            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8733                kvl.len = saved + j;
8734                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
8735                if !kv_lens_done {
8736                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8737                }
8738            }
8739            if let Some(rl) = cache.recur[il].as_mut() {
8740                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8741                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8742                };
8743                let geometry = linear.geometry;
8744                let d_state = geometry.key_head_dim as usize;
8745                let num_k = geometry.key_heads as usize;
8746                let num_v = geometry.value_heads as usize;
8747                let d_conv = geometry.conv_kernel as usize;
8748                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8749                let scale = 1.0 / (d_state as f32).sqrt();
8750                if let Some(st) = &ckpt.gdn[il] {
8751                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8752                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8753                    if let Some((acc, base, t_v)) = dev_j {
8754                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
8755                        e.ssm_conv_ring_rebuild_dc(
8756                            &st.qkv_mixed,
8757                            ring_old,
8758                            &mut rl.conv_state,
8759                            conv_dim,
8760                            acc,
8761                            base,
8762                            t_v,
8763                            d_conv,
8764                        )?;
8765                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
8766                        e.gdn_scan_s128_dc(
8767                            &st.q_l2,
8768                            &st.k_l2,
8769                            &st.v_g,
8770                            &st.g_log,
8771                            &st.beta,
8772                            state_in,
8773                            &mut rl.ssm_state,
8774                            &mut o,
8775                            num_v,
8776                            acc,
8777                            base,
8778                            t_v,
8779                            scale,
8780                        )?;
8781                    } else {
8782                        e.ssm_conv_ring_rebuild(
8783                            &st.qkv_mixed,
8784                            ring_old,
8785                            &mut rl.conv_state,
8786                            conv_dim,
8787                            j,
8788                            d_conv,
8789                        )?;
8790                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8791                        e.gdn_scan_s128(
8792                            &st.q_l2,
8793                            &st.k_l2,
8794                            &st.v_g,
8795                            &st.g_log,
8796                            &st.beta,
8797                            state_in,
8798                            &mut rl.ssm_state,
8799                            &mut o,
8800                            num_v,
8801                            j,
8802                            scale,
8803                        )?;
8804                    }
8805                } else if let Some(cols) = &ckpt.cols[il] {
8806                    if !batched_cols {
8807                        let (c, s) = &cols[j - 1];
8808                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8809                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8810                    }
8811                } else {
8812                    return Err(
8813                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8814                    );
8815                }
8816            }
8817        }
8818        cache.pos = snap.pos + j;
8819        Ok(())
8820    }
8821
8822    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8823    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8824    fn commit_verified_prefix_stream(
8825        &self,
8826        e: &Engine,
8827        cache: &mut Cache,
8828        snap: &crate::cache::CacheSnapshot,
8829        ckpt: &VerifyCkpt,
8830        acc: &CudaSlice<u32>,
8831        base: usize,
8832        t_v: usize,
8833    ) -> Result<(), Box<dyn std::error::Error>> {
8834        for il in 0..self.layers.len() {
8835            if let Some(rl) = cache.recur[il].as_mut() {
8836                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8837                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8838                };
8839                let geometry = linear.geometry;
8840                let d_state = geometry.key_head_dim as usize;
8841                let num_k = geometry.key_heads as usize;
8842                let num_v = geometry.value_heads as usize;
8843                let d_conv = geometry.conv_kernel as usize;
8844                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8845                let scale = 1.0 / (d_state as f32).sqrt();
8846                let st = ckpt.gdn[il]
8847                    .as_ref()
8848                    .ok_or("stream restore: batched-linear stash missing")?;
8849                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8850                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8851                e.ssm_conv_ring_rebuild_dc(
8852                    &st.qkv_mixed,
8853                    ring_old,
8854                    &mut rl.conv_state,
8855                    conv_dim,
8856                    acc,
8857                    base,
8858                    t_v,
8859                    d_conv,
8860                )?;
8861                let mut o = e.uninit(d_state * num_v * t_v)?;
8862                e.gdn_scan_s128_dc(
8863                    &st.q_l2,
8864                    &st.k_l2,
8865                    &st.v_g,
8866                    &st.g_log,
8867                    &st.beta,
8868                    state_in,
8869                    &mut rl.ssm_state,
8870                    &mut o,
8871                    num_v,
8872                    acc,
8873                    base,
8874                    t_v,
8875                    scale,
8876                )?;
8877            }
8878        }
8879        Ok(())
8880    }
8881
8882    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8883    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8884    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8885    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8886    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8887    pub fn decode_step_t_aux2(
8888        &self,
8889        e: &Engine,
8890        tokens: &[u32],
8891        pos0: usize,
8892        cache: &mut Cache,
8893        aux_layers: &[usize],
8894        pred_col: Option<usize>,
8895    ) -> Result<
8896        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8897        Box<dyn std::error::Error>,
8898    > {
8899        let cfg = &self.cfg;
8900        let n_embd = cfg.n_embd as usize;
8901        let eps = cfg.rms_eps;
8902        let t = tokens.len();
8903        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8904        let pos_d = e.htod_i32(&pos_vec)?;
8905        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
8906        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8907        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8908        let want_pred = pred_col.is_some();
8909
8910        for (il, layer) in self.layers.iter().enumerate() {
8911            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8912            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8913            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8914            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8915            if norm_fused {
8916                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8917            } else {
8918                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8919            }
8920            let mixed = match &layer.mixer {
8921                Mixer::Full(fa) => {
8922                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8923                }
8924                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8925                Mixer::Linear(la) => {
8926                    let mut out = e.zeros(t * n_embd)?;
8927                    for col in 0..t {
8928                        let mut h_col = e.zeros(n_embd)?;
8929                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
8930                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8931                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8932                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8933                    }
8934                    out
8935                }
8936            };
8937            let ffn_fuse = match &layer.ffn {
8938                crate::hybrid::Ffn::Dense {
8939                    ffn_gate, ffn_up, ..
8940                } => {
8941                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8942                        && e.uses_q8_1_fast(ffn_gate)
8943                        && e.uses_q8_1_fast(ffn_up)
8944                }
8945                crate::hybrid::Ffn::Moe(_) => false,
8946            };
8947            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8948            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8949            if ffn_fuse {
8950                e.add(&x, &mixed, &mut x1, t * n_embd)?;
8951                e.rms_norm_decode(
8952                    &x1,
8953                    layer.post_attn_norm.float_data(),
8954                    &mut z,
8955                    n_embd,
8956                    t,
8957                    eps,
8958                )?;
8959            } else {
8960                e.add_rms_norm(
8961                    &x,
8962                    &mixed,
8963                    layer.post_attn_norm.float_data(),
8964                    &mut x1,
8965                    &mut z,
8966                    n_embd,
8967                    t,
8968                    eps,
8969                )?;
8970            }
8971            let ffn_out = match &layer.ffn {
8972                crate::hybrid::Ffn::Dense {
8973                    ffn_gate,
8974                    ffn_up,
8975                    ffn_down,
8976                } => {
8977                    let n_ff = ffn_gate.out_features();
8978                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
8979                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
8980                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8981                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
8982                    Self::ffn_act_lim(
8983                        e,
8984                        &self.cfg,
8985                        &gate,
8986                        &up,
8987                        1.0,
8988                        1.0,
8989                        self.cfg.clamp_shexp_at(il as u32),
8990                        &mut act,
8991                        t * n_ff,
8992                    )?;
8993                    e.matmul_decode_exact(ffn_down, &act, t)?
8994                }
8995                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8996            };
8997            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8998            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8999            if aux_layers.contains(&il) {
9000                let mut a = e.zeros(n_embd)?;
9001                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9002                aux_last.push(a);
9003                if let Some(pc) = pred_col {
9004                    let mut ap = e.zeros(n_embd)?;
9005                    e.copy_view_into(
9006                        &mut ap,
9007                        0,
9008                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9009                        n_embd,
9010                    )?;
9011                    aux_pred.push(ap);
9012                }
9013            }
9014            x = x2;
9015        }
9016        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9017        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9018        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9019        let host = e.dtoh(&logits)?;
9020        cache.pos += t;
9021        Ok((
9022            host,
9023            aux_last,
9024            if want_pred { Some(aux_pred) } else { None },
9025        ))
9026    }
9027
9028    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9029    /// `step35_decode_attn`.
9030    ///
9031    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9032    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9033    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9034    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9035    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9036    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9037    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9038    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9039    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9040    /// position of each query row. A batched twin would have to reproduce all of that AND the
9041    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9042    /// take one `base_len`, not a per-row offset).
9043    ///
9044    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9045    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9046    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9047    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9048    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9049    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9050    /// step35 twin is a perf lane's job and must be gated against this arm.
9051    ///
9052    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9053    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9054    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9055    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9056    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9057    #[allow(clippy::too_many_arguments)]
9058    fn step35_verify(
9059        &self,
9060        e: &Engine,
9061        fa: &FullAttnLayer,
9062        h: &CudaSlice<f32>,
9063        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9064        t: usize,
9065        cache: &mut Cache,
9066        il: usize,
9067    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9068        let n_embd = self.cfg.n_embd as usize;
9069        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9070        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9071        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9072        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9073        // cannot regress it into silently reading an empty buffer.
9074        assert_eq!(
9075            h.len(),
9076            t * n_embd,
9077            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9078             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9079            h_q8.is_some()
9080        );
9081        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9082        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9083        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9084        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9085        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9086        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9087        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9088        for r in 0..t {
9089            // Absolute position of this query row. `cache.pos` is the committed length at round
9090            // start and every row before r has already been appended by this loop, so the r-th
9091            // verify token sits at cache.pos + r — the same position eager decode would give it.
9092            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9093            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9094            e.copy_view_into(
9095                &mut h_row,
9096                0,
9097                &h.slice(r * n_embd..(r + 1) * n_embd),
9098                n_embd,
9099            )?;
9100            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9101            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9102            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9103            debug_assert_eq!(
9104                o.len(),
9105                n_embd,
9106                "step35_decode_attn returns post-wo [n_embd]"
9107            );
9108            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9109        }
9110        Ok(out)
9111    }
9112
9113    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9114    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9115    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9116    #[allow(clippy::too_many_arguments)]
9117    fn full_attn_verify(
9118        &self,
9119        e: &Engine,
9120        fa: &FullAttnLayer,
9121        h: &CudaSlice<f32>,
9122        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9123        pos_d: &CudaSlice<i32>,
9124        t: usize,
9125        cache: &mut Cache,
9126        il: usize,
9127        stream_ctr: Option<&CudaSlice<i32>>,
9128    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9129        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9130        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9131        // its own arm. A verify that silently computes different attention than decode defeats the
9132        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9133        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9134        // shape and not laziness.
9135        if self.sliding_gated_moe_batch_program() {
9136            if stream_ctr.is_some() {
9137                return Err(
9138                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9139                            cannot express the SWA offset KV view; same root cause as the dc \
9140                            decode refusal) — run spec without the stream arm"
9141                        .into(),
9142                );
9143            }
9144            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9145        }
9146        let cfg = &self.cfg;
9147        let geometry = cfg.full_attention_geometry_at(il as u32);
9148        let n_head = geometry.n_head as usize;
9149        let n_head_kv = geometry.n_head_kv as usize;
9150        let head_dim = geometry.head_dim_k as usize;
9151        let eps = cfg.rms_eps;
9152        let scale = geometry.attention_scale();
9153        let n_embd = cfg.n_embd as usize;
9154
9155        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9156        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9157        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9158        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9159        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9160        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9161        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9162        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9163        let (qf, mut k, v) = {
9164            let mut fused = None;
9165            let qkv_fast =
9166                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9167            if t == 1 && qkv_fast {
9168                let (hq_o, hd_o);
9169                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9170                    Some(p) => p,
9171                    None => {
9172                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9173                        (&hq_o, &hd_o)
9174                    }
9175                };
9176                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9177            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9178                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9179                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9180                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9181                let (hq_o, hd_o);
9182                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9183                    Some(p) => p,
9184                    None => {
9185                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9186                        (&hq_o, &hd_o)
9187                    }
9188                };
9189                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9190            }
9191            match (fused, h_q8) {
9192                (Some(triple), _) => triple,
9193                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9194                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9195                (None, Some((hq, hd))) if qkv_fast => (
9196                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9197                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9198                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9199                ),
9200                (None, _) => (
9201                    e.matmul_decode_exact(&fa.wq, h, t)?,
9202                    e.matmul_decode_exact(&fa.wk, h, t)?,
9203                    e.matmul_decode_exact(&fa.wv, h, t)?,
9204                ),
9205            }
9206        };
9207        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9208        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9209        let (mut q, gate) = if gated {
9210            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9211            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9212            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9213            (q, Some(gate))
9214        } else {
9215            (qf, None)
9216        };
9217
9218        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9219        e.rms_norm(
9220            &q,
9221            fa.q_norm.float_data(),
9222            &mut qn,
9223            head_dim,
9224            n_head * t,
9225            eps,
9226        )?;
9227        q = qn;
9228        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9229        e.rms_norm(
9230            &k,
9231            fa.k_norm.float_data(),
9232            &mut kn,
9233            head_dim,
9234            n_head_kv * t,
9235            eps,
9236        )?;
9237        k = kn;
9238        let rope_dims = geometry.n_rot as usize;
9239        e.rope_neox(
9240            &mut q,
9241            pos_d,
9242            head_dim,
9243            rope_dims,
9244            n_head,
9245            t,
9246            geometry.rope_base,
9247            1.0,
9248        )?;
9249        e.rope_neox(
9250            &mut k,
9251            pos_d,
9252            head_dim,
9253            rope_dims,
9254            n_head_kv,
9255            t,
9256            geometry.rope_base,
9257            1.0,
9258        )?;
9259
9260        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9261        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9262        let kvl = cache.kv[il].as_mut().unwrap();
9263        let (kv_dim_k, kv_dim_v, ktb, vtb) =
9264            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9265        if let Some(ctr) = stream_ctr {
9266            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9267            // math on a (block, token) grid, documented byte-identical); host len is a stale
9268            // LOWER BOUND under pre-issue (drain reconciles it).
9269            e.append_kv_quantized_rows_dc(
9270                &k,
9271                &v,
9272                &mut kvl.k,
9273                &mut kvl.v,
9274                ctr,
9275                t,
9276                kv_dim_k,
9277                kv_dim_v,
9278                ktb,
9279                vtb,
9280                crate::Engine::kv_fp8_on(),
9281            )?;
9282        } else {
9283            for i in 0..t {
9284                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
9285                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
9286                e.append_kv_quantized_view(
9287                    &k_row,
9288                    &v_row,
9289                    &mut kvl.k,
9290                    &mut kvl.v,
9291                    kvl.len + i,
9292                    kv_dim_k,
9293                    kv_dim_v,
9294                    ktb,
9295                    vtb,
9296                    crate::Engine::kv_fp8_on(),
9297                )?;
9298            }
9299            kvl.len += t;
9300        }
9301
9302        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
9303        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
9304        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
9305        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
9306        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
9307        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
9308        // keys. The verify appends all T tokens first but bounds the key range per row.
9309        //
9310        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
9311        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
9312        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
9313        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
9314        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
9315        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
9316        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
9317        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
9318        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
9319        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
9320        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
9321        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
9322        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
9323        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
9324        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
9325        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
9326        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
9327        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
9328        if let Some(ctr) = stream_ctr {
9329            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
9330            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
9331            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
9332            let upper = kvl.len + t + 64;
9333            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
9334            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
9335            e.fa_decode_rows_dc(
9336                &q,
9337                &k_view,
9338                &v_view,
9339                &mut attn,
9340                head_dim,
9341                n_head,
9342                n_head_kv,
9343                ctr,
9344                upper.min(cache.max_ctx),
9345                t,
9346                scale,
9347                ktb,
9348                vtb,
9349                0,
9350                false,
9351            )?;
9352        } else if spec_lean() && t == 1 {
9353            let t_kv = base_len + 1;
9354            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
9355            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
9356            e.fa_decode_kvmod(
9357                &q,
9358                &k_view,
9359                &v_view,
9360                &mut attn,
9361                head_dim,
9362                n_head,
9363                n_head_kv,
9364                t_kv,
9365                scale,
9366                ktb,
9367                vtb,
9368                crate::Engine::kv_fp8_on(),
9369            )?;
9370        } else if e.fa_rows_eligible(base_len, head_dim) {
9371            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
9372            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
9373            e.fa_decode_rows(
9374                &q,
9375                &k_view,
9376                &v_view,
9377                &mut attn,
9378                head_dim,
9379                n_head,
9380                n_head_kv,
9381                base_len,
9382                t,
9383                scale,
9384                ktb,
9385                vtb,
9386                None,
9387                false,
9388                crate::Engine::kv_fp8_on(),
9389                None,
9390            )?;
9391        } else {
9392            for r in 0..t {
9393                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
9394                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
9395                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
9396                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
9397                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
9398                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
9399                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
9400                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
9401                e.fa_decode_kvmod(
9402                    &q_row,
9403                    &k_view_r,
9404                    &v_view_r,
9405                    &mut attn_row,
9406                    head_dim,
9407                    n_head,
9408                    n_head_kv,
9409                    t_kv_r,
9410                    scale,
9411                    ktb,
9412                    vtb,
9413                    crate::Engine::kv_fp8_on(),
9414                )?;
9415                e.copy_into(
9416                    &mut attn,
9417                    r * n_head * head_dim,
9418                    &attn_row,
9419                    n_head * head_dim,
9420                )?;
9421            }
9422        }
9423
9424        let attn_g = match &gate {
9425            Some(gate) => {
9426                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
9427                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
9428                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
9429                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
9430                ag
9431            }
9432            None => attn,
9433        };
9434        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
9435        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
9436        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
9437    }
9438
9439    /// Context-linear bytes for a plain serving session's trunk cache.
9440    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
9441        crate::cache::cache_bytes_per_token_for_plan(
9442            &self.cfg,
9443            &self.plan,
9444            0,
9445            self.plan.layers.len(),
9446        )
9447    }
9448
9449    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
9450    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
9451        (
9452            self.plain_session_kv_bytes_per_token(),
9453            crate::cache::cache_ring_bytes_per_token_for_plan(
9454                &self.cfg,
9455                &self.plan,
9456                0,
9457                self.plan.layers.len(),
9458            ),
9459            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
9460        )
9461    }
9462
9463    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
9464    /// scratch. With no MTP head this equals the plain coefficient.
9465    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
9466        let scratch = self
9467            .mtp
9468            .iter()
9469            .chain(self.mtp_extra.iter())
9470            .map(|mtp| {
9471                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9472                k + v
9473            })
9474            .sum::<usize>();
9475        self.plain_session_kv_bytes_per_token()
9476            .saturating_add(scratch)
9477    }
9478
9479    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
9480    /// capped by the same SWA ring rows as the trunk.
9481    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
9482        let total = self.spec_session_kv_bytes_per_token();
9483        let (_, mut ring, rows) = self.plain_session_kv_shape();
9484        if rows > 0 {
9485            ring = ring.saturating_add(
9486                self.mtp
9487                    .iter()
9488                    .chain(self.mtp_extra.iter())
9489                    .map(|mtp| {
9490                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9491                        k + v
9492                    })
9493                    .sum::<usize>(),
9494            );
9495        }
9496        (total, ring, rows)
9497    }
9498
9499    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
9500    /// the NextN head to draft K tokens then verifies them in one batched target forward.
9501    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
9502    /// acceptance rate. `k` = draft length per round.
9503    ///
9504    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
9505    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
9506    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
9507    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
9508    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
9509    /// captured graph references is event-free; the spec loop is strictly single-stream.
9510    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
9511    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
9512    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
9513    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
9514    /// generate_spec_inner2.
9515    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
9516    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
9517    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
9518    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
9519    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
9520    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
9521    pub fn new_session(
9522        &self,
9523        e: &Engine,
9524        max_ctx: usize,
9525    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
9526        Ok(SpecSession {
9527            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
9528            // is the SERVING spec-session path, and with the ppN door open across two cards a
9529            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
9530            // round — the wrong-card class already fixed on the two batched serving paths
9531            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
9532            // branch, same allocations), so single-device behavior is byte-unchanged.
9533            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
9534            scratch: self.new_mtp_scratch(e, max_ctx)?,
9535            committed: Vec::new(),
9536            last_h: None,
9537            next_pred: None,
9538            sctr: 0,
9539            uctr: 0,
9540            draft_ctx: None,
9541            pending_tok: None,
9542            turn_ckpt: None,
9543            telem: SpecTelemetryCounters::default(),
9544            capture_at: None,
9545            boundary_captures: Vec::new(),
9546            ckpt_at: None,
9547        })
9548    }
9549
9550    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
9551    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
9552    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
9553    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
9554    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
9555    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
9556    /// worker always receives a fully-warm continuation session (committed = whole
9557    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
9558    /// boundary logits on the empty-suffix shape).
9559    ///
9560    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
9561    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
9562    /// request, and plain feeds a carried suffix via eager `decode_step` below
9563    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
9564    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
9565    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
9566    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
9567    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
9568    /// burst prime.
9569    ///
9570    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
9571    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
9572    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
9573    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
9574    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
9575    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
9576    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
9577    /// cold session draws from the identical row at counter 0 and then runs its rounds from
9578    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
9579    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
9580    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
9581    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
9582    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
9583    ///
9584    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
9585    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
9586    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
9587    /// and are never routed here.
9588    ///
9589    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
9590    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
9591    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
9592    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
9593    /// entry stays published for the next request.
9594    #[allow(clippy::too_many_arguments)]
9595    pub fn spec_session_from_restored(
9596        &self,
9597        e: &Engine,
9598        mut cache: Cache,
9599        prefix: Vec<u32>,
9600        suffix: &[u32],
9601        draft_k: &CudaSlice<u8>,
9602        draft_v: &CudaSlice<u8>,
9603        draft_k_tok_bytes: usize,
9604        draft_v_tok_bytes: usize,
9605        draft_len: usize,
9606        last_h: &[f32],
9607        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
9608        // when a suffix follows — the feed's own logits are the boundary then.
9609        boundary_logits: &[f32],
9610        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
9611        // ONE place instead of being half-applied by the worker.
9612        sampling: Option<SpecSampling>,
9613        require_anchor: bool,
9614        max_ctx: usize,
9615        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
9616        // prompt position to split the suffix feed at and capture the extended-entry
9617        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
9618        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
9619        // WHY: the prompt-end capture below includes the template's live generation header
9620        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
9621        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
9622        // diverged from every future prompt and the hit boundary FROZE at the first
9623        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
9624        republish_at: Option<usize>,
9625    ) -> Result<SpecSession, (Option<Cache>, String)> {
9626        let pos = prefix.len();
9627        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
9628            Err((Some(cache), msg))
9629        };
9630        if self.mtp.is_none() {
9631            return fail(cache, "no MTP head attached (nothing to draft with)".into());
9632        }
9633        if pos == 0 {
9634            return fail(cache, "empty committed prefix".into());
9635        }
9636        if cache.pos != pos {
9637            let msg = format!(
9638                "restored cache pos {} != restored prefix len {pos}",
9639                cache.pos
9640            );
9641            return fail(cache, msg);
9642        }
9643        if draft_len != pos {
9644            return fail(
9645                cache,
9646                format!("draft plane len {draft_len} != restored prefix len {pos}"),
9647            );
9648        }
9649        if pos + suffix.len() >= max_ctx {
9650            return fail(
9651                cache,
9652                format!(
9653                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
9654                    pos + suffix.len(),
9655                ),
9656            );
9657        }
9658        let mut scratch = match MtpScratch::new(
9659            e,
9660            &self.cfg,
9661            &self.plan,
9662            max_ctx,
9663            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9664        ) {
9665            Ok(s) => s,
9666            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
9667        };
9668        if scratch.kv.ring.is_some() {
9669            return fail(
9670                cache,
9671                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
9672            );
9673        }
9674        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
9675            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
9676        {
9677            return fail(
9678                cache,
9679                format!(
9680                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
9681                     {}/{} bytes/token (stale entry across a format change)",
9682                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
9683                ),
9684            );
9685        }
9686        if pos > scratch.cap {
9687            return fail(
9688                cache,
9689                format!(
9690                    "draft plane rows {pos} exceed scratch capacity {}",
9691                    scratch.cap
9692                ),
9693            );
9694        }
9695        let kb = pos * draft_k_tok_bytes;
9696        let vb = pos * draft_v_tok_bytes;
9697        if draft_k.len() < kb || draft_v.len() < vb {
9698            return fail(
9699                cache,
9700                format!(
9701                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
9702                    draft_k.len(),
9703                    draft_v.len(),
9704                ),
9705            );
9706        }
9707        if kb > 0 {
9708            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
9709                return fail(cache, format!("draft K restore copy failed: {err}"));
9710            }
9711        }
9712        if vb > 0 {
9713            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
9714                return fail(cache, format!("draft V restore copy failed: {err}"));
9715            }
9716        }
9717        if let Err(err) = scratch.set_len(e, pos) {
9718            return fail(cache, format!("draft scratch len set failed: {err}"));
9719        }
9720        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9721            // anchor upload failure is acceptance-only when a suffix feed follows (fill
9722            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9723            // burst entry asserts committed + last_h + next_pred) — the caller says which.
9724            e.htod(last_h).ok()
9725        } else {
9726            None
9727        };
9728        if require_anchor && last_h_dev.is_none() {
9729            return fail(
9730                cache,
9731                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9732            );
9733        }
9734        let mut committed = prefix;
9735        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9736        // what the empty-suffix continuation assert in the burst entry requires.
9737        let next_pred;
9738        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9739        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9740        // drawing its own first token from the same row.
9741        let mut sctr = 0u32;
9742        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9743        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9744        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9745        // after the suffix joins `committed` below.
9746        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9747        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9748        if !suffix.is_empty() {
9749            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9750            // From here on the trunk cache mutates: failures return Err((None, _)) and
9751            // the worker serves the request cold-plain instead of reusing the carrier.
9752            let dirty =
9753                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9754            let n_embd = self.cfg.n_embd as usize;
9755            let t = suffix.len();
9756            let mut h_rows = match e.uninit(t * n_embd) {
9757                Ok(b) => b,
9758                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9759            };
9760            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9761            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9762            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9763            let b_rel = republish_at
9764                .and_then(|abs| abs.checked_sub(pos))
9765                .filter(|&r| r > 0 && r < t);
9766            let mut feed_logits = Vec::new();
9767            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9768                || e.frozen_cpu_experts_prefer_tokenwise_prime();
9769            let mut fed = 0usize;
9770            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9771                if seg_end <= fed {
9772                    continue;
9773                }
9774                let seg = &suffix[fed..seg_end];
9775                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9776                if batched {
9777                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
9778                    // queued after this segment ride `queued_after` so Step35 arm selection
9779                    // stays keyed to the request's end (tick-seg law).
9780                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9781                        Ok((l, _h_seed, hiddens)) => {
9782                            if let Err(err) =
9783                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9784                            {
9785                                return dirty(format!("suffix hidden copy: {err}"));
9786                            }
9787                            feed_logits = l;
9788                        }
9789                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
9790                    }
9791                } else {
9792                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9793                    for (i, &tok) in seg.iter().enumerate() {
9794                        match self.decode_step_h(e, tok, &mut cache) {
9795                            Ok((l, h)) => {
9796                                if let Err(err) =
9797                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9798                                {
9799                                    return dirty(format!("suffix hidden copy: {err}"));
9800                                }
9801                                feed_logits = l;
9802                            }
9803                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9804                        }
9805                    }
9806                }
9807                fed = seg_end;
9808                if Some(seg_end) == b_rel {
9809                    // The stable pre-generation boundary: capture the extended-entry
9810                    // publication AND this session's own turn checkpoint here instead of at
9811                    // prompt-end (both would otherwise carry the volatile live-header tail
9812                    // the next re-render replaces). Failure silent, turn_ckpt convention.
9813                    debug_assert_eq!(
9814                        cache.pos,
9815                        pos + seg_end,
9816                        "stable-boundary capture off the feed split"
9817                    );
9818                    if spec_restore_republish_on() {
9819                        if let Ok(snap) = cache.snapshot(e) {
9820                            boundary_captures.push(SpecBoundaryCapture {
9821                                snap,
9822                                pos: pos + seg_end,
9823                                logits: feed_logits.clone(),
9824                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9825                            });
9826                        }
9827                    }
9828                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9829                        e.uninit(n_embd).and_then(|mut a| {
9830                            e.copy_view_into(
9831                                &mut a,
9832                                0,
9833                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9834                                n_embd,
9835                            )?;
9836                            Ok(a)
9837                        });
9838                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9839                        restored_turn_ckpt = Some(SpecCheckpoint {
9840                            snap,
9841                            pos: pos + seg_end,
9842                            last_h,
9843                        });
9844                    }
9845                }
9846            }
9847            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9848            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9849            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9850            // with T). Fill failures are acceptance-only — truncate to the restored rows
9851            // and continue; the burst's own set_len keeps the invariant.
9852            let mtp = self.mtp.as_ref().expect("mtp checked above");
9853            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9854            let embd_gpu = if spec_host_embd() {
9855                None
9856            } else {
9857                Some(
9858                    self.embd_gpu
9859                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9860                )
9861            };
9862            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9863            let fill_chunk = 4096usize;
9864            let mut filled = true;
9865            let mut start = 0usize;
9866            'fill: while start < t {
9867                let end = (start + fill_chunk).min(t);
9868                let tc = end - start;
9869                let Ok(mut phs) = e.zeros(tc * n_embd) else {
9870                    filled = false;
9871                    break 'fill;
9872                };
9873                let (src_lo, dst_off, n_copy) = if start == 0 {
9874                    (0, n_embd, (tc - 1) * n_embd)
9875                } else {
9876                    ((start - 1) * n_embd, 0, tc * n_embd)
9877                };
9878                if start == 0 {
9879                    if let Some(lh) = last_h_dev.as_ref() {
9880                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
9881                            filled = false;
9882                            break 'fill;
9883                        }
9884                    }
9885                }
9886                if n_copy > 0
9887                    && e.copy_view_into(
9888                        &mut phs,
9889                        dst_off,
9890                        &h_rows.slice(src_lo..src_lo + n_copy),
9891                        n_copy,
9892                    )
9893                    .is_err()
9894                {
9895                    filled = false;
9896                    break 'fill;
9897                }
9898                if self
9899                    .mtp_kv_fill_all(
9900                        e,
9901                        &suffix[start..end],
9902                        &phs,
9903                        pos + start,
9904                        &mut scratch,
9905                        embd_dev,
9906                    )
9907                    .is_err()
9908                {
9909                    filled = false;
9910                    break 'fill;
9911                }
9912                start = end;
9913            }
9914            if !filled {
9915                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9916                // so keep only the restored rows resident and let verify arbitrate.
9917                if let Err(err) = scratch.set_len(e, pos) {
9918                    return dirty(format!("scratch truncation after failed fill: {err}"));
9919                }
9920            }
9921            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9922            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9923            // finding (d)). Pre-lane, publication was armed only for COLD sessions
9924            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9925            // non-continuation burst — but a converted hit's first burst IS a continuation,
9926            // so a growing conversation learned exactly ONE boundary and turn 3 could never
9927            // hit a longer prefix than turn 2 did.
9928            //
9929            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9930            // line — the trunk is primed over the whole prompt, nothing is generated, and the
9931            // draft plane rows [0..prompt) are filled just above. That is a complete
9932            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9933            // publishes; the worker's existing publication sweep picks it up because it is
9934            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9935            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9936            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9937            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9938            // publication is an optimization, never a correctness dependency.
9939            //
9940            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9941            // entry's tail is the live generation header the next re-render replaces, so on a
9942            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9943            // the stable-boundary capture above IS this publication, minus the poisoned tail.
9944            if spec_restore_republish_on() && boundary_captures.is_empty() {
9945                debug_assert_eq!(
9946                    cache.pos,
9947                    pos + t,
9948                    "extended-entry capture must sit at the restored session's prompt end",
9949                );
9950                if let Ok(snap) = cache.snapshot(e) {
9951                    boundary_captures.push(SpecBoundaryCapture {
9952                        snap,
9953                        pos: pos + t,
9954                        logits: feed_logits.clone(),
9955                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9956                    });
9957                }
9958            }
9959            // continuation seed: the feed's boundary logits ARE the plain path's boundary
9960            // logits (same program), so greedy's argmax here is plain's first emitted token,
9961            // and the sampled draw is the cold sampled session's own first token.
9962            next_pred = Some(if sampled {
9963                let sp = sampling.expect("sampled implies a sampler");
9964                // `committed` is still the restored prefix here; the suffix joins it below —
9965                // so this is the last-N window over the WHOLE prompt, exactly the cold
9966                // session's own window at its first token.
9967                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
9968                match sample_boundary_token(
9969                    e,
9970                    &feed_logits,
9971                    &sp,
9972                    &hist,
9973                    &mut sctr,
9974                    "restore-suffix-feed",
9975                ) {
9976                    Ok(t) => t,
9977                    // the trunk is already fed: hand nothing back, the worker serves the
9978                    // request cold-plain. Never fall back to an argmax — that would put a
9979                    // greedy token in a sampled stream to save a slow path.
9980                    Err(err) => {
9981                        return dirty(format!("boundary token draw failed: {err}"));
9982                    }
9983                }
9984            } else {
9985                argmax(&feed_logits) as u32
9986            });
9987            let mut lh = match e.uninit(n_embd) {
9988                Ok(b) => b,
9989                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
9990            };
9991            if let Err(err) = e.copy_view_into(
9992                &mut lh,
9993                0,
9994                &h_rows.slice((t - 1) * n_embd..t * n_embd),
9995                n_embd,
9996            ) {
9997                return dirty(format!("boundary hidden copy: {err}"));
9998            }
9999            last_h_dev = Some(lh);
10000            committed.extend_from_slice(suffix);
10001        } else {
10002            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10003            // ENTRY's boundary logits are the boundary row, and this is the token the cold
10004            // session emits from that same row. Owned here rather than in the worker so the
10005            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10006            if boundary_logits.is_empty() {
10007                return fail(
10008                    cache,
10009                    "full-cover restore without the entry's boundary logits".into(),
10010                );
10011            }
10012            next_pred = Some(if sampled {
10013                let sp = sampling.expect("sampled implies a sampler");
10014                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10015                match sample_boundary_token(
10016                    e,
10017                    boundary_logits,
10018                    &sp,
10019                    &hist,
10020                    &mut sctr,
10021                    "restore-full-cover",
10022                ) {
10023                    Ok(t) => t,
10024                    // nothing has been mutated on this shape — hand the carrier back and let
10025                    // the hit serve PLAIN (the banked pre-lane path).
10026                    Err(err) => {
10027                        return fail(cache, format!("boundary token draw failed: {err}"));
10028                    }
10029                }
10030            } else {
10031                argmax(boundary_logits) as u32
10032            });
10033        }
10034        Ok(SpecSession {
10035            cache,
10036            scratch,
10037            committed,
10038            last_h: last_h_dev,
10039            next_pred,
10040            sctr,
10041            uctr: 0,
10042            draft_ctx: None,
10043            pending_tok: None,
10044            // Stable-boundary capture from the split feed above (None on the legacy shape):
10045            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10046            // affinity probe declined ("no turn checkpoint retained") and the conversation
10047            // fell back to the frozen prefix entry forever.
10048            turn_ckpt: restored_turn_ckpt,
10049            telem: SpecTelemetryCounters::default(),
10050            capture_at: None,
10051            boundary_captures,
10052            ckpt_at: None,
10053        })
10054    }
10055
10056    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10057    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10058    /// snapshot, or draft-KV row that only corrupts the following round.
10059    pub fn optipipe_compare_session_state(
10060        &self,
10061        e: &Engine,
10062        reference: &SpecSession,
10063        candidate: &SpecSession,
10064    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10065        fn fail(what: &str) -> Box<dyn std::error::Error> {
10066            format!("optipipe state mismatch: {what}").into()
10067        }
10068        fn same_f32(a: &[f32], b: &[f32]) -> bool {
10069            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10070        }
10071        fn compare_layers(
10072            es: &Engine,
10073            range: std::ops::Range<usize>,
10074            reference: &SpecSession,
10075            candidate: &SpecSession,
10076            report: &mut OptiForkStateIdentity,
10077        ) -> Result<(), Box<dyn std::error::Error>> {
10078            for il in range {
10079                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10080                    (Some(a), Some(b)) => {
10081                        if a.len != b.len {
10082                            return Err(fail(&format!(
10083                                "layer {il} host KV len {} != {}",
10084                                a.len, b.len
10085                            )));
10086                        }
10087                        let ad = es.dtoh_i32(&a.len_d)?;
10088                        let bd = es.dtoh_i32(&b.len_d)?;
10089                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
10090                            return Err(fail(&format!(
10091                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10092                                a.len,
10093                            )));
10094                        }
10095                        let kb = a.len * a.k_tok_bytes;
10096                        let vb = a.len * a.v_tok_bytes;
10097                        if kb > 0 {
10098                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10099                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10100                            if ak != bk {
10101                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10102                                return Err(fail(&format!(
10103                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10104                                    at / a.k_tok_bytes,
10105                                    at % a.k_tok_bytes,
10106                                    ak[at],
10107                                    bk[at],
10108                                )));
10109                            }
10110                        }
10111                        if vb > 0 {
10112                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10113                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10114                            if av != bv {
10115                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10116                                return Err(fail(&format!(
10117                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10118                                    at / a.v_tok_bytes,
10119                                    at % a.v_tok_bytes,
10120                                    av[at],
10121                                    bv[at],
10122                                )));
10123                            }
10124                        }
10125                        report.trunk_kv_bytes += kb + vb;
10126                    }
10127                    (None, None) => {}
10128                    _ => return Err(fail(&format!("layer {il} KV presence"))),
10129                }
10130                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10131                    (Some(a), Some(b)) => {
10132                        let ac = es.dtoh(&a.conv_state)?;
10133                        let bc = es.dtoh(&b.conv_state)?;
10134                        if !same_f32(&ac, &bc) {
10135                            return Err(fail(&format!("layer {il} conv state")));
10136                        }
10137                        let as_ = es.dtoh(&a.ssm_state)?;
10138                        let bs = es.dtoh(&b.ssm_state)?;
10139                        if !same_f32(&as_, &bs) {
10140                            return Err(fail(&format!("layer {il} SSM state")));
10141                        }
10142                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10143                    }
10144                    (None, None) => {}
10145                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10146                }
10147            }
10148            Ok(())
10149        }
10150
10151        if reference.committed != candidate.committed {
10152            return Err(fail("committed token ids"));
10153        }
10154        if reference.cache.pos != candidate.cache.pos
10155            || reference.cache.max_ctx != candidate.cache.max_ctx
10156        {
10157            return Err(fail("cache pos/capacity"));
10158        }
10159        if reference.pending_tok != candidate.pending_tok
10160            || reference.next_pred != candidate.next_pred
10161            || reference.sctr != candidate.sctr
10162            || reference.uctr != candidate.uctr
10163        {
10164            return Err(fail("pending/prediction/counter tail"));
10165        }
10166
10167        let mut report = OptiForkStateIdentity::default();
10168        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10169            let rt = crate::pp::PpNRt::get(e)?;
10170            for stage in 0..rt.n_stages() {
10171                let _scope = rt.enter(stage);
10172                compare_layers(
10173                    rt.engine(stage, e),
10174                    fence[stage]..fence[stage + 1],
10175                    reference,
10176                    candidate,
10177                    &mut report,
10178                )?;
10179            }
10180        } else {
10181            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10182        }
10183
10184        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10185            return Err(fail("draft scratch plane count"));
10186        }
10187        for index in 0..reference.scratch.plane_count() {
10188            let (a, _) = reference.scratch.plane(index);
10189            let (b, _) = candidate.scratch.plane(index);
10190            if a.len != b.len
10191                || a.kv_dim_k != b.kv_dim_k
10192                || a.kv_dim_v != b.kv_dim_v
10193                || a.k_tok_bytes != b.k_tok_bytes
10194                || a.v_tok_bytes != b.v_tok_bytes
10195                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10196            {
10197                return Err(fail(&format!("draft scratch plane {index} length/layout")));
10198            }
10199            let kb = a.len * a.k_tok_bytes;
10200            let vb = a.len * a.v_tok_bytes;
10201            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10202                return Err(fail(&format!("draft scratch plane {index} K bytes")));
10203            }
10204            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10205                return Err(fail(&format!("draft scratch plane {index} V bytes")));
10206            }
10207            report.scratch_kv_bytes += kb + vb;
10208        }
10209
10210        match (&reference.last_h, &candidate.last_h) {
10211            (Some(a), Some(b)) => {
10212                let ah = e.dtoh(a)?;
10213                let bh = e.dtoh(b)?;
10214                if !same_f32(&ah, &bh) {
10215                    return Err(fail("last hidden/seed bytes"));
10216                }
10217                report.hidden_bytes = ah.len() * 4;
10218            }
10219            (None, None) => {}
10220            _ => return Err(fail("last hidden/seed presence")),
10221        }
10222        Ok(report)
10223    }
10224
10225    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10226    /// retained prompt-end checkpoint, so a request whose prompt matches
10227    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10228    ///
10229    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10230    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10231    /// restored from the device copy taken there, draft scratch length reset, `committed`
10232    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10233    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10234    /// every burst after it are identical to a cold run of the same token stream — the
10235    /// committed-tokens-authoritative contract.
10236    ///
10237    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10238    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10239    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10240    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10241    /// (the scratch KV, the resident embedding), none of which the rewind moves.
10242    ///
10243    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10244    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10245    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10246    pub fn spec_rewind_to_checkpoint(
10247        &self,
10248        e: &Engine,
10249        sess: &mut SpecSession,
10250    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10251        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10252            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10253        }) {
10254            return Err(
10255                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10256            );
10257        }
10258        let Some(ckpt) = sess.turn_ckpt.take() else {
10259            return Ok(None);
10260        };
10261        assert!(
10262            ckpt.pos <= sess.committed.len(),
10263            "checkpoint past committed ({} > {})",
10264            ckpt.pos,
10265            sess.committed.len()
10266        );
10267        // Restore through each layer's owning engine. A single primary-engine rollback is not
10268        // sufficient when the serving cache is stage-owned under cross-device PP.
10269        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
10270        debug_assert_eq!(
10271            sess.cache.pos, ckpt.pos,
10272            "rollback landed off the checkpoint"
10273        );
10274        sess.scratch.set_len(e, ckpt.pos)?;
10275        sess.committed.truncate(ckpt.pos);
10276        sess.last_h = Some(ckpt.last_h);
10277        sess.next_pred = None;
10278        sess.pending_tok = None;
10279        Ok(Some(ckpt.pos))
10280    }
10281
10282    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
10283    /// checkpoint without re-priming the checkpoint prefix.
10284    ///
10285    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
10286    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
10287    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
10288    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
10289    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
10290    ///
10291    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
10292    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
10293    pub fn spec_grow_and_rewind_to_checkpoint(
10294        &self,
10295        e: &Engine,
10296        sess: &mut SpecSession,
10297        target_cap: usize,
10298    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10299        if target_cap <= sess.cache.max_ctx {
10300            return self.spec_rewind_to_checkpoint(e, sess);
10301        }
10302        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
10303            return Ok(None);
10304        };
10305        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
10306            return Err(format!(
10307                "checkpoint pos {} outside committed length {}",
10308                ckpt.pos,
10309                sess.committed.len(),
10310            )
10311            .into());
10312        }
10313        if ckpt.pos > target_cap {
10314            return Err(format!(
10315                "checkpoint pos {} exceeds grown capacity {target_cap}",
10316                ckpt.pos,
10317            )
10318            .into());
10319        }
10320
10321        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
10322        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
10323        crate::pp::restore_cache_checkpoint(
10324            e,
10325            self,
10326            Some(&sess.cache),
10327            &mut grown_cache,
10328            &ckpt.snap,
10329        )?;
10330
10331        if sess.scratch.plane_count() != grown_scratch.plane_count() {
10332            return Err("checkpoint draft plane count mismatch".into());
10333        }
10334        for index in 0..sess.scratch.plane_count() {
10335            let (src, _) = sess.scratch.plane(index);
10336            let (dst, _) = grown_scratch.plane_mut(index);
10337            if ckpt.pos > src.len
10338                || src.kv_dim_k != dst.kv_dim_k
10339                || src.kv_dim_v != dst.kv_dim_v
10340                || src.k_tok_bytes != dst.k_tok_bytes
10341                || src.v_tok_bytes != dst.v_tok_bytes
10342            {
10343                return Err(format!(
10344                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
10345                    ckpt.pos, src.len,
10346                )
10347                .into());
10348            }
10349            match (&src.ring, dst.ring.as_ref()) {
10350                (Some(sring), Some(_)) => {
10351                    // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
10352                    // physical rows once lapped — same class as the trunk-KV restore panic
10353                    // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
10354                    let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
10355                        format!("checkpoint draft plane {index} SWA restore refused: {err}")
10356                    })?;
10357                    let rows = phys.len();
10358                    let kb = rows * src.k_tok_bytes;
10359                    let vb = rows * src.v_tok_bytes;
10360                    if kb > 0 {
10361                        e.copy_u8_range_into(
10362                            &mut dst.k,
10363                            0,
10364                            &src.k,
10365                            phys.start * src.k_tok_bytes,
10366                            kb,
10367                        )?;
10368                    }
10369                    if vb > 0 {
10370                        e.copy_u8_range_into(
10371                            &mut dst.v,
10372                            0,
10373                            &src.v,
10374                            phys.start * src.v_tok_bytes,
10375                            vb,
10376                        )?;
10377                    }
10378                    dst.ring
10379                        .as_mut()
10380                        .expect("ring presence checked above")
10381                        .apply_rebase(new_base);
10382                    if let Some(base_d) = dst.base_d.as_mut() {
10383                        e.set_i32_one(base_d, new_base as i32)?;
10384                    }
10385                }
10386                (None, None) => {
10387                    let kb = ckpt.pos * src.k_tok_bytes;
10388                    let vb = ckpt.pos * src.v_tok_bytes;
10389                    if kb > 0 {
10390                        e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
10391                    }
10392                    if vb > 0 {
10393                        e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
10394                    }
10395                }
10396                _ => {
10397                    return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
10398                }
10399            }
10400        }
10401        grown_scratch.set_len(e, ckpt.pos)?;
10402        // The old scratch is dropped immediately after publication below. Bound its D2D reads
10403        // first; growth happens once per rewritten turn, outside the decode hot loop.
10404        e.stream().synchronize()?;
10405
10406        let ckpt = sess
10407            .turn_ckpt
10408            .take()
10409            .expect("checkpoint remained present through transactional grow");
10410        let pos = ckpt.pos;
10411        sess.cache = grown_cache;
10412        sess.scratch = grown_scratch;
10413        sess.committed.truncate(pos);
10414        sess.last_h = Some(ckpt.last_h);
10415        sess.next_pred = None;
10416        sess.pending_tok = None;
10417        sess.draft_ctx = None;
10418        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
10419        debug_assert!(
10420            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
10421            "grown draft rewind landed off checkpoint"
10422        );
10423        Ok(Some(pos))
10424    }
10425
10426    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
10427    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
10428    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
10429    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
10430    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
10431    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
10432    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
10433    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
10434    /// park-time flush is a future request whose sampler is not knowable here (residual
10435    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
10436    pub fn spec_flush_pending(
10437        &self,
10438        e: &Engine,
10439        sess: &mut SpecSession,
10440        sampling: Option<SpecSampling>,
10441    ) -> Result<(), Box<dyn std::error::Error>> {
10442        let Some(b) = sess.pending_tok.take() else {
10443            return Ok(());
10444        };
10445        if self.mtp.is_none() {
10446            return Err("pending carry requires an MTP head".into());
10447        }
10448        let n_embd = self.cfg.n_embd as usize;
10449        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10450        let embd_gpu = if spec_host_embd() {
10451            None
10452        } else {
10453            Some(
10454                self.embd_gpu
10455                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10456            )
10457        };
10458        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10459        let pos_b = sess.cache.pos;
10460        sess.scratch.set_len(e, pos_b)?;
10461        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
10462        sess.next_pred = Some(match sampling {
10463            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
10464                // window includes `b` itself: it is committed by this pass, and the pre-lane
10465                // code never counted a boundary token in the penalty history at all.
10466                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
10467                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
10468            }
10469            _ => argmax(&lg_b) as u32,
10470        });
10471        let anchor = sess
10472            .last_h
10473            .as_ref()
10474            .expect("pending carry requires last_h (the predecessor-row anchor)");
10475        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
10476        sess.last_h = Some(hb);
10477        sess.committed.push(b);
10478        Ok(())
10479    }
10480
10481    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
10482    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
10483    /// rounds through that same graph. Other model families keep their eager T=1 contract.
10484    fn spec_target_step_h(
10485        &self,
10486        e: &Engine,
10487        token: u32,
10488        cache: &mut Cache,
10489    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10490        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
10491            return self.decode_step_h(e, token, cache);
10492        }
10493        let pos0 = cache.pos;
10494        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
10495        Ok((e.dtoh(&logits)?, hidden))
10496    }
10497
10498    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
10499    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
10500    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
10501    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
10502    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
10503    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
10504    /// dispatch sites cannot drift apart again.
10505    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
10506    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
10507    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
10508    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
10509    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
10510    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
10511    fn mtp_graph_capturable(&self) -> bool {
10512        self.mtp
10513            .as_ref()
10514            .map(|m| match &m.ffn {
10515                crate::hybrid::Ffn::Dense { .. } => true,
10516                crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
10517            })
10518            .unwrap_or(false)
10519    }
10520
10521    fn batched_serving_numeric_class(&self) -> bool {
10522        self.plan
10523            .trunk_operations()
10524            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
10525    }
10526
10527    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
10528    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
10529    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
10530    /// keeping the engine's own version structural rather than name-based means a new
10531    /// checkpoint of the same shape inherits the default, and a different shape does not.
10532    fn vgraph_family_default(&self) -> bool {
10533        let has_linear = self
10534            .layers
10535            .iter()
10536            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
10537        let has_moe = self
10538            .layers
10539            .iter()
10540            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
10541        has_linear && has_moe
10542    }
10543
10544    fn sliding_gated_moe_batch_program(&self) -> bool {
10545        self.uses_sliding_gated_moe_program()
10546    }
10547
10548    fn gemma_batch_program(&self) -> bool {
10549        self.uses_gemma_program()
10550    }
10551
10552    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
10553    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
10554    /// session already exist.
10555    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
10556        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
10557            || !spec_devacc()
10558            || spec_replay_env_enabled()
10559            || spec_stream()
10560            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
10561            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
10562            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
10563            || std::env::var("MEMRA_SPEC_PMIN")
10564                .ok()
10565                .and_then(|v| v.parse::<f32>().ok())
10566                .unwrap_or(0.0)
10567                > 0.0
10568            || self.is_gemma4_e4b()
10569            || self.gemma_batch_program()
10570            || self.mtp.is_none()
10571            || !self.mtp_extra.is_empty()
10572        {
10573            return false;
10574        }
10575        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
10576            return false;
10577        };
10578        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
10579            return false;
10580        }
10581        crate::pp::PpNRt::get(e)
10582            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
10583            .unwrap_or(false)
10584    }
10585
10586    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
10587    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
10588    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
10589    #[allow(clippy::too_many_arguments)]
10590    pub fn generate_spec_session_pair(
10591        &self,
10592        e: &Engine,
10593        sess_a: &mut SpecSession,
10594        max_new_a: usize,
10595        k_a: usize,
10596        sess_b: &mut SpecSession,
10597        max_new_b: usize,
10598        k_b: usize,
10599    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
10600    {
10601        if !self.spec_pipe_available(e) {
10602            return Err("two-session speculative pipeline is outside its reduced matrix".into());
10603        }
10604        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
10605            return Err(
10606                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
10607            );
10608        }
10609        for sess in [&*sess_a, &*sess_b] {
10610            if sess.committed.is_empty()
10611                || sess.last_h.is_none()
10612                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
10613            {
10614                return Err("two-session speculative pipeline requires warm continuations".into());
10615            }
10616        }
10617
10618        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10619            && !spec_host_embd()
10620            && self.mtp_graph_capturable()
10621            && self.mtp_extra.is_empty()
10622            && !crate::model::full_prec_enabled();
10623        let graph_a = graph_ok && k_a + 2 < 96;
10624        let graph_b = graph_ok && k_b + 2 < 96;
10625        let was_tracking = e.ctx().is_event_tracking();
10626        if (graph_a || graph_b) && was_tracking {
10627            unsafe {
10628                e.ctx().disable_event_tracking();
10629            }
10630        }
10631
10632        static LOGGED: std::sync::Once = std::sync::Once::new();
10633        LOGGED.call_once(|| {
10634            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
10635        });
10636        let sync = std::sync::Arc::new(SpecPipeSync::new());
10637        let lane_a = SpecPipeLane {
10638            sync: sync.clone(),
10639            lane: 0,
10640        };
10641        let lane_b = SpecPipeLane { sync, lane: 1 };
10642        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
10643        let (result_a, result_b) = std::thread::scope(|scope| {
10644            let b = scope.spawn(move || {
10645                let mut finish = SpecPipeFinish::new(&lane_b);
10646                let sess_b = unsafe { sess_b_ptr.get_mut() };
10647                let result = e
10648                    .ctx()
10649                    .bind_to_thread()
10650                    .map_err(|err| err.to_string())
10651                    .and_then(|_| {
10652                        self.generate_spec_inner2(
10653                            e,
10654                            &[],
10655                            max_new_b,
10656                            k_b,
10657                            graph_b,
10658                            Some(sess_b),
10659                            None,
10660                            None,
10661                            None,
10662                            None,
10663                            Some(&lane_b),
10664                        )
10665                        .map_err(|err| err.to_string())
10666                    });
10667                finish.close(result.is_err());
10668                result
10669            });
10670            let mut finish = SpecPipeFinish::new(&lane_a);
10671            let result_a = self.generate_spec_inner2(
10672                e,
10673                &[],
10674                max_new_a,
10675                k_a,
10676                graph_a,
10677                Some(sess_a),
10678                None,
10679                None,
10680                None,
10681                None,
10682                Some(&lane_a),
10683            );
10684            finish.close(result_a.is_err());
10685            let result_b = b
10686                .join()
10687                .map_err(|_| "paired speculative session B panicked".to_string())
10688                .and_then(|r| r);
10689            (result_a, result_b)
10690        });
10691
10692        if (graph_a || graph_b) && was_tracking {
10693            unsafe {
10694                e.ctx().enable_event_tracking();
10695            }
10696        }
10697        let result_a = result_a?;
10698        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
10699        Ok((result_a, result_b))
10700    }
10701
10702    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
10703    /// message rendered through the chat template continuation). Returns (new tokens emitted,
10704    /// drafted, accepted); session.committed grows by suffix + emitted.
10705    pub fn generate_spec_session(
10706        &self,
10707        e: &Engine,
10708        sess: &mut SpecSession,
10709        suffix: &[u32],
10710        max_new: usize,
10711        k: usize,
10712    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10713        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
10714    }
10715
10716    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
10717    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
10718    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
10719    /// for the filtered target (feat/filtered-spec).
10720    ///
10721    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
10722    /// output — once right after the prime's first token, then once per round commit — so a
10723    /// streaming caller can flush text at round cadence instead of once per burst. The slices
10724    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
10725    /// timing only: token bytes, session state, and exactness are untouched.
10726    ///
10727    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
10728    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
10729    /// the caller's scheduler regains control without waiting the burst out. Burst size is
10730    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
10731    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
10732    /// drains and the defensive tail flush can land with nothing new committed).
10733    #[allow(clippy::too_many_arguments)]
10734    pub fn generate_spec_session_sampled(
10735        &self,
10736        e: &Engine,
10737        sess: &mut SpecSession,
10738        suffix: &[u32],
10739        max_new: usize,
10740        k: usize,
10741        sampling: Option<SpecSampling>,
10742        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10743    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10744        self.generate_spec_session_sampled_prime_split(
10745            e, sess, suffix, max_new, k, sampling, None, on_commit,
10746        )
10747    }
10748
10749    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
10750    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
10751    /// pass `None` and stay on the existing zero-prime path.
10752    #[allow(clippy::too_many_arguments)]
10753    pub fn generate_spec_session_sampled_prime_split(
10754        &self,
10755        e: &Engine,
10756        sess: &mut SpecSession,
10757        suffix: &[u32],
10758        max_new: usize,
10759        k: usize,
10760        sampling: Option<SpecSampling>,
10761        prime_split: Option<usize>,
10762        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10763    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10764        self.generate_spec_session_constrained_prime_split(
10765            e,
10766            sess,
10767            suffix,
10768            max_new,
10769            k,
10770            sampling,
10771            None,
10772            prime_split,
10773            on_commit,
10774        )
10775    }
10776
10777    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
10778    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
10779    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
10780    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
10781    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
10782    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
10783    /// may drop (drafter is unconstrained); that is measured, not hidden.
10784    #[allow(clippy::too_many_arguments)]
10785    pub fn generate_spec_session_constrained(
10786        &self,
10787        e: &Engine,
10788        sess: &mut SpecSession,
10789        suffix: &[u32],
10790        max_new: usize,
10791        k: usize,
10792        sampling: Option<SpecSampling>,
10793        constraint: Option<&mut dyn SpecConstraint>,
10794        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10795    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10796        self.generate_spec_session_constrained_prime_split(
10797            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
10798        )
10799    }
10800
10801    #[allow(clippy::too_many_arguments)]
10802    pub fn generate_spec_session_constrained_prime_split(
10803        &self,
10804        e: &Engine,
10805        sess: &mut SpecSession,
10806        suffix: &[u32],
10807        max_new: usize,
10808        k: usize,
10809        sampling: Option<SpecSampling>,
10810        constraint: Option<&mut dyn SpecConstraint>,
10811        prime_split: Option<usize>,
10812        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10813    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10814        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
10815            return Err(
10816                "constrained spec decode is greedy-only (worker routes sampled \
10817                        constrained to plain decode)"
10818                    .into(),
10819            );
10820        }
10821        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
10822        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
10823        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
10824        // serve continuation case — consume the carry in-loop with zero solo passes.
10825        if sess.pending_tok.is_some()
10826            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
10827        {
10828            self.spec_flush_pending(e, sess, sampling)?;
10829        }
10830
10831        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
10832        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
10833        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
10834        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10835            && !spec_host_embd()
10836            && self.mtp_graph_capturable()
10837            && self.mtp_extra.is_empty()
10838            && k + 2 < 96
10839            && !crate::model::full_prec_enabled();
10840        let was_tracking = e.ctx().is_event_tracking();
10841        if graph_draft && was_tracking {
10842            unsafe {
10843                e.ctx().disable_event_tracking();
10844            }
10845        }
10846        let r = self.generate_spec_inner2(
10847            e,
10848            suffix,
10849            max_new,
10850            k,
10851            graph_draft,
10852            Some(sess),
10853            sampling,
10854            constraint,
10855            on_commit,
10856            prime_split,
10857            None,
10858        );
10859        if graph_draft && was_tracking {
10860            unsafe {
10861                e.ctx().enable_event_tracking();
10862            }
10863        }
10864        let (out, d, a) = r?;
10865        Ok((out, d, a))
10866    }
10867
10868    pub fn generate_spec(
10869        &self,
10870        e: &Engine,
10871        prompt: &[u32],
10872        max_new: usize,
10873        k: usize,
10874    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10875        if crate::pp::pp_cuts(self.layers.len()).is_some()
10876            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
10877        {
10878            return Err("pipeline rewrite is not qualified for speculative decode".into());
10879        }
10880        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
10881            return Err("speculative rewrite is not qualified for this ModelPlan".into());
10882        }
10883        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
10884        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
10885        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10886            && !spec_host_embd()
10887            && self.mtp_graph_capturable()
10888            && self.mtp_extra.is_empty()
10889            && k + 2 < 96
10890            && !crate::model::full_prec_enabled();
10891        if !graph_draft {
10892            return self.generate_spec_inner2(
10893                e, prompt, max_new, k, false, None, None, None, None, None, None,
10894            );
10895        }
10896        let was_tracking = e.ctx().is_event_tracking();
10897        if was_tracking {
10898            unsafe {
10899                e.ctx().disable_event_tracking();
10900            }
10901        }
10902        let r = self.generate_spec_inner2(
10903            e, prompt, max_new, k, true, None, None, None, None, None, None,
10904        );
10905        if was_tracking {
10906            unsafe {
10907                e.ctx().enable_event_tracking();
10908            }
10909        }
10910        r
10911    }
10912
10913    fn generate_spec_inner2(
10914        &self,
10915        e: &Engine,
10916        prompt: &[u32],
10917        max_new: usize,
10918        k: usize,
10919        graph_draft: bool,
10920        mut sess: Option<&mut SpecSession>,
10921        sampling: Option<SpecSampling>,
10922        mut constraint: Option<&mut dyn SpecConstraint>,
10923        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10924        prime_split: Option<usize>,
10925        pipe: Option<&SpecPipeLane>,
10926    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10927        assert!(k >= 1, "k must be >= 1");
10928        if let Some(p) = pipe {
10929            p.setup_begin()?;
10930        }
10931        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
10932        let mut flushed = 0usize;
10933        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
10934        // at the next round boundary (same exit as max_new reached — the session tail runs).
10935        // Initialized by the unconditional post-prime flush below.
10936        let mut keep_going;
10937        let mtp = self
10938            .mtp
10939            .as_ref()
10940            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
10941        let n_vocab = self.output.out_features();
10942        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
10943        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
10944        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
10945        let d_vocab = mtp
10946            .shared_head_head
10947            .as_ref()
10948            .unwrap_or(&self.output)
10949            .out_features();
10950        if !self.mtp_extra.is_empty() {
10951            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
10952                || self.plan.mtp_blocks.len() != self.mtp_head_count()
10953            {
10954                return Err(
10955                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
10956                );
10957            }
10958            // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
10959            // token-frequency and head-independent, and every downstream remap (per-step argmax,
10960            // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
10961            // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
10962            for (offset, head) in self.mtp_extra.iter().enumerate() {
10963                if head.d2t != mtp.d2t
10964                    || head
10965                        .shared_head_head
10966                        .as_ref()
10967                        .unwrap_or(&self.output)
10968                        .out_features()
10969                        != d_vocab
10970                {
10971                    return Err(format!(
10972                        "embedded MTP head {} has incompatible draft vocabulary",
10973                        offset + 1
10974                    )
10975                    .into());
10976                }
10977            }
10978            eprintln!(
10979                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
10980                self.mtp_head_count()
10981            );
10982        }
10983        let n_embd = self.cfg.n_embd as usize;
10984        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
10985        // already committed (their state is in the caches); 0 = fresh single-shot call.
10986        let session_mode = sess.is_some();
10987        let max_ctx = match sess.as_ref() {
10988            Some(s) => s.cache.max_ctx,
10989            None => prompt.len() + max_new + k + 8,
10990        };
10991        let mut own_cache;
10992        let mut own_scratch;
10993        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
10994        // (requested split, destination list). Single-shot per burst; fresh calls have none.
10995        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
10996        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
10997        // committed-length position; consumed one-shot like `capture_at`. None = legacy
10998        // prompt-end capture below.
10999        let mut ckpt_req: Option<usize> = None;
11000        let (
11001            cache,
11002            scratch,
11003            mut sess_tail,
11004            mut sess_draft_slot,
11005            mut sess_pending_slot,
11006            sess_ckpt_slot,
11007            sess_telem,
11008        ): (
11009            &mut Cache,
11010            &mut MtpScratch,
11011            Option<(
11012                &mut Vec<u32>,
11013                &mut Option<CudaSlice<f32>>,
11014                &mut Option<u32>,
11015                &mut u32,
11016                &mut u32,
11017            )>,
11018            Option<&mut Option<DraftGraphCtx>>,
11019            Option<&mut Option<u32>>,
11020            Option<&mut Option<SpecCheckpoint>>,
11021            Option<&SpecTelemetryCounters>,
11022        ) = match sess.take() {
11023            Some(sr) => {
11024                let SpecSession {
11025                    cache,
11026                    scratch,
11027                    committed,
11028                    last_h,
11029                    next_pred,
11030                    sctr: s_sctr,
11031                    uctr: s_uctr,
11032                    draft_ctx,
11033                    pending_tok,
11034                    turn_ckpt,
11035                    telem,
11036                    capture_at,
11037                    boundary_captures,
11038                    ckpt_at,
11039                } = sr;
11040                sess_capture = Some((capture_at.take(), boundary_captures));
11041                ckpt_req = ckpt_at.take();
11042                (
11043                    cache,
11044                    scratch,
11045                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11046                    Some(draft_ctx),
11047                    Some(pending_tok),
11048                    Some(turn_ckpt),
11049                    Some(telem),
11050                )
11051            }
11052            None => {
11053                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11054                // `Cache::new` verbatim.
11055                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11056                // Persistent scratch = max_ctx rows (~2KB/token quantized).
11057                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11058                (
11059                    &mut own_cache,
11060                    &mut own_scratch,
11061                    None,
11062                    None,
11063                    None,
11064                    None,
11065                    None,
11066                )
11067            }
11068        };
11069        if scratch.plane_count() != self.mtp_head_count() {
11070            return Err(format!(
11071                "MTP scratch/head count mismatch ({}/{})",
11072                scratch.plane_count(),
11073                self.mtp_head_count()
11074            )
11075            .into());
11076        }
11077        let base = cache.pos;
11078        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11079        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11080        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11081        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11082        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11083        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11084        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11085        // acceptance-only — exactness is verify's job either way).
11086        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11087        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11088        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11089        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11090        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11091        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11092        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11093        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11094        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11095        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11096        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11097        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11098        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11099        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11100        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11101        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11102        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11103        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11104        // + fallback seam).
11105        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11106        // bar — the retained verify-state commit proven equivalent to sequential serving —
11107        // was waiting on this arch running the serving batched verify class, which the
11108        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11109        // replay-free commit consumes is now produced by the SAME serving-class verify that
11110        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11111        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11112        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11113        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11114        // rollback + A/B seam.
11115        let spec_replay = spec_replay_env_enabled();
11116        if constraint.is_some() && spec_replay {
11117            return Err(
11118                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11119                        (legacy replay commits an unmasked bonus)"
11120                    .into(),
11121            );
11122        }
11123        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11124        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11125        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11126        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11127        if !refresh && !self.mtp_extra.is_empty() {
11128            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11129        }
11130
11131        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11132        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11133        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11134        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11135        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11136        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11137        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11138        // generation exactly where the last turn stopped — no prime at all. The stashed
11139        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11140        // committed.last() by the same rule this entry applies to a cold prime's last row —
11141        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11142        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11143        // where the sampler and the session's Philox counters were live). `last_h` seeds the
11144        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11145        let continuation = prompt.is_empty();
11146        if continuation {
11147            assert!(session_mode, "empty prompt requires a session");
11148            assert!(
11149                sess_tail
11150                    .as_ref()
11151                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
11152                        && lh.is_some()
11153                        && (np.is_some() || carried_pending.is_some())),
11154                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11155            );
11156        }
11157        let mut prime_logits;
11158        let mut prompt_h: Option<CudaSlice<f32>> = None;
11159        let t_prime = std::time::Instant::now();
11160        let batched_prime = !continuation
11161            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11162            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11163            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11164        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11165        if prime_split.is_some() && continuation {
11166            return Err("spec prime split requires a non-empty prime".into());
11167        }
11168        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11169        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11170        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11171        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11172        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11173        // cannot honor (outside this prime's range) silently drops the capture — the
11174        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11175        let ckpt_rel = if continuation {
11176            None
11177        } else {
11178            ckpt_req
11179                .and_then(|abs| abs.checked_sub(base))
11180                .filter(|&r| r > 0 && r < prompt.len())
11181        };
11182        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11183        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11184        // the legacy single-split program, byte-for-byte.
11185        let mut stops: Vec<usize> = Vec::new();
11186        for b in [prime_split, ckpt_rel].into_iter().flatten() {
11187            if !stops.contains(&b) {
11188                stops.push(b);
11189            }
11190        }
11191        stops.sort_unstable();
11192        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11193        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11194        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11195        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11196        if continuation {
11197            prime_logits = Vec::new();
11198        } else if !stops.is_empty() {
11199            if let Some(&first) = stops.first() {
11200                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
11201                    return Err(format!(
11202                        "spec prime split {first} is below PRIME_MIN_T {}",
11203                        crate::hybrid_forward::PRIME_MIN_T,
11204                    )
11205                    .into());
11206                }
11207            }
11208            // Mirror the plain worker's boundary stops exactly. Each segment is a
11209            // request-level prime (`queued_after` keeps Step35 arm selection independent of
11210            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
11211            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
11212            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
11213            // coherent prompt.
11214            let mut h_all = e.uninit(prompt.len() * n_embd)?;
11215            prime_logits = Vec::new();
11216            let mut prev = 0usize;
11217            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
11218                if seg_end <= prev {
11219                    continue;
11220                }
11221                let seg = &prompt[prev..seg_end];
11222                let is_final = seg_end == prompt.len();
11223                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
11224                    && (!is_final
11225                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11226                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
11227                if batched_seg {
11228                    let (l, _, h_seg) =
11229                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
11230                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
11231                    prime_logits = l;
11232                } else {
11233                    for (i, &tok) in seg.iter().enumerate() {
11234                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
11235                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
11236                        prime_logits = l;
11237                    }
11238                }
11239                prev = seg_end;
11240                if is_final {
11241                    break;
11242                }
11243                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
11244                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
11245                // states are about to be advanced in place by the next segment, so this is
11246                // the ONLY moment the boundary's recurrent state exists. Capture iff the
11247                // worker requested exactly this stop (cold sessions only — `capture_at` is
11248                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
11249                // publication is an optimization, never a correctness dependency.
11250                if base == 0 {
11251                    if let Some((requested, slot)) = sess_capture.as_mut() {
11252                        // Publish at the requested miss-LCP stop (the shared-prefix class)
11253                        // AND at the stable-boundary stop (the next-turn re-render class,
11254                        // lane/frspec-multiturn-cache) — the same boundary set the plain
11255                        // prefill tick learns. Without the second entry, the turn after a
11256                        // cold re-park could only hit the OLDER lcp entry (the measured
11257                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
11258                        // rewound to 15222). Dedupe is the worker sweep's has_key.
11259                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
11260                            if let Ok(snap) = cache.snapshot(e) {
11261                                slot.push(SpecBoundaryCapture {
11262                                    snap,
11263                                    pos: seg_end,
11264                                    logits: prime_logits.clone(),
11265                                    // rows [0..seg_end) of h_all are primed — the following
11266                                    // segments append, never overwrite.
11267                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
11268                                });
11269                            }
11270                        }
11271                    }
11272                }
11273                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
11274                // same snapshot mechanics, installed post-prime in place of the prompt-end
11275                // capture the re-render class always diverged below.
11276                if ckpt_rel == Some(seg_end) {
11277                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11278                        e.uninit(n_embd).and_then(|mut a| {
11279                            e.copy_view_into(
11280                                &mut a,
11281                                0,
11282                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
11283                                n_embd,
11284                            )?;
11285                            Ok(a)
11286                        });
11287                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
11288                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
11289                            snap,
11290                            pos: base + seg_end,
11291                            last_h,
11292                        }),
11293                        _ => None,
11294                    });
11295                }
11296            }
11297            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
11298                eprintln!(
11299                    "[spec-prime] stops={stops:?} tail={}",
11300                    prompt.len() - stops.last().copied().unwrap_or(0)
11301                );
11302            }
11303            prompt_h = Some(h_all);
11304        } else if batched_prime {
11305            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
11306            prime_logits = l;
11307            prompt_h = Some(hiddens);
11308        } else {
11309            prime_logits = Vec::new();
11310            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
11311            for (i, &tok) in prompt.iter().enumerate() {
11312                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
11313                if let Some(ph) = prompt_h.as_mut() {
11314                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
11315                }
11316                prime_logits = l;
11317            }
11318        }
11319        e.stream().synchronize()?;
11320        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
11321        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
11322        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
11323        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
11324        // prime_split. The mid-prompt capture above already consumed the request if it matched.
11325        if !continuation && base == 0 {
11326            if let Some((requested, slot)) = sess_capture.as_mut() {
11327                if *requested == Some(prompt.len()) && slot.is_empty() {
11328                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
11329                    if let Ok(snap) = cache.snapshot(e) {
11330                        slot.push(SpecBoundaryCapture {
11331                            snap,
11332                            pos: prompt.len(),
11333                            logits: prime_logits.clone(),
11334                            last_h: prompt_h
11335                                .as_ref()
11336                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
11337                                .unwrap_or_default(),
11338                        });
11339                    }
11340                }
11341            }
11342        }
11343        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
11344        // prime-subtraction hack.
11345        crate::PRIME_NANOS.store(
11346            t_prime.elapsed().as_nanos() as u64,
11347            std::sync::atomic::Ordering::Relaxed,
11348        );
11349
11350        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11351        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
11352        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
11353        let host_embd = spec_host_embd();
11354        let embd_gpu = if host_embd {
11355            None
11356        } else {
11357            Some(
11358                self.embd_gpu
11359                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11360            )
11361        };
11362        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11363        if host_embd {
11364            eprintln!(
11365                "[spec] host-row embedding: {} bytes kept off HBM",
11366                self.embd.raw.len()
11367            );
11368        }
11369        let mut out: Vec<u32> = Vec::with_capacity(max_new);
11370        let mut total_drafted = 0usize;
11371        let mut total_accepted = 0usize;
11372
11373        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
11374        // The sampler config, the session's Philox counters and the penalty window are parsed
11375        // HERE, above the boundary-token selection, because the boundary token must be drawn
11376        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
11377        // selection, which is the whole mechanical reason the boundary token was an argmax:
11378        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
11379        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
11380        // below takes the argmax path it always took).
11381        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
11382        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
11383        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
11384        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
11385        let sp = sampling.unwrap_or_else(|| SpecSampling {
11386            temp: std::env::var("MEMRA_SPEC_TEMP")
11387                .ok()
11388                .and_then(|v| v.parse().ok())
11389                .unwrap_or(0.0),
11390            seed: std::env::var("MEMRA_SEED")
11391                .ok()
11392                .and_then(|v| v.parse().ok())
11393                .unwrap_or(42),
11394            top_k: std::env::var("MEMRA_TOP_K")
11395                .ok()
11396                .and_then(|v| v.parse().ok())
11397                .unwrap_or(0),
11398            top_p: std::env::var("MEMRA_TOP_P")
11399                .ok()
11400                .and_then(|v| v.parse().ok())
11401                .unwrap_or(1.0),
11402            min_p: std::env::var("MEMRA_MIN_P")
11403                .ok()
11404                .and_then(|v| v.parse().ok())
11405                .unwrap_or(0.0),
11406            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
11407                .ok()
11408                .and_then(|v| v.parse().ok())
11409                .unwrap_or(0),
11410            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
11411                .ok()
11412                .and_then(|v| v.parse().ok())
11413                .unwrap_or(1.0),
11414            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
11415                .ok()
11416                .and_then(|v| v.parse().ok())
11417                .unwrap_or(0.0),
11418            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
11419                .ok()
11420                .and_then(|v| v.parse().ok())
11421                .unwrap_or(0.0),
11422        });
11423        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
11424        let sampled = sp_temp > 0.0;
11425        // Counters resume from the session (burst continuity: randomness must never repeat
11426        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
11427        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
11428        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
11429        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
11430        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
11431        // for the penalized+filtered target). History = generated tokens, host-tracked window.
11432        let pen_on = sampled
11433            && sp.penalty_last_n > 0
11434            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
11435        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
11436        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
11437        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
11438        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
11439        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
11440        // which is what the API contract says and what the plain sampler's own `history` does.
11441        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
11442        let mut pen_hist: Vec<u32> = if pen_on {
11443            let sess_hist: &[u32] = if spec_pen_session_on() {
11444                sess_tail
11445                    .as_ref()
11446                    .map(|(c, ..)| c.as_slice())
11447                    .unwrap_or(&[])
11448            } else {
11449                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
11450            };
11451            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
11452        } else {
11453            Vec::new()
11454        };
11455        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
11456        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
11457        // request's own filtered/penalized target through the session's Philox stream
11458        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
11459        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
11460        // Emit it, then FEED it to establish the loop invariant below.
11461        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
11462        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
11463        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
11464        // prompt's last logits (plain constrained-greedy identity); a continuation without
11465        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
11466        // worker never resumes constrained sessions from the pool, so this cannot fire).
11467        if let Some(c) = constraint.as_deref_mut() {
11468            if continuation && carried_pending.is_none() {
11469                return Err("constrained spec continuation requires a carried pending \
11470                            (pool resume is unconstrained-only)"
11471                    .into());
11472            }
11473            if !continuation {
11474                c.mask_logits(&mut prime_logits)
11475                    .map_err(|e2| format!("constraint: {e2}"))?;
11476            }
11477        }
11478        let mut last_token = if let Some(b) = carried_pending {
11479            b
11480        } else if continuation {
11481            // A continuation's boundary token was DRAWN by the burst that stashed it (the
11482            // session tail below), or by `spec_session_from_restored` for a converted
11483            // prefix-cache hit — in both cases from the correct logits row with this same
11484            // session's Philox stream, which is why it can be consumed here as-is.
11485            sess_tail.as_ref().unwrap().2.unwrap()
11486        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
11487            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
11488        } else {
11489            // greedy (byte contract), the rollback door, or constrained (masked-argmax
11490            // identity — the worker routes sampled+constrained to the plain path, and this
11491            // function refuses the combination outright above).
11492            argmax(&prime_logits) as u32
11493        };
11494        if pen_on {
11495            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
11496            // emitted token into its penalty history, and pre-lane the burst's first token
11497            // was invisible to penalties forever (never pushed, and never in `committed`
11498            // until this burst's tail). Covers the carry/continuation seeds too — neither is
11499            // in `committed` yet.
11500            pen_hist.push(last_token);
11501        }
11502        if carried_pending.is_none() {
11503            out.push(last_token);
11504            // grammar advances with every emitted token (carried pendings were consumed
11505            // by the burst that emitted them).
11506            if let Some(c) = constraint.as_deref_mut() {
11507                c.consume(last_token)
11508                    .map_err(|e2| format!("constraint: {e2}"))?;
11509            }
11510        }
11511        if continuation {
11512            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
11513            // overhang so the chain's first append lands at slot base (== committed.len()).
11514            scratch.set_len(e, base)?;
11515        }
11516        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
11517        // concatenating to the full `out`). Called after the prime's first token and after each
11518        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
11519        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
11520        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
11521        fn flush_commit(
11522            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
11523            out: &[u32],
11524            flushed: &mut usize,
11525        ) -> bool {
11526            if let Some(f) = cb.as_mut() {
11527                let keep = f(&out[*flushed..]);
11528                *flushed = out.len();
11529                keep
11530            } else {
11531                true
11532            }
11533        }
11534        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11535        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
11536        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
11537        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
11538        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
11539        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
11540        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
11541        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
11542        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
11543        // those, so their residual mass is p(x), correct by construction).
11544        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
11545            match &mtp.d2t {
11546                Some(map) => Some(e.htod_u32_v(map)?),
11547                None => None,
11548            }
11549        } else {
11550            None
11551        };
11552        let mut q_full_buf: Option<CudaSlice<f32>> = None;
11553        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
11554        // dspark sampled-admission walk); byte-identical to the closure it replaces.
11555        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
11556        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
11557        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
11558        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
11559        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
11560        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
11561        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
11562        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
11563        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
11564        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
11565        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
11566        let t_ent = std::time::Instant::now();
11567
11568        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
11569        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
11570        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
11571        // the one that matters (a history-rewriting client mutates what the session GENERATED,
11572        // so the next turn's prompt agrees with this one up to exactly here).
11573        //
11574        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
11575        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
11576        // hold exactly `base + prompt.len()` rows and nothing generated.
11577        //
11578        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
11579        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
11580        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
11581        // `<think>` block the client strips, so every later turn's diff diverged exactly one
11582        // token below the checkpoint and affinity declined 100% of the time. Measured on the
11583        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
11584        // whole mechanism inert while looking, from the outside, like a working
11585        // correctness-declines-safely path — hence the decline log carries the offsets.
11586        //
11587        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
11588        // state (the reason a spec session could not rewind before). The draft scratch needs no
11589        // copy: rows below the boundary are rewritten by the next turn's own fill.
11590        //
11591        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
11592        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
11593        // checkpoint rather than replacing it with a strictly worse one.
11594        //
11595        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
11596        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
11597        // fail the burst that is already running — so the error is swallowed, loud only under
11598        // MEMRA_DEBUG_SPEC.
11599        //
11600        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
11601        // posture above was DISPROVED for the think-posture template class — the prompt's own
11602        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
11603        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
11604        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
11605        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
11606        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
11607        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
11608        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
11609        if let Some(slot) = sess_ckpt_slot {
11610            if let Some(early) = ckpt_early {
11611                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11612                    eprintln!(
11613                        "[spec] stable-boundary turn checkpoint skipped; \
11614                               next turn re-primes in full"
11615                    );
11616                }
11617                *slot = early;
11618            } else if !continuation {
11619                let pos = cache.pos;
11620                debug_assert_eq!(
11621                    pos,
11622                    base + prompt.len(),
11623                    "turn checkpoint must sit at the prompt end, before the init feed"
11624                );
11625                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11626                    if let Some(ph) = &prompt_h {
11627                        // hidden of the LAST primed row = the predecessor anchor at this
11628                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
11629                        // last_h, and what the next prime's fill reads for its first row).
11630                        let np = prompt.len();
11631                        e.uninit(n_embd).and_then(|mut a| {
11632                            e.copy_view_into(
11633                                &mut a,
11634                                0,
11635                                &ph.slice((np - 1) * n_embd..np * n_embd),
11636                                n_embd,
11637                            )?;
11638                            Ok(a)
11639                        })
11640                    } else {
11641                        Err("no prompt hiddens".into())
11642                    };
11643                match (cache.snapshot(e), anchor) {
11644                    (Ok(snap), Ok(last_h)) => {
11645                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
11646                    }
11647                    (s, a) => {
11648                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
11649                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11650                            let err = s
11651                                .err()
11652                                .map(|e| e.to_string())
11653                                .or_else(|| a.err().map(|e| e.to_string()))
11654                                .unwrap_or_default();
11655                            eprintln!(
11656                                "[spec] turn checkpoint skipped ({err}); \
11657                                       next turn re-primes in full"
11658                            );
11659                        }
11660                    }
11661                }
11662            }
11663        }
11664        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
11665        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
11666        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
11667        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
11668        let mut last_pred = 0u32;
11669        let mut last_col_logits: Option<CudaSlice<f32>> = None;
11670        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
11671        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
11672        let mut init_logits_host: Option<Vec<f32>> = None;
11673        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
11674            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
11675            last_pred = argmax(&init_logits) as u32;
11676            if constraint.is_some() {
11677                init_logits_host = Some(init_logits.clone());
11678            }
11679            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
11680            if sampled {
11681                last_col_logits = Some(e.htod(&init_logits)?);
11682            }
11683            h
11684        } else {
11685            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
11686            let lh = sess_tail
11687                .as_ref()
11688                .unwrap()
11689                .1
11690                .as_ref()
11691                .expect("pending carry requires last_h");
11692            e.clone_dtod(lh)?
11693        };
11694        let t_init = t_ent.elapsed();
11695        let mut last_col_stats: Option<(f32, f32, f32)> = None;
11696        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
11697        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
11698        // stable pointer for the graph-draft round-start copy.
11699        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
11700        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
11701        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
11702        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
11703        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
11704        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
11705        // overwritten below).
11706        let mut fill_prev = e.clone_dtod(&h_seed0)?;
11707        {
11708            if let Some(ph) = &prompt_h {
11709                let np = prompt.len();
11710                e.copy_view_into(
11711                    &mut h_seed_buf,
11712                    0,
11713                    &ph.slice((np - 1) * n_embd..np * n_embd),
11714                    n_embd,
11715                )?;
11716            } else if continuation {
11717                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11718                    if let Some(lh) = lh.as_ref() {
11719                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
11720                    }
11721                }
11722            }
11723        }
11724        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
11725        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
11726
11727        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
11728        let fork_mode = OptiForkGateMode::configured();
11729        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
11730        // the end. Metric normalization vs the reference engine: BOTH engines count
11731        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
11732        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
11733        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
11734        let mut st_drafted = vec![0usize; k];
11735        let mut st_accepted = vec![0usize; k];
11736        let mut st_len_hist = vec![0usize; k + 1];
11737        let mut st_full = 0usize;
11738        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
11739        // stop the draft chain early when the head's softmax confidence in its own pick drops
11740        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
11741        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
11742        let p_min = *PMIN.get_or_init(|| {
11743            std::env::var("MEMRA_SPEC_PMIN")
11744                .ok()
11745                .and_then(|v| v.parse().ok())
11746                .unwrap_or(0.0)
11747        });
11748        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
11749        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
11750        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
11751        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
11752        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
11753        // verify batch is not); the j==0 exemption stays for pending-less rounds.
11754        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
11755            .map(|v| v == "1")
11756            .unwrap_or(false);
11757
11758        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
11759        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
11760        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
11761        // cuBLAS path in an exotic head) falls back to the eager draft chain.
11762        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
11763        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
11764        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
11765        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
11766        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
11767        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
11768        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
11769        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
11770        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
11771            Some(c) => c,
11772            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
11773        };
11774        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
11775        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
11776        if sampled && dctx.g_q.len() < d_vocab {
11777            dctx.g_q = e.zeros(d_vocab)?;
11778            dctx.g_perturb = e.zeros(d_vocab)?;
11779        }
11780        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
11781        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
11782        // truncation (the correctness backstop) stops cutting every tight-schema round.
11783        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
11784        // shape, so a parked graph of the other shape is dropped and recaptured.
11785        let dmask_on = constraint
11786            .as_deref()
11787            .is_some_and(|c| c.draft_mask_enabled());
11788        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
11789        if dmask_on && dctx.g_dmask.len() < dmask_words {
11790            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
11791            dctx.graph = None; // the old capture baked the old (or no) mask pointer
11792            dctx.failed.clear_greedy();
11793            dctx.keeper.clear();
11794        }
11795        if dctx.graph.is_some() && dctx.graph_masked != dmask_on {
11796            dctx.graph = None;
11797            dctx.failed.clear_greedy();
11798            dctx.keeper.clear();
11799        }
11800        if graph_draft && !sampled && dctx.graph.is_none() && !dctx.failed.greedy_failed() {
11801            // dcw door: the capture warmups append device-counter rows the capture body cannot
11802            // rebase for; pre-arm ring headroom host-side (no-op on flat planes / room-enough
11803            // rings, and the door-off path is untouched).
11804            if step35_draft_dcw_on() {
11805                scratch.ensure_dcw_headroom(e, k + 2)?;
11806            }
11807            let DraftGraphCtx {
11808                g_tok,
11809                g_pos,
11810                g_seed,
11811                g_p,
11812                g_dmask,
11813                ..
11814            } = &mut dctx;
11815            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
11816            // host uploads the position's real words, so the warmups stay grammar-free.
11817            if dmask_on {
11818                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11819            }
11820            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11821            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
11822            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
11823            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
11824            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
11825            // passes (and, in serve, other sessions) recycle those addresses and the replay then
11826            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
11827            let cap_res = e.capture_graph_retained(|e| {
11828                self.mtp_head_forward_cap(
11829                    e,
11830                    mtp,
11831                    g_tok,
11832                    g_pos,
11833                    g_seed,
11834                    g_p,
11835                    &mut *scratch,
11836                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
11837                    true,
11838                    embd_gpu.expect("graph draft requires resident embedding"),
11839                    embd_qt,
11840                    embd_rb,
11841                    d_vocab,
11842                    None,
11843                    None,
11844                    if dmask_on {
11845                        Some((g_dmask_ro, dmask_words))
11846                    } else {
11847                        None
11848                    },
11849                )
11850            });
11851            match cap_res {
11852                Ok((g, keep)) => {
11853                    scratch.set_len(e, base)?;
11854                    dctx.graph = Some(g);
11855                    dctx.graph_masked = dmask_on;
11856                    dctx.keeper = keep;
11857                }
11858                Err(err) => {
11859                    scratch.set_len(e, base)?;
11860                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
11861                    // silent. Once per flip — mark returns None on an already-failed ctx.
11862                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
11863                        eprintln!("{line}");
11864                    }
11865                }
11866            }
11867        }
11868        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
11869        // graph object, built only when sampled && graph-eligible — the greedy capture above is
11870        // untouched (and skipped when sampled: its graph would never be launched). Same head
11871        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
11872        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
11873        // once per round); the raw head logits land in the persistent g_q for the host's
11874        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
11875        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
11876        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
11877        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
11878        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
11879        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
11880        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
11881        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
11882        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
11883        // this compare misses at most ONCE per resumed request — the first burst recaptures
11884        // and every later burst in that request replays. A client that wants the parked graph
11885        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
11886        // stable across its whole conversation.
11887        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
11888        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
11889        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
11890        // force the eager draft (which computes stats/penalties per row).
11891        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
11892        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
11893        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
11894        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
11895        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
11896        // the request shape the vendor-default flip makes the majority).
11897        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
11898        let pure_temp = s_key.pure_temp();
11899        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
11900            dctx.graph_s = None;
11901            dctx.failed.clear_sampled();
11902            dctx.s_key = None;
11903            dctx.q_slots.clear();
11904            dctx.keeper_s.clear();
11905        }
11906        if graph_draft
11907            && sampled
11908            && pure_temp
11909            && dctx.graph_s.is_none()
11910            && !dctx.failed.sampled_failed()
11911        {
11912            // dcw door: same warmup headroom pre-arm as the greedy capture above.
11913            if step35_draft_dcw_on() {
11914                scratch.ensure_dcw_headroom(e, k + 2)?;
11915            }
11916            let DraftGraphCtx {
11917                g_tok,
11918                g_pos,
11919                g_seed,
11920                g_p,
11921                g_ctr,
11922                g_perturb,
11923                g_q,
11924                ..
11925            } = &mut dctx;
11926            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
11927            let cap_res = e.capture_graph_retained(|e| {
11928                self.mtp_head_forward_cap(
11929                    e,
11930                    mtp,
11931                    g_tok,
11932                    g_pos,
11933                    g_seed,
11934                    g_p,
11935                    &mut *scratch,
11936                    p_min > 0.0,
11937                    true,
11938                    embd_gpu.expect("graph draft requires resident embedding"),
11939                    embd_qt,
11940                    embd_rb,
11941                    d_vocab,
11942                    Some((g_ctr, g_perturb, g_q, sp_seed, sp_temp)),
11943                    None,
11944                    None, // constrained spec is greedy-only — sampled never carries a hook
11945                )
11946            });
11947            match cap_res {
11948                Ok((g, keep)) => {
11949                    scratch.set_len(e, base)?;
11950                    for _ in 0..k {
11951                        dctx.q_slots.push(e.zeros(d_vocab)?);
11952                    }
11953                    dctx.graph_s = Some(g);
11954                    dctx.s_key = Some(s_key);
11955                    dctx.keeper_s = keep;
11956                }
11957                Err(err) => {
11958                    scratch.set_len(e, base)?;
11959                    // LOUD flip (audit Q2): same contract as the greedy capture above.
11960                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
11961                        eprintln!("{line}");
11962                    }
11963                }
11964            }
11965        }
11966        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819) ----
11967        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
11968        // captured under this request's exact regime, and capture requires `pure_temp` — so a
11969        // parked graph implies `pure_temp`. That implication is the whole exactness argument for
11970        // the graph arm, so it is asserted here rather than assumed: a future change that widens
11971        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
11972        // fails LOUDLY at this line instead of silently drafting from the raw softmax while the
11973        // verify applies filter stats. Release builds refuse the graph (drop it, draft eager)
11974        // rather than launching it; the launch site re-tests `pure_temp` independently.
11975        if sampled && !pure_temp && dctx.graph_s.is_some() {
11976            debug_assert!(
11977                false,
11978                "sampled draft graph parked under {:?} survived into a FILTERED request \
11979                 (top_k={} top_p={} min_p={} pen_on={}): the in-graph chain draws from the RAW \
11980                 softmax, so the verify's filtered q would test a distribution the draft was \
11981                 never sampled from",
11982                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11983            );
11984            eprintln!(
11985                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
11986                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={}); drafting \
11987                 EAGER — the key must carry every field that shapes q",
11988                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on,
11989            );
11990            dctx.graph_s = None;
11991            dctx.s_key = None;
11992            dctx.q_slots.clear();
11993            dctx.keeper_s.clear();
11994        }
11995        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
11996        // on — is this request sampled, is it in the pure-temp regime the sampled graph is only
11997        // legal in, and is a graph PARKED from an earlier request of the same session? The launch
11998        // arms below print which chain actually ran, so the probe never restates the condition.
11999        if skey_probe() {
12000            eprintln!(
12001                "[skey] burst sampled={} pure_temp={} temp={} top_k={} top_p={} min_p={} \
12002                 pen_on={} k={} graph_draft={} graph_s_parked={} s_key_parked={:?}",
12003                sampled as u8,
12004                pure_temp as u8,
12005                sp_temp,
12006                sp.top_k,
12007                sp.top_p,
12008                sp.min_p,
12009                pen_on as u8,
12010                k,
12011                graph_draft as u8,
12012                dctx.graph_s.is_some() as u8,
12013                dctx.s_key,
12014            );
12015        }
12016        let t_cap = t_ent.elapsed();
12017        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
12018        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
12019        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
12020        // fill: the first chain step processes it and appends its entry at slot prompt.len().
12021        if let Some(ph) = &prompt_h {
12022            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
12023            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
12024            // global positions [base..base+tp). Fresh call: base==0, identical to before.
12025            scratch.set_len(e, base)?;
12026            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
12027            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
12028            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
12029            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
12030            let tp = prompt.len();
12031            let fill_chunk: usize = if crate::cache::swa_ring_on() {
12032                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
12033            } else {
12034                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
12035                // meaning one monolithic fill.
12036                std::env::var("MEMRA_PRIME_CHUNK")
12037                    .ok()
12038                    .and_then(|v| v.parse().ok())
12039                    .unwrap_or(4096)
12040            };
12041            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
12042            let mut start = 0usize;
12043            while start < tp {
12044                let end = (start + fill_chunk).min(tp);
12045                let tc = end - start;
12046                {
12047                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
12048                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
12049                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
12050                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
12051                    let mut phs = e.zeros(tc * n_embd)?;
12052                    let (src_lo, dst_off) = if start == 0 {
12053                        (0, n_embd)
12054                    } else {
12055                        ((start - 1) * n_embd, 0)
12056                    };
12057                    let n_copy = if start == 0 {
12058                        (tc - 1) * n_embd
12059                    } else {
12060                        tc * n_embd
12061                    };
12062                    if start == 0 {
12063                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
12064                            if let Some(lh) = lh.as_ref() {
12065                                e.copy_into(&mut phs, 0, lh, n_embd)?;
12066                            }
12067                        }
12068                    }
12069                    if n_copy > 0 {
12070                        e.copy_view_into(
12071                            &mut phs,
12072                            dst_off,
12073                            &ph.slice(src_lo..src_lo + n_copy),
12074                            n_copy,
12075                        )?;
12076                    }
12077                    self.mtp_kv_fill_all(
12078                        e,
12079                        &prompt[start..end],
12080                        &phs,
12081                        base + start,
12082                        &mut *scratch,
12083                        embd_dev,
12084                    )?;
12085                }
12086                start = end;
12087            }
12088        }
12089        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
12090        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
12091        // (=1 brackets the whole call in run_spec.rs, prime included.)
12092        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
12093            unsafe extern "C" {
12094                fn cudaProfilerStart() -> i32;
12095            }
12096            unsafe {
12097                cudaProfilerStart();
12098            }
12099        }
12100        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
12101        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
12102        // consume each other's device outputs; the host drains the ring every M rounds. v1
12103        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
12104        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
12105        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
12106        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
12107        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
12108        let stream_on = crate::spec::spec_stream()
12109            && !sampled
12110            && !spec_replay
12111            && self.mtp_extra.is_empty()
12112            && constraint.is_none()
12113            && !session_mode
12114            && embd_gpu.is_some()
12115            && !crate::model::full_prec_enabled()
12116            && k + 2 < 96;
12117        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
12118        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
12119        if stream_on {
12120            let cap = e.capture_graph(|e| {
12121                for j in 0..k.max(1) {
12122                    self.mtp_head_forward_cap(
12123                        e,
12124                        mtp,
12125                        &mut dctx.g_tok,
12126                        &mut dctx.g_pos,
12127                        &mut dctx.g_seed,
12128                        &mut dctx.g_p,
12129                        &mut *scratch,
12130                        true,
12131                        true,
12132                        embd_gpu.expect("round stream requires resident embedding"),
12133                        embd_qt,
12134                        embd_rb,
12135                        d_vocab,
12136                        None,
12137                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
12138                        None, // round-stream requires constraint.is_none() (see stream_on)
12139                    )?;
12140                }
12141                Ok(())
12142            });
12143            match cap {
12144                Ok(g) => {
12145                    scratch.set_len(e, 0)?;
12146                    stream_graph = Some(g);
12147                }
12148                Err(err) => {
12149                    scratch.set_len(e, 0)?;
12150                    if debug_spec {
12151                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
12152                    }
12153                }
12154            }
12155        }
12156        let stream_active = stream_on && stream_graph.is_some();
12157        if debug_spec {
12158            eprintln!(
12159                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
12160                crate::spec::spec_stream(),
12161                dctx.graph.is_some(),
12162                stream_graph.is_some()
12163            );
12164        }
12165        let t_v_s = k + 1;
12166        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
12167        // module (extracted 2026-07-12; the gemma burst reuses them).
12168        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
12169        let crate::round_stream::StreamBufs {
12170            mut vtok_d,
12171            mut brk_d,
12172            mut pend_d,
12173            last_pred_d,
12174            mut pos_ctr,
12175            mut pos_start_d,
12176            mut ring_d,
12177            acc_d: mut stream_acc,
12178            m_rounds,
12179            k: _,
12180        } = sb;
12181        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
12182            Some(crate::round_stream::kv_len_ptr_table(
12183                e,
12184                cache,
12185                Some(&pos_ctr),
12186            )?)
12187        } else {
12188            None
12189        };
12190
12191        let t_fill = t_ent.elapsed();
12192        let mut round = 0usize;
12193        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
12194        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
12195        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
12196        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
12197        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
12198        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
12199        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
12200        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
12201        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
12202        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
12203        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
12204        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
12205        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
12206        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
12207        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
12208        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
12209        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
12210        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
12211        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
12212        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
12213        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
12214        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
12215        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
12216        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
12217        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
12218        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
12219        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
12220        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
12221        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
12222        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
12223            .ok()
12224            .and_then(|v| v.parse().ok());
12225        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
12226            4
12227        } else if self.cfg.n_embd as usize >= 2500 {
12228            2
12229        } else {
12230            1
12231        };
12232        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
12233        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
12234        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
12235        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
12236        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
12237            .ok()
12238            .and_then(|v| v.parse().ok())
12239            .unwrap_or(1024);
12240        let floor_at = |pos: usize| -> usize {
12241            if adapt_floor_env.is_some() || pos < floor_ctx {
12242                adapt_floor
12243            } else if adapt_floor >= 4 {
12244                1
12245            } else {
12246                adapt_floor
12247            }
12248        };
12249        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
12250        // fixed-K default path is untouched by this whole block.
12251        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
12252            .ok()
12253            .and_then(|v| v.parse().ok())
12254            .unwrap_or(7);
12255        let k_cap = k.min(cap_max).max(1);
12256        let mut kc = k_cap;
12257        let mut opti_fork: Option<OptiForkState> = None;
12258        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
12259        if fork_mode != OptiForkGateMode::Disabled {
12260            let fence = crate::pp::pp_cuts(self.layers.len());
12261            let refusal = if !session_mode {
12262                Some("not-session")
12263            } else if k != 1 || adapt {
12264                Some("requires-fixed-k1")
12265            } else if sampled || constraint.is_some() || spec_replay {
12266                Some("sampled-constrained-or-replay")
12267            } else if pipe.is_some() {
12268                Some("two-session-pipeline")
12269            } else if !spec_devacc() {
12270                Some("requires-device-accept")
12271            } else if stream_active || crate::spec::spec_stream() {
12272                Some("round-stream")
12273            } else if !self.mtp_extra.is_empty() {
12274                Some("multi-head-mtp")
12275            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
12276                Some("swa-ring")
12277            } else if crate::pp::pp_host_bounce_active() {
12278                Some("host-bounce")
12279            } else if fork_mode == OptiForkGateMode::Controller
12280                && cache.recur.iter().any(Option::is_some)
12281            {
12282                Some("controller-requires-zero-recurrent-state")
12283            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
12284                Some("requires-pp2")
12285            } else {
12286                None
12287            };
12288            if let Some(reason) = refusal {
12289                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12290                eprintln!("[opti-fork] refused reason={reason}");
12291            } else {
12292                let fence = fence.expect("validated PP-2 fence");
12293                let rt = crate::pp::PpNRt::get(e)?;
12294                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
12295                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
12296                let primary_supported =
12297                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
12298                if !rt.cross_device() || !primary_supported {
12299                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12300                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
12301                } else {
12302                    // Both recurrent snapshots and both seed generations are allocated before
12303                    // the first fork, each through its owning PP stage. Allocation failure
12304                    // therefore happens before any optimistic state mutation can occur.
12305                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12306                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12307                    let fork = OptiForkState::new(
12308                        e,
12309                        cache,
12310                        fork_mode,
12311                        alternate_snapshot,
12312                        &h_seed_buf,
12313                        &fill_prev,
12314                        rt,
12315                        fence[1],
12316                        self.layers.len(),
12317                    )?;
12318                    eprintln!(
12319                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
12320                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
12321                        fence[1],
12322                        fork.logical_payload_bytes[0],
12323                        fork.logical_payload_bytes[1],
12324                        fork.controller.map_or(0.0, |policy| policy.threshold),
12325                    );
12326                    fork_snapshot = Some(current_snapshot);
12327                    opti_fork = Some(fork);
12328                }
12329            }
12330        }
12331        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
12332        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
12333        let mut snap = match fork_snapshot {
12334            Some(snapshot) => snapshot,
12335            None => cache.snapshot(e)?,
12336        };
12337        let mut carried_opti: Option<OptiControllerTicket> = None;
12338        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
12339        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
12340        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
12341            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
12342        } else {
12343            None
12344        };
12345        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
12346        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
12347        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
12348        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
12349        // pass of any kind). Verify still
12350        // checks every emitted token against the target -> exactness holds by construction; only
12351        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
12352        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
12353        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
12354        let mut pending: Option<u32> = carried_pending;
12355        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
12356        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
12357        // the verify accept readback). Printed once at loop end via spec-stats.
12358        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
12359        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
12360        // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
12361        // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
12362        // under it) and `verify-wait` is only the residual drain at the accept readback: one
12363        // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
12364        // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
12365        // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
12366        // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
12367        // which is what says the queueing time was hidden and is not a target. Diagnostic only.
12368        let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
12369        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
12370        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
12371        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
12372        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
12373        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
12374        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
12375        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
12376        let mut ph_wait = 0f64;
12377        let mut ph_commit = 0f64;
12378        let mut ph_t = std::time::Instant::now();
12379        let mut ph_mark = |acc: &mut f64, on: bool| {
12380            if on {
12381                let now = std::time::Instant::now();
12382                *acc += (now - ph_t).as_secs_f64();
12383                ph_t = now;
12384            }
12385        };
12386        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
12387        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
12388        // arm holds it — the slab stash is live verify -> commit inside a round, and the
12389        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
12390        // the model (rebuilding per call re-captures the pool per prompt, which is the
12391        // measured way to lose more than the launches cost); the captured bodies are
12392        // cache-independent, every state read going through per-round refreshed pointer
12393        // tables. None = the eager walk, byte-identical.
12394        //
12395        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
12396        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
12397        // whenever the stream is live rather than relying on that refusal.
12398        // The lock is taken ONLY when the door is armed: with the flag off this whole block
12399        // is inert, so the default path cannot serialize two spec generations behind a mutex
12400        // it never reads.
12401        let vg_armed =
12402            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
12403        let mut vg_guard = if vg_armed && !stream_active {
12404            let mut g = self.dspark_vgraphs.lock().unwrap();
12405            if g.is_none() {
12406                // Size by the WIDEST verify this run can present, which is k+1 and NOT
12407                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
12408                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
12409                // panic in the sampled ON arm, measured before this line said k+1).
12410                let vt_cap = (k.max(k_cap) + 1).max(2);
12411                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
12412                if g.is_some() {
12413                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
12414                    // than trusting that a flag set means a pool built.
12415                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
12416                } else {
12417                    eprintln!(
12418                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
12419                         non-uniform state, or vt_cap < 2) — eager walk"
12420                    );
12421                }
12422            }
12423            Some(g)
12424        } else {
12425            None
12426        };
12427        // Capacity fail-safe: a round wider than the pool was built for must take the eager
12428        // walk, not slice the stash past its rows. The sizing above already covers every
12429        // round this run can present; this keeps a future caller (or a k that grows behind
12430        // the pool's back) on the byte-identical fallback instead of a panic.
12431        let vg_t_cap = vg_guard
12432            .as_ref()
12433            .and_then(|g| g.as_ref())
12434            .map(|g| g.t_capacity())
12435            .unwrap_or(0);
12436        if let Some(p) = pipe {
12437            p.setup_end();
12438        }
12439        while keep_going && out.len() < max_new {
12440            // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
12441            // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
12442            // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
12443            // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
12444            // step37 TP2 stack. This prints where the other ~150 ms lives.
12445            let round_prof = ROUND_PROF
12446                .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
12447            let round_t0 = round_prof.then(std::time::Instant::now);
12448            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
12449            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
12450            if let (true, Some(sg), Some(ptrs)) = (
12451                stream_active && round >= 1 && pending.is_some(),
12452                &stream_graph,
12453                &stream_ptrs,
12454            ) {
12455                if debug_spec {
12456                    static ONCE: std::sync::Once = std::sync::Once::new();
12457                    ONCE.call_once(|| {
12458                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
12459                    });
12460                }
12461                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
12462                e.set_u32_one(&mut pend_d, pending.unwrap())?;
12463                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
12464                for _mi in 0..m_rounds {
12465                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
12466                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
12467                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
12468                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
12469                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
12470                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12471                    sg.launch()?;
12472                    e.spec_assemble_verify(
12473                        &g_tokp2k,
12474                        &pend_d,
12475                        d2t_dev.as_ref(),
12476                        &mut vtok_d,
12477                        &mut brk_d,
12478                        p_min,
12479                        k,
12480                        pmin0,
12481                    )?;
12482                    let mut ck = VerifyCkpt::new(self.layers.len());
12483                    let dummy = vec![0u32; t_v_s];
12484                    let (tl_d, vx) = self.decode_step_t_core_stream(
12485                        e,
12486                        &dummy,
12487                        0,
12488                        &mut *cache,
12489                        embd_dev,
12490                        Some(&mut ck),
12491                        Some((&vtok_d, &pos_ctr)),
12492                        None,
12493                        None,
12494                        None,
12495                    )?;
12496                    for j in 0..t_v_s {
12497                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12498                    }
12499                    e.spec_accept_greedy_dc(
12500                        &preds_d,
12501                        &vtok_d,
12502                        &last_pred_d,
12503                        &brk_d,
12504                        &mut stream_acc,
12505                    )?;
12506                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12507                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12508                    self.commit_verified_prefix_stream(
12509                        e,
12510                        &mut *cache,
12511                        &snap,
12512                        &ck,
12513                        &stream_acc,
12514                        1,
12515                        t_v_s,
12516                    )?;
12517                    e.spec_rollback_stream(
12518                        ptrs,
12519                        &pos_start_d,
12520                        &stream_acc,
12521                        1,
12522                        self.layers.len() + 1,
12523                    )?;
12524                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12525                }
12526                e.stream().synchronize()?;
12527                let ring_h = e.dtoh_u32(&ring_d)?;
12528                let cnt = ring_h[0] as usize;
12529                for i in 0..cnt {
12530                    if out.len() < max_new {
12531                        out.push(ring_h[1 + i]);
12532                    }
12533                }
12534                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12535                for il in 0..self.layers.len() {
12536                    if let Some(kvl) = cache.kv[il].as_mut() {
12537                        kvl.len = pos_h;
12538                    }
12539                }
12540                cache.pos = pos_h;
12541                scratch.kv.len = pos_h;
12542                pending = Some(ring_h[cnt]); // last drained token = the live bonus
12543                last_token = ring_h[cnt];
12544                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12545                total_accepted += cnt.saturating_sub(m_rounds);
12546                if let Some(t) = sess_telem {
12547                    // totals only — the burst's per-round accept counts stayed on device
12548                    // (that is the point of the round-stream arm). pos_* untouched.
12549                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12550                }
12551                round += m_rounds;
12552                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12553                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12554                continue;
12555            }
12556            let pipe_draft = match pipe {
12557                Some(p) => Some(p.draft_begin(round)?),
12558                None => None,
12559            };
12560            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12561            let mut current_opti = carried_opti.take();
12562            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
12563                match opti_fork.as_mut() {
12564                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
12565                    None => None,
12566                    Some(_) => None,
12567                }
12568            } else {
12569                None
12570            };
12571            if current_opti.is_none() {
12572                if let Some(fork) = opti_fork.as_ref() {
12573                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
12574                } else {
12575                    cache.snapshot_into(e, &mut snap)?;
12576                }
12577            } else if snap.pos != pos {
12578                return Err(format!(
12579                    "optipipe carried snapshot pos {} != current pos {pos}",
12580                    snap.pos
12581                )
12582                .into());
12583            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
12584            ph_mark(&mut ph_rest, phase_on);
12585
12586            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12587            // p-min semantics (both paths): stop the chain early when the head's confidence in
12588            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12589            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12590            let base0 = if pending.is_some() { 1usize } else { 0usize };
12591            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12592            // accepted run + 1 (the gemma law — see the setup block above the loop).
12593            let k_this = if adapt { kc } else { k };
12594            let mut draft: Vec<u32> = Vec::with_capacity(k);
12595            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12596            let mut controller_draft_prob: Option<f32> = None;
12597            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
12598            if let Some(ticket) = current_opti.as_mut() {
12599                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
12600                if ticket.verify_tokens[0] != carried_pending {
12601                    return Err(format!(
12602                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
12603                        ticket.verify_tokens[0],
12604                    )
12605                    .into());
12606                }
12607                draft.push(ticket.verify_tokens[1]);
12608                controller_draft_prob = Some(ticket.draft_prob);
12609                controller_eager_state = ticket
12610                    .take_eager_seed()
12611                    .map(|seed| (ticket.verify_tokens[1], seed));
12612            } else {
12613                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12614                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12615                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12616                // rejected drafts and p-min extras via the len mechanism).
12617                scratch.set_len(e, pos + base0 - 1)?;
12618                // dcw door: a captured chain appends k_this device-counter rows (plus the
12619                // pseudo-seed replay) with no host intervention; any ring rebase those appends
12620                // could need happens HERE, host-side, before the replays. The eager arm keeps
12621                // its own per-step prepare, so this is graph-path-only work.
12622                if step35_draft_dcw_on() && (dctx.graph.is_some() || dctx.graph_s.is_some()) {
12623                    scratch.ensure_dcw_headroom(e, k_this + 2)?;
12624                }
12625                if pen_on {
12626                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
12627                    // device dedup: the serve window is already PEN_WINDOW_MAX, and this
12628                    // defensive min also bounds non-server callers.
12629                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
12630                    let w0 = pen_hist.len().saturating_sub(win);
12631                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
12632                }
12633                if sampled {
12634                    draft_logits.clear();
12635                    draft_stats.clear();
12636                }
12637                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
12638                // position's mask is computed on that clone and advanced by the PROPOSED token. The
12639                // real state moves only on emission (verify's job), so the emitted stream is
12640                // unchanged — the mask only removes tokens the verify would have truncated anyway.
12641                let mut dmask_live = dmask_on;
12642                if dmask_live {
12643                    let t_c = std::time::Instant::now();
12644                    constraint
12645                        .as_deref_mut()
12646                        .unwrap()
12647                        .draft_begin()
12648                        .map_err(|e2| format!("constraint: {e2}"))?;
12649                    dm_clone_ns += t_c.elapsed().as_nanos();
12650                    dm_rounds += 1;
12651                }
12652                if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
12653                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
12654                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
12655                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
12656                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12657                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
12658                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12659                    for j in 0..k_this {
12660                        // per-position mask upload (contents only — the graph's baked pointer is
12661                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
12662                        // mask node degrades to a no-op ban instead of needing a second graph.
12663                        if dmask_live
12664                            && !upload_draft_mask(
12665                                e,
12666                                constraint.as_deref_mut().unwrap(),
12667                                &mut dctx.g_dmask,
12668                                mtp.d2t.as_ref(),
12669                                d_vocab,
12670                                dmask_words,
12671                            )?
12672                        {
12673                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
12674                            // genuinely miss the legal set): neutralize the captured mask node and
12675                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
12676                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12677                            dmask_live = false;
12678                        }
12679                        gr.launch()?;
12680                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12681                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12682                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
12683                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
12684                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
12685                        // replay's embed node, and the MMU fault kills the CUDA context for the
12686                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
12687                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
12688                        // buffer (g_seed = the verify-side handoff vs head-side compute).
12689                        if (idx as usize) >= d_vocab {
12690                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
12691                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
12692                            // seed, untouched since the round-start copy — the pair discriminates
12693                            // "seed arrived poisoned" from "head forward produced NaN".
12694                            let seed_h = e.dtoh(&dctx.g_seed)?;
12695                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12696                            let in_h = e.dtoh(&h_seed_buf)?;
12697                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
12698                            return Err(format!(
12699                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12700                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
12701                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
12702                             the embed row (#87 trap)"
12703                            )
12704                            .into());
12705                        }
12706                        // trimmed draft vocab -> target token id (identity when no d2t map)
12707                        let d = match &mtp.d2t {
12708                            Some(map) => map[idx as usize],
12709                            None => idx,
12710                        };
12711                        let draft_p = if p_min > 0.0
12712                            || opti_fork
12713                                .as_ref()
12714                                .is_some_and(|fork| fork.controller.is_some())
12715                        {
12716                            Some(e.dtoh(&dctx.g_p)?[0])
12717                        } else {
12718                            None
12719                        };
12720                        if j == 0 {
12721                            controller_draft_prob = draft_p;
12722                        }
12723                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
12724                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12725                                break;
12726                            }
12727                        }
12728                        draft.push(d);
12729                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
12730                        // index the argmax wrote — patch the persistent token buffer (4B htod).
12731                        if d != idx {
12732                            e.set_u32_one(&mut dctx.g_tok, d)?;
12733                        }
12734                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
12735                        // unmasked drafting for the remaining positions (verify still arbitrates).
12736                        // speculative advance; a chain the grammar can no longer follow (EOS
12737                        // proposed) ends here. The captured mask node always runs, so a dead chain
12738                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
12739                        if dmask_live
12740                            && !constraint
12741                                .as_deref_mut()
12742                                .unwrap()
12743                                .draft_advance(d)
12744                                .map_err(|e2| format!("constraint: {e2}"))?
12745                        {
12746                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
12747                            break;
12748                        }
12749                    }
12750                // PURE-TEMP RE-TEST (lane/graph-s-key-exactness-20260819): the sampled graph is
12751                // legal ONLY in the regime it was captured in. The condition used to read
12752                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything else —
12753                // which it could not, because the key omitted the filters. Both halves are now
12754                // enforced: the key drops a stale graph, and this site refuses to launch one.
12755                } else if let (true, Some(gr)) = (sampled && pure_temp, &dctx.graph_s) {
12756                    if skey_probe() {
12757                        eprintln!(
12758                            "[skey] chain=graph_s round={round} pure_temp={} top_k={} \
12759                             top_p={} min_p={} s_key_parked={:?}",
12760                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12761                        );
12762                    }
12763                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
12764                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
12765                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
12766                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
12767                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
12768                    // stream. Host sctr advances in lockstep (computed, no readback needed).
12769                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
12770                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
12771                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12772                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
12773                    for j in 0..k_this {
12774                        gr.launch()?;
12775                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
12776                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
12777                        // counts the p-min-discarded token too)
12778                        // q retention: ONE async D2D of the persistent head-logits buffer into this
12779                        // round's slot j (stream-ordered after the replay, before the next one).
12780                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
12781                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
12782                        // #87 SENTINEL TRAP (see the greedy graph arm above).
12783                        if (idx as usize) >= d_vocab {
12784                            let seed_h = e.dtoh(&dctx.g_seed)?;
12785                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12786                            return Err(format!(
12787                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
12788                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
12789                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
12790                             (#87 trap)"
12791                            )
12792                            .into());
12793                        }
12794                        let d = match &mtp.d2t {
12795                            Some(map) => map[idx as usize],
12796                            None => idx,
12797                        };
12798                        draft_idx.push(idx);
12799                        if p_min > 0.0 {
12800                            let p = e.dtoh(&dctx.g_p)?[0];
12801                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12802                                break;
12803                            }
12804                        }
12805                        draft.push(d);
12806                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
12807                        if d != idx {
12808                            e.set_u32_one(&mut dctx.g_tok, d)?;
12809                        }
12810                    }
12811                    // uniform accept path: fill draft_stats per used slot (pure-temp regime — the
12812                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot, tiny).
12813                    for j in 0..draft.len().max(draft_idx.len()) {
12814                        let rows0 = e.htod_i32(&[0])?;
12815                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12816                        e.filter_stats(
12817                            &dctx.q_slots[j],
12818                            d_vocab,
12819                            &rows0,
12820                            &mut th_d,
12821                            &mut z_d,
12822                            &mut mx_d,
12823                            d_vocab,
12824                            1,
12825                            sp_temp,
12826                            sp.top_k,
12827                            sp.top_p,
12828                            sp.min_p,
12829                        )?;
12830                        draft_stats.push((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
12831                    }
12832                } else {
12833                    if skey_probe() && sampled {
12834                        eprintln!(
12835                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
12836                             top_p={} min_p={} s_key_parked={:?}",
12837                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
12838                        );
12839                    }
12840                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
12841                    let chain_heads = !self.mtp_extra.is_empty();
12842                    let mut e_tok = last_token;
12843                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
12844                    let mut chain_tokens = if chain_heads {
12845                        vec![last_token]
12846                    } else {
12847                        Vec::new()
12848                    };
12849                    let mut chain_seeds = if chain_heads {
12850                        vec![e.clone_dtod(&h_seed_buf)?]
12851                    } else {
12852                        Vec::new()
12853                    };
12854                    for j in 0..k_this {
12855                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
12856                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
12857                        let mtp_pos = pos + base0 + j;
12858                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
12859                        // A position with no legal draft-vocab row drops to unmasked drafting for
12860                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
12861                        if dmask_live {
12862                            dmask_live = upload_draft_mask(
12863                                e,
12864                                constraint.as_deref_mut().unwrap(),
12865                                &mut dctx.g_dmask,
12866                                mtp.d2t.as_ref(),
12867                                d_vocab,
12868                                dmask_words,
12869                            )?;
12870                        }
12871                        let mask = if dmask_live {
12872                            Some((&dctx.g_dmask, dmask_words))
12873                        } else {
12874                            None
12875                        };
12876                        let (dl_d, h_nextn) = if chain_heads {
12877                            if debug_spec {
12878                                eprintln!(
12879                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
12880                                    mtp_chain_head_index(j, self.mtp_head_count()),
12881                                    chain_tokens.len(),
12882                                );
12883                            }
12884                            self.mtp_chain_forward_dev(
12885                                e,
12886                                &chain_tokens,
12887                                &chain_seeds,
12888                                &mut *scratch,
12889                                pos + base0 - 1,
12890                                embd_dev,
12891                                mask,
12892                            )?
12893                        } else {
12894                            self.mtp_head_forward_dev(
12895                                e,
12896                                mtp,
12897                                e_tok,
12898                                &d_seed,
12899                                &mut *scratch,
12900                                mtp_pos,
12901                                embd_dev,
12902                                mask,
12903                            )?
12904                        };
12905                        let tok_d = if sampled {
12906                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
12907                            // the filtered softmax (filters off => th=0, exact v1 semantics).
12908                            if perturb_buf.is_none() {
12909                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
12910                            }
12911                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
12912                            if pen_on {
12913                                let h = pen_hist_d.as_ref().unwrap();
12914                                let nh = h.len();
12915                                e.penalize_logits(
12916                                    &mut q_row,
12917                                    h,
12918                                    nh,
12919                                    sp.penalty_repeat,
12920                                    sp.penalty_freq,
12921                                    sp.penalty_present,
12922                                    d_vocab,
12923                                )?;
12924                            }
12925                            let rows0 = e.htod_i32(&[0])?;
12926                            let (mut th_d, mut z_d, mut mx_d) =
12927                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
12928                            e.filter_stats(
12929                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
12930                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
12931                            )?;
12932                            let (th, z, mx) =
12933                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
12934                            let pb = perturb_buf.as_mut().unwrap();
12935                            e.gumbel_perturb_filtered(
12936                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
12937                            )?;
12938                            sctr += 1;
12939                            draft_logits.push(q_row);
12940                            draft_stats.push((mx, th, z));
12941                            e.argmax_token_device(pb, d_vocab)?
12942                        } else {
12943                            e.argmax_token_device(&dl_d, d_vocab)?
12944                        };
12945                        let idx = e.dtoh_u32_one(&tok_d)?;
12946                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
12947                        // here because the eager chain's operands are all readable: dl_d (the head
12948                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
12949                        if (idx as usize) >= d_vocab {
12950                            let dl_h = e.dtoh(&dl_d)?;
12951                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
12952                            let seed_h = if chain_heads {
12953                                e.dtoh(chain_seeds.last().unwrap())?
12954                            } else {
12955                                e.dtoh(&d_seed)?
12956                            };
12957                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
12958                            return Err(format!(
12959                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
12960                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
12961                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
12962                             embed row (#87 trap)"
12963                            )
12964                            .into());
12965                        }
12966                        let d = match &mtp.d2t {
12967                            Some(map) => map[idx as usize],
12968                            None => idx,
12969                        };
12970                        if sampled {
12971                            draft_idx.push(idx);
12972                        }
12973                        let draft_p = if p_min > 0.0
12974                            || opti_fork
12975                                .as_ref()
12976                                .is_some_and(|fork| fork.controller.is_some())
12977                        {
12978                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
12979                            Some(e.dtoh(&p_d)?[0])
12980                        } else {
12981                            None
12982                        };
12983                        if j == 0 {
12984                            controller_draft_prob = draft_p;
12985                        }
12986                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
12987                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
12988                                break;
12989                            }
12990                        }
12991                        draft.push(d);
12992                        if chain_heads {
12993                            chain_tokens.push(d);
12994                            chain_seeds.push(h_nextn);
12995                        } else {
12996                            e_tok = d;
12997                            d_seed = h_nextn;
12998                        }
12999                        // speculative advance; a chain the grammar can no longer follow (EOS
13000                        // proposed) ends here — the prefix already proposed still rides verify.
13001                        if dmask_live
13002                            && !constraint
13003                                .as_deref_mut()
13004                                .unwrap()
13005                                .draft_advance(d)
13006                                .map_err(|e2| format!("constraint: {e2}"))?
13007                        {
13008                            break;
13009                        }
13010                    }
13011                    if !chain_heads
13012                        && opti_fork
13013                            .as_ref()
13014                            .is_some_and(|fork| fork.controller.is_some())
13015                    {
13016                        controller_eager_state = Some((e_tok, d_seed));
13017                    }
13018                }
13019            }
13020            let k_round = draft.len();
13021            if let Some(p) = pipe {
13022                p.draft_end(round);
13023            }
13024            drop(pipe_draft);
13025
13026            ph_mark(&mut ph_draft, phase_on);
13027            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
13028            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
13029            let verify_tokens: Vec<u32> = match pending {
13030                Some(b) => {
13031                    let mut v = Vec::with_capacity(k_round + 1);
13032                    v.push(b);
13033                    v.extend_from_slice(&draft);
13034                    v
13035                }
13036                None => draft.clone(),
13037            };
13038            let base = if pending.is_some() { 1 } else { 0 };
13039            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
13040            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
13041            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
13042                Some(ticket.take_ckpt())
13043            } else if spec_replay {
13044                None
13045            } else {
13046                Some(VerifyCkpt::new(self.layers.len()))
13047            };
13048            let controller_can_probe = base == 1
13049                && k_round == 1
13050                && out.len().saturating_add(2) < max_new
13051                && controller_draft_prob.is_some()
13052                && opti_fork
13053                    .as_ref()
13054                    .and_then(|fork| fork.controller.as_ref())
13055                    .is_some_and(|policy| !policy.breaker_tripped);
13056            let mut successor_attempt: Option<OptiControllerTicket> = None;
13057            let mut rejected_probe: Option<(f32, u32)> = None;
13058            let mut controller_prepared: Option<OptiControllerPrepared> = None;
13059            if controller_can_probe {
13060                // Prepare d2/q and, on admission, d3 before either current verify half is
13061                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
13062                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
13063                // the primary stream after N stage 1 would serialize the supposed pipeline.
13064                let eager_pos = scratch.kv.len + 1;
13065                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
13066                    e,
13067                    mtp,
13068                    &mut dctx,
13069                    &mut *scratch,
13070                    d_vocab,
13071                    &mut controller_eager_state,
13072                    eager_pos,
13073                    embd_dev,
13074                )?;
13075                let first_probability = controller_draft_prob
13076                    .ok_or("optipipe controller probe lost first-token probability")?;
13077                let q_proxy = first_probability * pending_probability;
13078                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13079                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13080                let admitted = opti_fork
13081                    .as_ref()
13082                    .and_then(|fork| fork.controller.as_ref())
13083                    .ok_or("optipipe controller policy disappeared")?
13084                    .admit(q_proxy);
13085                if admitted {
13086                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13087                    let eager_pos = scratch.kv.len + 1;
13088                    let (optimistic_draft, optimistic_draft_probability) = self
13089                        .opti_controller_draft_step(
13090                            e,
13091                            mtp,
13092                            &mut dctx,
13093                            &mut *scratch,
13094                            d_vocab,
13095                            &mut controller_eager_state,
13096                            eager_pos,
13097                            embd_dev,
13098                        )?;
13099                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13100                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
13101                        debug_assert_eq!(token, optimistic_draft);
13102                        seed
13103                    });
13104                    controller_prepared = Some(OptiControllerPrepared {
13105                        verify_tokens: [optimistic_pending, optimistic_draft],
13106                        draft_prob: optimistic_draft_probability,
13107                        eager_seed,
13108                        q_proxy,
13109                        scratch_len: scratch.kv.len,
13110                    });
13111                } else {
13112                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13113                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13114                    rejected_probe = Some((q_proxy, optimistic_pending));
13115                    eprintln!(
13116                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
13117                        opti_fork
13118                            .as_ref()
13119                            .and_then(|fork| fork.controller.as_ref())
13120                            .expect("controller policy")
13121                            .threshold,
13122                    );
13123                }
13124            }
13125            let fork_attempt = match fork_generation.take() {
13126                Some(generation) if base == 1 && k_round == 1 => Some(generation),
13127                Some(generation) => {
13128                    opti_fork
13129                        .as_mut()
13130                        .expect("fork generation without fork state")
13131                        .retire(generation)?;
13132                    None
13133                }
13134                None => None,
13135            };
13136            let (tlogits_d, vx) = if let Some(p) = pipe {
13137                self.decode_step_t_core_pipelined(
13138                    e,
13139                    &verify_tokens,
13140                    pos,
13141                    &mut *cache,
13142                    embd_dev,
13143                    ckpt.as_mut(),
13144                    p,
13145                    round,
13146                )?
13147            } else if controller_can_probe {
13148                let fence = opti_fork
13149                    .as_ref()
13150                    .ok_or("optipipe controller probe lost fork state")?
13151                    .fence;
13152                let boundary = match current_opti.as_mut() {
13153                    Some(ticket) => ticket.take_boundary(),
13154                    None => self.verify_stage0_issue(
13155                        e,
13156                        &verify_tokens,
13157                        pos,
13158                        &mut *cache,
13159                        embd_dev,
13160                        ckpt.as_mut(),
13161                        None,
13162                        &fence,
13163                        Some(true),
13164                        None,
13165                    )?,
13166                };
13167                if let Some(prepared) = controller_prepared.take() {
13168                    let generation = {
13169                        let fork = opti_fork
13170                            .as_mut()
13171                            .ok_or("optipipe controller admission lost fork state")?;
13172                        let generation = fork.reserve_successor()?;
13173                        let rt = fork.rt;
13174                        let snapshot_fence = fork.fence;
13175                        opti_snapshot_one_stage_owned_into(
13176                            e,
13177                            cache,
13178                            rt,
13179                            &snapshot_fence,
13180                            0,
13181                            fork.successor_snapshot_mut(),
13182                        )?;
13183                        generation
13184                    };
13185                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
13186                    let successor_boundary = self.verify_stage0_issue(
13187                        e,
13188                        &prepared.verify_tokens,
13189                        pos + verify_tokens.len(),
13190                        &mut *cache,
13191                        embd_dev,
13192                        Some(&mut successor_ckpt),
13193                        None,
13194                        &fence,
13195                        Some(false),
13196                        None,
13197                    )?;
13198                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13199                    let fork = opti_fork
13200                        .as_ref()
13201                        .ok_or("optipipe controller ticket lost fork state")?;
13202                    successor_attempt = Some(fork.controller_ticket(
13203                        generation,
13204                        successor_boundary,
13205                        successor_ckpt,
13206                        prepared.verify_tokens,
13207                        prepared.draft_prob,
13208                        prepared.eager_seed,
13209                        prepared.q_proxy,
13210                        prepared.scratch_len,
13211                    ));
13212                    eprintln!(
13213                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
13214                         verify={:?}",
13215                        generation.id,
13216                        prepared.q_proxy,
13217                        fork.controller.expect("controller policy").threshold,
13218                        prepared.verify_tokens,
13219                    );
13220                }
13221                let result = self.verify_stage1_finish(
13222                    e,
13223                    boundary,
13224                    &mut *cache,
13225                    ckpt.as_mut(),
13226                    None,
13227                    &fence,
13228                    successor_attempt.is_none(),
13229                )?;
13230                if let Some(ticket) = current_opti.as_mut() {
13231                    ticket.settle();
13232                }
13233                if successor_attempt.is_some() {
13234                    let fork = opti_fork
13235                        .as_mut()
13236                        .ok_or("optipipe successor snapshot lost fork state")?;
13237                    let rt = fork.rt;
13238                    let snapshot_fence = fork.fence;
13239                    opti_snapshot_one_stage_owned_into(
13240                        e,
13241                        cache,
13242                        rt,
13243                        &snapshot_fence,
13244                        1,
13245                        fork.successor_snapshot_mut(),
13246                    )?;
13247                    // Publish N only after both independent successor-state queues are complete.
13248                    fork.rt.publish_to(1, &e.stream())?;
13249                }
13250                result
13251            } else if let Some(ticket) = current_opti.as_mut() {
13252                let fork = opti_fork
13253                    .as_mut()
13254                    .ok_or("optipipe carried controller ticket lost fork state")?;
13255                let boundary = ticket.take_boundary();
13256                let result = self.verify_stage1_finish(
13257                    e,
13258                    boundary,
13259                    &mut *cache,
13260                    ckpt.as_mut(),
13261                    None,
13262                    &fork.fence,
13263                    true,
13264                )?;
13265                ticket.settle();
13266                result
13267            } else if let Some(generation) = fork_attempt {
13268                let fork = opti_fork
13269                    .as_mut()
13270                    .expect("fork generation without fork state");
13271                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
13272                let action = fork.mode.action(generation.id);
13273                let boundary = self.verify_stage0_issue(
13274                    e,
13275                    &verify_tokens,
13276                    pos,
13277                    &mut *cache,
13278                    embd_dev,
13279                    ckpt.as_mut(),
13280                    None,
13281                    &fork.fence,
13282                    Some(true),
13283                    None,
13284                )?;
13285                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13286                let mut ticket = fork.ticket(generation, boundary);
13287                if action == OptiForkAction::Abort {
13288                    return Err(format!(
13289                        "optipipe forced abort with generation {} stage0 in flight",
13290                        generation.id,
13291                    )
13292                    .into());
13293                }
13294                fork.reconcile(
13295                    e,
13296                    &mut *cache,
13297                    &mut *scratch,
13298                    &snap,
13299                    &mut h_seed_buf,
13300                    &mut fill_prev,
13301                    generation,
13302                    action,
13303                    verify_tokens[0],
13304                )?;
13305                let result = if action == OptiForkAction::Hit {
13306                    let boundary = ticket.take_boundary();
13307                    self.verify_stage1_finish(
13308                        e,
13309                        boundary,
13310                        &mut *cache,
13311                        ckpt.as_mut(),
13312                        None,
13313                        &fork.fence,
13314                        true,
13315                    )?
13316                } else {
13317                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
13318                    // verify only after E_restart published the restored stage-0 state.
13319                    self.decode_step_t_core(
13320                        e,
13321                        &verify_tokens,
13322                        pos,
13323                        &mut *cache,
13324                        embd_dev,
13325                        ckpt.as_mut(),
13326                    )?
13327                };
13328                ticket.settle();
13329                debug_assert_eq!(ticket.generation, generation);
13330                fork.retire(generation)?;
13331                result
13332            } else {
13333                // The serial verify every non-fork round takes — the MTP route's
13334                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
13335                // a pool above, and then the walk replays the captured trunk instead of
13336                // re-issuing it launch by launch.
13337                let vg_round = if verify_tokens.len() <= vg_t_cap {
13338                    vg_guard.as_mut().and_then(|g| g.as_mut())
13339                } else {
13340                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
13341                        // The commit reads this flag to pick its arm; a round that declines
13342                        // the pool must not inherit a stale `true` from the round before it.
13343                        g.round_slab = false;
13344                    }
13345                    None
13346                };
13347                self.decode_step_t_core_vg(
13348                    e,
13349                    &verify_tokens,
13350                    pos,
13351                    &mut *cache,
13352                    embd_dev,
13353                    ckpt.as_mut(),
13354                    vg_round,
13355                )?
13356            };
13357            let pipe_accept = match pipe {
13358                Some(p) => Some(p.accept_begin(round)?),
13359                None => None,
13360            };
13361
13362            if phase_sync {
13363                e.stream().synchronize()?;
13364            }
13365            ph_mark(&mut ph_verify, phase_on);
13366            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
13367            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
13368            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
13369            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
13370            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
13371            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
13372            // (== the bonus), so every index shifts by `base` and last_pred is unused.
13373            let t_v = verify_tokens.len();
13374            let mut preds: Vec<u32> = Vec::new();
13375            if !sampled {
13376                for j in 0..t_v {
13377                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
13378                }
13379                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
13380                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
13381                // next round's last_token = the next chain's embed lookup. Catch it at the
13382                // source with the column named — an all-NaN VERIFY column implicates the
13383                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
13384                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
13385                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
13386                    let mut probe = e.zeros(n_vocab)?;
13387                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
13388                    let col_h = e.dtoh(&probe)?;
13389                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
13390                    return Err(format!(
13391                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
13392                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
13393                         — the verify TRUNK produced a poisoned column (#87 trap). Run \
13394                         MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
13395                         that layer into attention and routed MoE). NOT the draft head, and NOT \
13396                         the PP stage split this message used to name: pp_cuts() returns None \
13397                         without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
13398                         that variable is set.",
13399                        preds[bad]
13400                    )
13401                    .into());
13402                }
13403            }
13404            ph_mark(&mut ph_wait, phase_on);
13405            let t_pred = |j: usize| -> u32 {
13406                if j == 0 && base == 0 {
13407                    last_pred
13408                } else {
13409                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
13410                    // used to call this from the sampled arm and panicked the worker; it now goes
13411                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
13412                    // out-of-range pred is a real bug, not something to paper over.
13413                    debug_assert!(
13414                        !sampled,
13415                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
13416                    );
13417                    preds[base + j - 1]
13418                }
13419            };
13420            let mut devacc_seeded = false;
13421            let mut devacc_acc: Option<CudaSlice<u32>> = None;
13422            let (n_acc, bonus) = if !sampled {
13423                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
13424                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
13425                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
13426                // gated on token identity vs the host walk (the arms below are bit-equal rules).
13427                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
13428                {
13429                    let draft_d = e.htod_u32_v(&draft)?;
13430                    let mut acc_out = e.alloc_u32_zeroed(2)?;
13431                    e.spec_accept_greedy(
13432                        &preds_d,
13433                        &draft_d,
13434                        last_pred,
13435                        base,
13436                        k_round,
13437                        &mut acc_out,
13438                    )?;
13439                    devacc_acc = Some(acc_out.clone());
13440                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
13441                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
13442                    // non-replay commit arms skip their host-offset seed copies (guarded below);
13443                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
13444                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
13445                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
13446                    // the update lands after the arms (devacc_seeded guard below).
13447                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
13448                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
13449                    // unified rule; full accept rewrites the verify-left value). Host mirrors
13450                    // update after the readback; commit_verified_prefix skips its len_d writes.
13451                    if let Some(successor) = successor_attempt.as_ref() {
13452                        opti_fork
13453                            .as_mut()
13454                            .ok_or("optipipe successor reconcile lost fork state")?
13455                            .queue_actual_reconcile(
13456                                e,
13457                                &snap,
13458                                &acc_out,
13459                                successor.verify_tokens[0],
13460                                base,
13461                            )?;
13462                    } else if let Some(ptrs) = &kv_len_ptrs {
13463                        let saved: Vec<i32> = (0..self.layers.len())
13464                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
13465                            .collect();
13466                        let saved_d = e.htod_i32(&saved)?;
13467                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
13468                    }
13469                    devacc_seeded = true;
13470                    let ab = e.dtoh_u32(&acc_out)?;
13471                    (ab[0] as usize, ab[1])
13472                } else {
13473                    let mut n_acc = 0usize;
13474                    for j in 0..k_round {
13475                        if t_pred(j) == draft[j] {
13476                            n_acc += 1;
13477                        } else {
13478                            break;
13479                        }
13480                    }
13481                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
13482                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
13483                    (n_acc, t_pred(n_acc))
13484                }
13485            } else {
13486                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
13487                if col_buf.is_none() {
13488                    col_buf = Some(e.zeros(n_vocab)?);
13489                }
13490                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
13491                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
13492                let mut pj = vec![0f32; k_round.max(1)];
13493                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
13494                if k_round > 0 {
13495                    let mut ids: Vec<u32> = Vec::new();
13496                    let mut rows: Vec<i32> = Vec::new();
13497                    for j in 0..k_round {
13498                        if j > 0 || base == 1 {
13499                            ids.push(draft[j]);
13500                            rows.push((base + j) as i32 - 1);
13501                        }
13502                    }
13503                    if !ids.is_empty() {
13504                        let nr = rows.len();
13505                        // penalties: materialize the used columns into one contiguous penalized
13506                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
13507                        // penalties: materialize used columns contiguously, penalize all rows in
13508                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
13509                        let p_rows: Vec<i32> = if pen_on {
13510                            (0..nr as i32).collect()
13511                        } else {
13512                            rows.clone()
13513                        };
13514                        if pen_on {
13515                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
13516                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
13517                            }
13518                            let pc = pcol_buf.as_mut().unwrap();
13519                            for (i2, &r) in rows.iter().enumerate() {
13520                                let c = r as usize;
13521                                e.copy_view_into(
13522                                    pc,
13523                                    i2 * n_vocab,
13524                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
13525                                    n_vocab,
13526                                )?;
13527                            }
13528                            let h = pen_hist_d.as_ref().unwrap();
13529                            let nh = h.len();
13530                            e.penalize_logits_rows(
13531                                pc,
13532                                h,
13533                                nh,
13534                                sp.penalty_repeat,
13535                                sp.penalty_freq,
13536                                sp.penalty_present,
13537                                n_vocab,
13538                                nr,
13539                            )?;
13540                        }
13541                        let p_src: &CudaSlice<f32> = if pen_on {
13542                            pcol_buf.as_ref().unwrap()
13543                        } else {
13544                            &tlogits_d
13545                        };
13546                        let rowsd = e.htod_i32(&p_rows)?;
13547                        let (mut th_d, mut z_d, mut mx_d) =
13548                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
13549                        e.filter_stats(
13550                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
13551                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
13552                        )?;
13553                        let idsd = e.htod_u32_v(&ids)?;
13554                        let mut outd = e.zeros(nr)?;
13555                        e.softmax_gather_filtered(
13556                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
13557                            sp_temp,
13558                        )?;
13559                        let outv = e.dtoh(&outd)?;
13560                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
13561                        let mut oi = 0usize;
13562                        for j in 0..k_round {
13563                            if j > 0 || base == 1 {
13564                                pj[j] = outv[oi];
13565                                oi += 1;
13566                            }
13567                        }
13568                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
13569                    }
13570                    if base == 0 {
13571                        let lc: &CudaSlice<f32> = if pen_on {
13572                            if col_buf.is_none() {
13573                                col_buf = Some(e.zeros(n_vocab)?);
13574                            }
13575                            let cb = col_buf.as_mut().unwrap();
13576                            e.copy_into(
13577                                cb,
13578                                0,
13579                                last_col_logits
13580                                    .as_ref()
13581                                    .expect("sampled: last_col_logits unset"),
13582                                n_vocab,
13583                            )?;
13584                            let h = pen_hist_d.as_ref().unwrap();
13585                            let nh = h.len();
13586                            e.penalize_logits(
13587                                cb,
13588                                h,
13589                                nh,
13590                                sp.penalty_repeat,
13591                                sp.penalty_freq,
13592                                sp.penalty_present,
13593                                n_vocab,
13594                            )?;
13595                            col_buf.as_ref().unwrap()
13596                        } else {
13597                            last_col_logits
13598                                .as_ref()
13599                                .expect("sampled: last_col_logits unset")
13600                        };
13601                        let rows0 = e.htod_i32(&[0])?;
13602                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13603                        e.filter_stats(
13604                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13605                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
13606                        )?;
13607                        let idsd = e.htod_u32_v(&[draft[0]])?;
13608                        let mut outd = e.zeros(1)?;
13609                        e.softmax_gather_filtered(
13610                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
13611                        )?;
13612                        pj[0] = e.dtoh(&outd)?[0];
13613                        last_col_stats =
13614                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
13615                    }
13616                }
13617                // q source: the graph arm retained the head logits in the persistent q_slots;
13618                // the eager arm in per-round draft_logits clones. Same raw-logit values either way.
13619                // FILTERED q_j: stats from draft_stats (eager pushes in-chain; the graph arm
13620                // computes them post-replay — graph engages only filter/penalty-free, so the
13621                // stats degenerate to th=0/full-Z there, keeping ONE accept path).
13622                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() {
13623                    &dctx.q_slots
13624                } else {
13625                    &draft_logits
13626                };
13627                let mut n_acc = 0usize;
13628                for j in 0..k_round {
13629                    let (qmx, qth, qz) = draft_stats[j];
13630                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
13631                    let rowsd = e.htod_i32(&[0])?;
13632                    let thd = e.htod(&[qth])?;
13633                    let zd = e.htod(&[qz])?;
13634                    let _ = qmx;
13635                    let mut outd = e.zeros(1)?;
13636                    e.softmax_gather_filtered(
13637                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
13638                        sp_temp,
13639                    )?;
13640                    let qj = e.dtoh(&outd)?[0];
13641                    let u = host_u01(sp_seed, uctr);
13642                    uctr += 1;
13643                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
13644                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
13645                    // exactness signature (see `skey_probe`). Impossible when the draft was
13646                    // drawn from the same filtered distribution the verify reconstructs here;
13647                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
13648                    if skey_probe() && qj == 0.0 {
13649                        eprintln!(
13650                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
13651                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
13652                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
13653                        );
13654                    }
13655                    if accept {
13656                        n_acc += 1;
13657                    } else {
13658                        break;
13659                    }
13660                }
13661                let bonus = if n_acc == k_round {
13662                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
13663                    let col = base + k_round - 1;
13664                    let cb = col_buf.as_mut().unwrap();
13665                    e.copy_view_into(
13666                        cb,
13667                        0,
13668                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13669                        n_vocab,
13670                    )?;
13671                    if pen_on {
13672                        let h = pen_hist_d.as_ref().unwrap();
13673                        let nh = h.len();
13674                        e.penalize_logits(
13675                            cb,
13676                            h,
13677                            nh,
13678                            sp.penalty_repeat,
13679                            sp.penalty_freq,
13680                            sp.penalty_present,
13681                            n_vocab,
13682                        )?;
13683                    }
13684                    if perturb_buf.is_none() {
13685                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13686                    }
13687                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
13688                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
13689                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
13690                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
13691                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
13692                    // last gathered column, in both base arms. `th` is a threshold in e-units of
13693                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
13694                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
13695                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
13696                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
13697                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
13698                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
13699                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
13700                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
13701                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
13702                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
13703                    // and row_max is unused once nothing is masked), so this fix is a byte-level
13704                    // no-op for the untruncated serve default. One extra one-block filter_stats
13705                    // per full-accept round is the whole cost.
13706                    let (mx, th) = {
13707                        let rows0 = e.htod_i32(&[0])?;
13708                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13709                        let cb0 = col_buf.as_ref().unwrap();
13710                        e.filter_stats(
13711                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
13712                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
13713                        )?;
13714                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
13715                    };
13716                    let pb = perturb_buf.as_mut().unwrap();
13717                    let cb2 = col_buf.as_ref().unwrap();
13718                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
13719                    sctr += 1;
13720                    let td = e.argmax_token_device(pb, n_vocab)?;
13721                    e.dtoh_u32_one(&td)?
13722                } else {
13723                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
13724                    let cb = col_buf.as_mut().unwrap();
13725                    if n_acc > 0 || base == 1 {
13726                        let col = base + n_acc - 1;
13727                        e.copy_view_into(
13728                            cb,
13729                            0,
13730                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
13731                            n_vocab,
13732                        )?;
13733                    } else {
13734                        let lc = last_col_logits.as_ref().unwrap();
13735                        e.copy_into(cb, 0, lc, n_vocab)?;
13736                    }
13737                    if pen_on {
13738                        let h = pen_hist_d.as_ref().unwrap();
13739                        let nh = h.len();
13740                        e.penalize_logits(
13741                            cb,
13742                            h,
13743                            nh,
13744                            sp.penalty_repeat,
13745                            sp.penalty_freq,
13746                            sp.penalty_present,
13747                            n_vocab,
13748                        )?;
13749                    }
13750                    let cb2 = col_buf.as_ref().unwrap();
13751                    let sc = sctr;
13752                    sctr += 1;
13753                    // p-stats for the reject column: from col_stats when the col was gathered,
13754                    // else (j==0&&base==0) from last_col_stats.
13755                    let p_stats = if n_acc > 0 || base == 1 {
13756                        // col index within the gathered set == number of gathered cols before n_acc
13757                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
13758                        col_stats.get(gi).copied().unwrap_or_else(|| {
13759                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
13760                        })
13761                    } else {
13762                        last_col_stats.expect("sampled: last_col_stats unset at reject")
13763                    };
13764                    let q_stats = draft_stats[n_acc];
13765                    if let Some(map) = &d2t_dev {
13766                        if q_full_buf.is_none() {
13767                            q_full_buf = Some(e.zeros(n_vocab)?);
13768                        }
13769                        let qf = q_full_buf.as_mut().unwrap();
13770                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
13771                        let qf2 = q_full_buf.as_ref().unwrap();
13772                        e.residual_sample_filtered(
13773                            cb2,
13774                            Some(qf2),
13775                            n_vocab,
13776                            sp_temp,
13777                            sp_seed,
13778                            sc,
13779                            p_stats,
13780                            q_stats,
13781                            &mut sample_tok,
13782                        )?;
13783                    } else {
13784                        e.residual_sample_filtered(
13785                            cb2,
13786                            Some(&q_bufs[n_acc]),
13787                            n_vocab,
13788                            sp_temp,
13789                            sp_seed,
13790                            sc,
13791                            p_stats,
13792                            q_stats,
13793                            &mut sample_tok,
13794                        )?;
13795                    }
13796                    e.dtoh_u32(&sample_tok)?[0]
13797                };
13798                (
13799                    n_acc,
13800                    guard_vocab_token(
13801                        bonus,
13802                        n_vocab,
13803                        &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
13804                    )?,
13805                )
13806            };
13807            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
13808            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
13809            // ordering). Walk the accepted drafts through the grammar in commit order; the
13810            // first illegal token truncates acceptance at its slot, and that slot's emission
13811            // is recomputed as the MASKED argmax of the target's own verify column — token-
13812            // identical to constrained plain greedy decode (an unmasked argmax that is
13813            // grammar-legal IS the masked argmax: masking only removes competitors). The
13814            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
13815            // measured in acceptance numbers, never hidden.
13816            let (n_acc, bonus) = match constraint.as_deref_mut() {
13817                None => (n_acc, bonus),
13818                Some(c) => {
13819                    fn ce(e2: String) -> Box<dyn std::error::Error> {
13820                        format!("constraint: {e2}").into()
13821                    }
13822                    let mut na = n_acc;
13823                    let mut cut = false;
13824                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
13825                        if c.is_allowed(d).map_err(ce)? {
13826                            c.consume(d).map_err(ce)?;
13827                        } else {
13828                            na = j;
13829                            cut = true;
13830                            dm_cut_tokens += n_acc - j;
13831                            break;
13832                        }
13833                    }
13834                    if cut {
13835                        dm_cuts += 1;
13836                    }
13837                    let mut bo = bonus;
13838                    if cut || !c.is_allowed(bo).map_err(ce)? {
13839                        let mut row = if na == 0 && base == 0 {
13840                            init_logits_host
13841                                .clone()
13842                                .ok_or("constraint: init logits missing (round-0 cut)")?
13843                        } else {
13844                            e.dtoh_view(
13845                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
13846                            )?
13847                        };
13848                        c.mask_logits(&mut row).map_err(ce)?;
13849                        bo = argmax(&row) as u32;
13850                    }
13851                    c.consume(bo).map_err(ce)?;
13852                    (na, bo)
13853                }
13854            };
13855            let mut successor_valid = false;
13856            if let Some((q_proxy, expected_d2)) = rejected_probe {
13857                let v_n = n_acc == 1 && bonus == expected_d2;
13858                eprintln!(
13859                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
13860                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
13861                );
13862            }
13863            if let Some(successor) = successor_attempt.as_ref() {
13864                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
13865                let generation = successor.generation;
13866                let q_proxy = successor.q_proxy;
13867                let expected_pending = successor.verify_tokens[0];
13868                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
13869                let fork = opti_fork
13870                    .as_mut()
13871                    .ok_or("optipipe successor resolution lost fork state")?;
13872                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
13873                if successor_valid {
13874                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13875                } else {
13876                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13877                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13878                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
13879                }
13880                let breaker_tripped = fork
13881                    .controller
13882                    .as_mut()
13883                    .expect("controller policy")
13884                    .resolve(successor_valid);
13885                if breaker_tripped {
13886                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13887                }
13888                eprintln!(
13889                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
13890                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
13891                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
13892                    generation.id, successor_valid, !successor_valid, breaker_tripped,
13893                );
13894                if !successor_valid {
13895                    let mut successor = successor_attempt
13896                        .take()
13897                        .expect("controller successor disappeared on miss");
13898                    successor.settle();
13899                    fork.retire(generation)?;
13900                }
13901            }
13902            total_drafted += k_round;
13903            total_accepted += n_acc;
13904            if let Some(t) = sess_telem {
13905                // Greedy, rejection-sampling, and grammar truncation all converge here after
13906                // the accept decision is already on host. Fixed-size relaxed atomics only.
13907                t.record_round(k_round, n_acc);
13908            }
13909            if spec_stats {
13910                st_len_hist[k_round] += 1;
13911                for j in 0..k_round {
13912                    st_drafted[j] += 1;
13913                }
13914                for j in 0..n_acc {
13915                    st_accepted[j] += 1;
13916                }
13917                if n_acc == k_round {
13918                    st_full += 1;
13919                }
13920            }
13921
13922            if debug_spec {
13923                eprintln!(
13924                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
13925                    out.len(),
13926                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
13927                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
13928                    // the GPU worker thread — a debug flag that killed the exact regime you would
13929                    // set it to investigate. See `debug_t_pred0`.
13930                    debug_t_pred0(sampled, base, last_pred, &preds)
13931                );
13932            }
13933
13934            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
13935            let commit_started = std::time::Instant::now();
13936            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
13937            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
13938            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
13939            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
13940            for j in 0..n_acc {
13941                if !session_mode && out.len() >= max_new {
13942                    break;
13943                }
13944                out.push(draft[j]);
13945            }
13946            if pen_on {
13947                pen_hist.extend_from_slice(&draft[0..n_acc]);
13948                pen_hist.push(bonus);
13949            }
13950            let bonus_emitted = session_mode || out.len() < max_new;
13951            if bonus_emitted {
13952                out.push(bonus);
13953            }
13954            last_token = bonus;
13955
13956            // --- 5. ROLLBACK + advance (§C) ---
13957            if n_acc == k_round && !spec_replay {
13958                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
13959                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
13960                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
13961                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
13962                // last_pred is dead in the pending path (t_pred reads verify col 0).
13963                //
13964                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
13965                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
13966                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
13967                // trunk hidden (the last verify column). set_len first: a p-min break may have
13968                // left one extra chain append at that slot. Partial accepts need NO fill (the
13969                // chain already covered every accepted position; round-start set_len truncates).
13970                let mut vh_seed = e.zeros(n_embd)?;
13971                e.copy_view_into(
13972                    &mut vh_seed,
13973                    0,
13974                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
13975                    n_embd,
13976                )?;
13977                if refresh {
13978                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
13979                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
13980                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
13981                    // the full stack (vx) is already resident from the verify. Replaces both the
13982                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
13983                    // (draft attention quality); exactness stays the verify's job.
13984                    scratch.set_len(e, pos)?;
13985                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
13986                    // (hidden of the last committed row before this verify batch).
13987                    let mut vxs = e.zeros(t_v * n_embd)?;
13988                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
13989                    if t_v > 1 {
13990                        e.copy_view_into(
13991                            &mut vxs,
13992                            n_embd,
13993                            &vx.slice(0..(t_v - 1) * n_embd),
13994                            (t_v - 1) * n_embd,
13995                        )?;
13996                    }
13997                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
13998                } else {
13999                    scratch.set_len(e, pos + base + k_round - 1)?;
14000                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
14001                    let mut hp = e.zeros(n_embd)?;
14002                    if t_v >= 2 {
14003                        e.copy_view_into(
14004                            &mut hp,
14005                            0,
14006                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
14007                            n_embd,
14008                        )?;
14009                    } else {
14010                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
14011                    }
14012                    self.mtp_kv_fill_all(
14013                        e,
14014                        &[draft[k_round - 1]],
14015                        &hp,
14016                        pos + base + k_round - 1,
14017                        &mut *scratch,
14018                        embd_dev,
14019                    )?;
14020                }
14021                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
14022                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
14023                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
14024                // col). Saves one MTP-block pass per round on top of the pairing fix.
14025                if !devacc_seeded {
14026                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
14027                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
14028                }
14029                pending = Some(bonus);
14030                if debug_spec {
14031                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
14032                }
14033            } else if !spec_replay && base + n_acc >= 1 {
14034                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
14035                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
14036                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
14037                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
14038                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
14039                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
14040                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
14041                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
14042                // accept (never compounds: the next verify recomputes true hiddens for all
14043                // committed columns).
14044                let j = base + n_acc;
14045                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
14046                // column stash was written into the graphs ctx's persistent slabs as in-graph
14047                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
14048                // commit must take the slab twin (same semantics, slab-addressed sources). The
14049                // ctx states which of the two this round produced via `round_slab`; trusting the
14050                // flag rather than the env keeps a round that fell back to the eager walk (a
14051                // capture that declined, a t the pool never captured) on the cols arm.
14052                let slab_commit = vg_guard
14053                    .as_ref()
14054                    .and_then(|g| g.as_ref())
14055                    .map(|g| g.round_slab)
14056                    .unwrap_or(false);
14057                if slab_commit {
14058                    self.dspark_commit_prefix_slab(
14059                        e,
14060                        &mut *cache,
14061                        &snap,
14062                        vg_guard
14063                            .as_ref()
14064                            .and_then(|g| g.as_ref())
14065                            .expect("slab_commit implies a graphs ctx"),
14066                        j,
14067                    )?;
14068                } else {
14069                    self.commit_verified_prefix(
14070                        e,
14071                        &mut *cache,
14072                        &snap,
14073                        ckpt.as_ref().unwrap(),
14074                        j,
14075                        devacc_seeded,
14076                        if devacc_seeded {
14077                            devacc_acc.as_ref().map(|a| (a, base, t_v))
14078                        } else {
14079                            None
14080                        },
14081                    )?;
14082                }
14083                let mut seed = e.zeros(n_embd)?;
14084                e.copy_view_into(
14085                    &mut seed,
14086                    0,
14087                    &vx.slice((j - 1) * n_embd..j * n_embd),
14088                    n_embd,
14089                )?;
14090                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
14091                // branch); without it the chain entries stand and only the tail truncates. Either
14092                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
14093                // (persistent mode), rope pos+j+1 (chain convention).
14094                if refresh {
14095                    scratch.set_len(e, pos)?;
14096                    let mut vxs = e.zeros(j * n_embd)?;
14097                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14098                    if j > 1 {
14099                        e.copy_view_into(
14100                            &mut vxs,
14101                            n_embd,
14102                            &vx.slice(0..(j - 1) * n_embd),
14103                            (j - 1) * n_embd,
14104                        )?;
14105                    }
14106                    self.mtp_kv_fill_all(
14107                        e,
14108                        &verify_tokens[0..j],
14109                        &vxs,
14110                        pos,
14111                        &mut *scratch,
14112                        embd_dev,
14113                    )?;
14114                } else {
14115                    scratch.set_len(e, pos + j)?;
14116                }
14117                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
14118                // bonus's predecessor (verify col j-1); no pseudo pass.
14119                if !devacc_seeded {
14120                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
14121                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
14122                }
14123                pending = Some(bonus);
14124                if debug_spec {
14125                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
14126                }
14127            } else if !spec_replay {
14128                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
14129                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
14130                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
14131                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
14132                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
14133                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
14134                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
14135                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
14136                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
14137                cache.rollback(e, &snap, 0)?;
14138                scratch.set_len(e, pos)?;
14139                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14140                pending = Some(bonus);
14141                if debug_spec {
14142                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
14143                }
14144            } else {
14145                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
14146                // this round survives, only possible before the first pending exists, ~round 0):
14147                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
14148                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
14149                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
14150                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
14151                // trunk hidden.
14152                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
14153                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
14154                if let Some(b) = pending.take() {
14155                    replay.push(b);
14156                }
14157                replay.extend_from_slice(&draft[0..n_acc]);
14158                replay.push(bonus);
14159                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
14160                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
14161                // last col exactly as before (byte-identical to the old _h_emb_dev call).
14162                let (rl_d, rx) = if self.batched_serving_numeric_class() {
14163                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
14164                    let mut hidden = e.uninit(replay.len() * n_embd)?;
14165                    for (row, &token) in replay.iter().enumerate() {
14166                        let (row_logits, row_hidden) =
14167                            self.spec_target_step_h(e, token, &mut *cache)?;
14168                        logits.extend_from_slice(&row_logits);
14169                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
14170                    }
14171                    (e.htod(&logits)?, hidden)
14172                } else {
14173                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
14174                };
14175                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
14176                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
14177                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
14178                last_pred = guard_vocab_token(
14179                    e.dtoh_u32(&preds_d)?[0],
14180                    n_vocab,
14181                    &format!("replay last_pred at round {round} pos={pos}"),
14182                )?;
14183                if sampled {
14184                    let lr0 = replay.len();
14185                    let lc = last_col_logits
14186                        .as_mut()
14187                        .expect("sampled: last_col_logits unset");
14188                    e.copy_view_into(
14189                        lc,
14190                        0,
14191                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
14192                        n_vocab,
14193                    )?;
14194                }
14195                let lr = replay.len();
14196                if lr >= 2 {
14197                    e.copy_view_into(
14198                        &mut h_seed_buf,
14199                        0,
14200                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
14201                        n_embd,
14202                    )?;
14203                } else {
14204                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
14205                    // last_token, whose own-row hidden fill_prev still holds.
14206                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14207                }
14208                // the bonus is COMMITTED here — it becomes the last committed row.
14209                let mut rh_last = e.zeros(n_embd)?;
14210                e.copy_view_into(
14211                    &mut rh_last,
14212                    0,
14213                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
14214                    n_embd,
14215                )?;
14216                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
14217                if debug_spec {
14218                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
14219                }
14220            }
14221            if devacc_seeded {
14222                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
14223                // consumed the old value (both slots carry the same value in every non-replay arm).
14224                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
14225            }
14226            if successor_valid {
14227                let optimistic_scratch_len = successor_attempt
14228                    .as_ref()
14229                    .expect("valid controller successor disappeared")
14230                    .scratch_len;
14231                // The normal current-round commit refreshed/truncated the logical scratch tail.
14232                // Its optimistic successor row was already written physically, so restoring only
14233                // the retained logical length makes that row live for the carried round.
14234                scratch.set_len(e, optimistic_scratch_len)?;
14235            }
14236            if let Some(current) = current_opti.take() {
14237                opti_fork
14238                    .as_mut()
14239                    .ok_or("optipipe current retirement lost fork state")?
14240                    .retire(current.generation)?;
14241            }
14242            if successor_valid {
14243                let successor = successor_attempt
14244                    .take()
14245                    .expect("valid controller successor disappeared before promotion");
14246                let generation = successor.generation;
14247                opti_fork
14248                    .as_mut()
14249                    .ok_or("optipipe successor promotion lost fork state")?
14250                    .promote_successor_snapshot(&mut snap, generation);
14251                carried_opti = Some(successor);
14252            }
14253            if anatomy_on {
14254                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
14255                // only for this diagnostic so it does not disappear into the following draft's
14256                // first token readback.
14257                e.stream().synchronize()?;
14258                ph_commit += commit_started.elapsed().as_secs_f64();
14259            }
14260            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
14261            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
14262            // final position — the floor's position key reads the committed depth). Burst
14263            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
14264            // like gemma's burst arm.
14265            if adapt {
14266                let fl_now = floor_at(cache.pos);
14267                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
14268            }
14269            ph_mark(&mut ph_rest, phase_on);
14270            if let Some(p) = pipe {
14271                p.accept_end(round);
14272            }
14273            drop(pipe_accept);
14274            if let Some(t0) = round_t0 {
14275                let ms = t0.elapsed().as_secs_f64() * 1e3;
14276                ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
14277                let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
14278                if n % 32 == 0 {
14279                    eprintln!(
14280                        "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
14281                        ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
14282                        out.len()
14283                    );
14284                }
14285            }
14286            round += 1;
14287            // sse-cadence: this round's accepted drafts + bonus are committed (out is
14288            // append-only past step 4) — flush at round cadence.
14289            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
14290        }
14291        if let Some(mut ticket) = carried_opti.take() {
14292            opti_fork
14293                .as_mut()
14294                .ok_or("optipipe tail drain lost fork state")?
14295                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
14296        }
14297        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
14298        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
14299        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
14300
14301        if spec_stats {
14302            let per_slot: Vec<String> = (0..k)
14303                .map(|j| {
14304                    if st_drafted[j] > 0 {
14305                        format!(
14306                            "{}/{}={:.3}",
14307                            st_accepted[j],
14308                            st_drafted[j],
14309                            st_accepted[j] as f64 / st_drafted[j] as f64
14310                        )
14311                    } else {
14312                        "0/0".into()
14313                    }
14314                })
14315                .collect();
14316            let acc = if total_drafted > 0 {
14317                total_accepted as f64 / total_drafted as f64
14318            } else {
14319                0.0
14320            };
14321            eprintln!(
14322                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
14323                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
14324                       tok_per_round={:.3}",
14325                per_slot.join(" "),
14326                (total_accepted + round) as f64 / round.max(1) as f64
14327            );
14328        }
14329        if constraint.is_some() {
14330            eprintln!(
14331                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
14332                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
14333                dm_clone_ns as f64 / 1e6,
14334                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
14335            );
14336        }
14337        if phase_on {
14338            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
14339            eprintln!(
14340                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
14341                ph_draft * 1e3,
14342                ph_draft / tot * 100.0,
14343                ph_verify * 1e3,
14344                ph_verify / tot * 100.0,
14345                ph_wait * 1e3,
14346                ph_wait / tot * 100.0,
14347                ph_rest * 1e3,
14348                ph_rest / tot * 100.0
14349            );
14350        }
14351        if anatomy_on {
14352            let rounds_f = round.max(1) as f64;
14353            let other = (ph_rest - ph_commit).max(0.0);
14354            eprintln!(
14355                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
14356                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
14357                ph_draft * 1e3 / rounds_f,
14358                ph_verify * 1e3 / rounds_f,
14359                ph_wait * 1e3 / rounds_f,
14360                ph_commit * 1e3 / rounds_f,
14361                other * 1e3 / rounds_f,
14362            );
14363        }
14364        let _pipe_tail = pipe.map(|p| p.primary());
14365        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
14366        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
14367        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
14368        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
14369        if let Some(slot) = sess_draft_slot.take() {
14370            *slot = Some(dctx);
14371        }
14372        let t_rounds = t_ent.elapsed();
14373        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
14374            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
14375            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
14376            // HERE, where the sampler, the session Philox counters and the penalty window are
14377            // all live and the boundary logits row still exists — that is the "make the state
14378            // available" half of the fix; the consuming burst then just emits it. `sctr` is
14379            // written to the session BELOW the draws so the advance is never lost.
14380            *next_pred_slot = Some(last_pred);
14381            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
14382            let mut stashed_pending = false;
14383            if let Some(b) = pending.take() {
14384                if !sampled {
14385                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
14386                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
14387                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
14388                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
14389                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
14390                    // OUT of `committed` (cache rows == committed); the consuming call
14391                    // prepends it once its verify commits the row. next_pred is unknowable
14392                    // without the commit pass — None; callers gate on pending_tok too.
14393                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
14394                    if let Some(slot) = sess_pending_slot.take() {
14395                        *slot = Some(b);
14396                    }
14397                    *next_pred_slot = None;
14398                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
14399                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
14400                    *last_h = Some(e.clone_dtod(&fill_prev)?);
14401                    stashed_pending = true;
14402                } else {
14403                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
14404                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
14405                    let pos_b = cache.pos;
14406                    scratch.set_len(e, pos_b)?;
14407                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
14408                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
14409                    // itself — the prediction AFTER the bonus never materialized; it would have
14410                    // been the next round's verify col 0). The commit's logits ARE that
14411                    // prediction — so they are also the row the next burst's boundary token
14412                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
14413                    *next_pred_slot = Some(if sample_boundary {
14414                        sample_boundary_token(
14415                            e,
14416                            &lg_b,
14417                            &sp,
14418                            &pen_hist,
14419                            &mut sctr,
14420                            "burst-tail-commit",
14421                        )?
14422                    } else {
14423                        argmax(&lg_b) as u32
14424                    });
14425                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
14426                    *last_h = Some(hb);
14427                }
14428            } else {
14429                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
14430                *last_h = Some(e.clone_dtod(&fill_prev)?);
14431                if sample_boundary {
14432                    // No pending to commit, so the boundary row is the one `last_pred` was
14433                    // argmaxed from and the sampled path keeps it on device: the init feed's
14434                    // logits when the burst ran zero rounds, else the legacy-replay path's
14435                    // last verify column (both predict the token AFTER the last committed
14436                    // row). It is retained precisely because round 0's accept test needs it,
14437                    // so the draw costs no extra D2H of the [n_vocab] row.
14438                    match last_col_logits.as_ref() {
14439                        Some(lc) => {
14440                            *next_pred_slot = Some(sample_boundary_token_dev(
14441                                e,
14442                                lc,
14443                                n_vocab,
14444                                &sp,
14445                                &pen_hist,
14446                                &mut sctr,
14447                                "burst-tail-nopending",
14448                            )?);
14449                        }
14450                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
14451                        // burst always feeds or replays, so the row exists — but if it ever
14452                        // is, the stream takes a greedy token and SAYS so rather than
14453                        // silently regressing to the pre-lane behaviour.
14454                        None => eprintln!(
14455                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
14456                             (reason: no retained boundary logits row)"
14457                        ),
14458                    }
14459                }
14460            }
14461            *sctr_slot = sctr;
14462            *uctr_slot = uctr;
14463            committed.extend_from_slice(prompt);
14464            if let Some(cb) = carried_pending {
14465                // the consumed carry's cache row landed in round 0's verify (every pending
14466                // round commits col 0) — it joins `committed` here, in sequence order.
14467                committed.push(cb);
14468            }
14469            if stashed_pending {
14470                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
14471                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
14472                // 18446744073709551615 out of range for slice of length 0", killing the
14473                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
14474                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
14475                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
14476                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
14477                // did). So a burst that stashes a pending without emitting anything of its own —
14478                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
14479                // guard skipping every token under a tight budget — arrives here with
14480                // out.len() == 0 and stashed_pending == true.
14481                //
14482                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
14483                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
14484                // just above is already accounted. Saturating, not a min/assert: an empty `out`
14485                // here is a legitimate burst shape, not a corrupt state.
14486                let emitted = out.len().saturating_sub(1);
14487                committed.extend_from_slice(&out[..emitted]);
14488            } else {
14489                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
14490            }
14491            debug_assert_eq!(
14492                cache.pos,
14493                committed.len(),
14494                "session invariant: cache rows == committed tokens"
14495            );
14496            if setup_trace {
14497                e.stream().synchronize()?; // bound the async tail fill in the trace
14498                let t_tail = t_ent.elapsed();
14499                eprintln!(
14500                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
14501                    t_init.as_secs_f64() * 1e3,
14502                    (t_cap - t_init).as_secs_f64() * 1e3,
14503                    (t_fill - t_cap).as_secs_f64() * 1e3,
14504                    (t_rounds - t_fill).as_secs_f64() * 1e3,
14505                    (t_tail - t_rounds).as_secs_f64() * 1e3,
14506                    t_tail.as_secs_f64() * 1e3,
14507                    out.len(),
14508                    continuation
14509                );
14510            }
14511            return Ok((out, total_drafted, total_accepted));
14512        }
14513        out.truncate(max_new);
14514        Ok((out, total_drafted, total_accepted))
14515    }
14516
14517    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
14518    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
14519    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
14520    pub fn extract_dspark_anchors(
14521        &self,
14522        e: &Engine,
14523        tokens: &[u32],
14524        anchor_positions: &[usize],
14525        gamma: usize,
14526        top_k: usize,
14527        chunk: usize,
14528        temperature: f32,
14529    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
14530        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
14531            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
14532        }
14533        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
14534            return Err("DSpark anchor positions must be sorted and unique".into());
14535        }
14536        for &position in anchor_positions {
14537            if position == 0 || position + gamma >= tokens.len() {
14538                return Err(format!(
14539                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
14540                    tokens.len()
14541                )
14542                .into());
14543            }
14544        }
14545
14546        let n_vocab = self.output.out_features();
14547        let n_embd = self.cfg.n_embd as usize;
14548        let mut cache =
14549            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
14550        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
14551        let embd_gpu = if spec_host_embd() {
14552            None
14553        } else {
14554            Some(
14555                self.embd_gpu
14556                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
14557            )
14558        };
14559        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
14560
14561        struct PendingRecord {
14562            position: usize,
14563            hidden: Option<Vec<f32>>,
14564            tokens: Vec<u32>,
14565            target_top_ids: Vec<Option<Vec<u32>>>,
14566            target_top_logits: Vec<Option<Vec<f32>>>,
14567            target_top_probs: Vec<Option<Vec<f32>>>,
14568            target_tail_probs: Vec<Option<f32>>,
14569        }
14570
14571        let mut pending: Vec<PendingRecord> = anchor_positions
14572            .iter()
14573            .map(|&position| PendingRecord {
14574                position,
14575                hidden: None,
14576                tokens: tokens[position..=position + gamma].to_vec(),
14577                target_top_ids: vec![None; gamma],
14578                target_top_logits: vec![None; gamma],
14579                target_top_probs: vec![None; gamma],
14580                target_tail_probs: vec![None; gamma],
14581            })
14582            .collect();
14583
14584        let mut start = 0usize;
14585        while start < tokens.len() {
14586            let end = (start + chunk).min(tokens.len());
14587            let chunk_tokens = &tokens[start..end];
14588            let (target_logits, hidden_rows) =
14589                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
14590            for record in &mut pending {
14591                let hidden_position = record.position - 1;
14592                if hidden_position >= start && hidden_position < end {
14593                    let local = hidden_position - start;
14594                    record.hidden = Some(
14595                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
14596                    );
14597                }
14598                for slot in 0..gamma {
14599                    let target_row = record.position + slot;
14600                    if target_row < start || target_row >= end {
14601                        continue;
14602                    }
14603                    let local = target_row - start;
14604                    let logits =
14605                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
14606                    let (ids, top_logits, probs, tail) =
14607                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
14608                    record.target_top_ids[slot] = Some(ids);
14609                    record.target_top_logits[slot] = Some(top_logits);
14610                    record.target_top_probs[slot] = Some(probs);
14611                    record.target_tail_probs[slot] = Some(tail);
14612                }
14613            }
14614            start = end;
14615        }
14616
14617        pending
14618            .into_iter()
14619            .map(|record| {
14620                let hidden = record
14621                    .hidden
14622                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
14623                let target_top_ids =
14624                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
14625                let target_top_logits = flatten_dspark_rows(
14626                    record.target_top_logits,
14627                    record.position,
14628                    "target logits",
14629                )?;
14630                let target_top_probs =
14631                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
14632                let target_tail_probs = record
14633                    .target_tail_probs
14634                    .into_iter()
14635                    .enumerate()
14636                    .map(|(slot, value)| {
14637                        value.ok_or_else(|| {
14638                            format!("missing DSpark tail at {} slot {slot}", record.position)
14639                        })
14640                    })
14641                    .collect::<Result<Vec<_>, _>>()?;
14642                Ok(DsparkAnchorRecord {
14643                    position: record.position,
14644                    hidden,
14645                    tokens: record.tokens,
14646                    target_top_ids,
14647                    target_top_logits,
14648                    target_top_probs,
14649                    target_tail_probs,
14650                })
14651            })
14652            .collect()
14653    }
14654
14655    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
14656    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
14657    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
14658    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
14659    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
14660    /// quant-induced head/hidden-state mismatch from text drift.
14661    ///
14662    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
14663    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
14664    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
14665    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
14666    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
14667    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
14668    ///              conditions on the corpus — deterministic and arm-comparable by design.
14669    ///
14670    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
14671    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
14672    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
14673    ///
14674    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
14675    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
14676    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
14677    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
14678    /// agreement vs this path — not usable as a training-data source).
14679    pub fn replay_acceptance(
14680        &self,
14681        e: &Engine,
14682        tokens: &[u32],
14683        k: usize,
14684        stride: usize,
14685        chunk: usize,
14686        mut hdump: Option<&mut std::fs::File>,
14687    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
14688        assert!(k >= 1 && stride >= 1 && chunk >= 2);
14689        let mtp = self
14690            .mtp
14691            .as_ref()
14692            .expect("replay_acceptance requires an MTP head");
14693        let n_vocab = self.output.out_features();
14694        let d_vocab = mtp
14695            .shared_head_head
14696            .as_ref()
14697            .unwrap_or(&self.output)
14698            .out_features();
14699        let n_embd = self.cfg.n_embd as usize;
14700        let t_total = tokens.len();
14701        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
14702        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
14703        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
14704        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
14705        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
14706        let embd_gpu = if spec_host_embd() {
14707            None
14708        } else {
14709            Some(
14710                self.embd_gpu
14711                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
14712            )
14713        };
14714        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
14715
14716        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
14717        let mut bg: Vec<u32> = vec![0; t_total + 1];
14718        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
14719        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
14720        let mut seed_buf = e.zeros(n_embd)?;
14721        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
14722        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
14723        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
14724        let mut s = 0usize;
14725        while s < t_total {
14726            let cend = (s + chunk).min(t_total);
14727            let tc = cend - s;
14728            let ch = &tokens[s..cend];
14729            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
14730            //    the chunk's true hiddens.
14731            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
14732            for j in 0..tc {
14733                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
14734            }
14735            let preds = e.dtoh_u32(&preds_d)?;
14736            for j in 0..tc {
14737                bg[s + j + 1] = preds[j];
14738            }
14739            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
14740            // checkpoint-quality metric (position j's logits score the GOLD next token).
14741            if nll_on {
14742                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
14743                if jmax > 0 {
14744                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
14745                    let rows: Vec<i32> = (0..jmax as i32).collect();
14746                    let idsd = e.htod_u32_v(&ids)?;
14747                    let rowsd = e.htod_i32(&rows)?;
14748                    let mut outd = e.zeros(jmax)?;
14749                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
14750                    for pr in e.dtoh(&outd)? {
14751                        nll_sum += -((pr.max(1e-30)) as f64).ln();
14752                        nll_cnt += 1;
14753                    }
14754                }
14755            }
14756            if let Some(f) = hdump.as_deref_mut() {
14757                use std::io::Write;
14758                let host: Vec<f32> = e.dtoh(&vx)?;
14759                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
14760                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
14761                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
14762                for v in &host[..tc * n_embd] {
14763                    let b = v.to_bits();
14764                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
14765                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
14766                }
14767                f.write_all(&bytes)?;
14768            }
14769            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
14770            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
14771            // per token saved; the forced trunk pass + hdump is all the mode needs).
14772            let chainless = stride > t_total;
14773            if chainless {
14774                e.copy_view_into(
14775                    &mut prev_last_h,
14776                    0,
14777                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
14778                    n_embd,
14779                )?;
14780                s = cend;
14781                continue;
14782            }
14783            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
14784            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
14785            let mut vxs = e.zeros(tc * n_embd)?;
14786            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
14787            if tc > 1 {
14788                e.copy_view_into(
14789                    &mut vxs,
14790                    n_embd,
14791                    &vx.slice(0..(tc - 1) * n_embd),
14792                    (tc - 1) * n_embd,
14793                )?;
14794            }
14795            scratch.set_len(e, s)?;
14796            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14797            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
14798            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
14799            //    truncates those approximate appends before they can ever be read.
14800            let ps: Vec<usize> = (s..cend)
14801                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
14802                .collect();
14803            for &p in ps.iter().rev() {
14804                scratch.set_len(e, p)?;
14805                if p == s {
14806                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
14807                } else {
14808                    e.copy_view_into(
14809                        &mut seed_buf,
14810                        0,
14811                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
14812                        n_embd,
14813                    )?;
14814                }
14815                let mut e_tok = tokens[p];
14816                let mut d_seed = e.clone_dtod(&seed_buf)?;
14817                let chain_heads = !self.mtp_extra.is_empty();
14818                let mut chain_tokens = if chain_heads {
14819                    vec![tokens[p]]
14820                } else {
14821                    Vec::new()
14822                };
14823                let mut chain_seeds = if chain_heads {
14824                    vec![e.clone_dtod(&seed_buf)?]
14825                } else {
14826                    Vec::new()
14827                };
14828                let mut drafts: Vec<u32> = Vec::with_capacity(k);
14829                for j in 0..k {
14830                    let (dl_d, h_nextn) = if chain_heads {
14831                        self.mtp_chain_forward_dev(
14832                            e,
14833                            &chain_tokens,
14834                            &chain_seeds,
14835                            &mut scratch,
14836                            p,
14837                            embd_dev,
14838                            None,
14839                        )?
14840                    } else {
14841                        self.mtp_head_forward_dev(
14842                            e,
14843                            mtp,
14844                            e_tok,
14845                            &d_seed,
14846                            &mut scratch,
14847                            p + 1 + j,
14848                            embd_dev,
14849                            None,
14850                        )?
14851                    };
14852                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
14853                    let idx = e.dtoh_u32_one(&tok_d)?;
14854                    let d = match &mtp.d2t {
14855                        Some(map) => map[idx as usize],
14856                        None => idx,
14857                    };
14858                    drafts.push(d);
14859                    if chain_heads {
14860                        chain_tokens.push(d);
14861                        chain_seeds.push(h_nextn);
14862                    } else {
14863                        e_tok = d;
14864                        d_seed = h_nextn;
14865                    }
14866                }
14867                // targets may live in a LATER chunk's bg — resolved after the walk.
14868                rows.push((p, drafts, Vec::new()));
14869            }
14870            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
14871            //    expect scratch.len == cend with exact rows).
14872            scratch.set_len(e, s)?;
14873            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
14874            e.copy_view_into(
14875                &mut prev_last_h,
14876                0,
14877                &vx.slice((tc - 1) * n_embd..tc * n_embd),
14878                n_embd,
14879            )?;
14880            s = cend;
14881        }
14882        for (p, drafts, targets) in rows.iter_mut() {
14883            for j in 0..drafts.len() {
14884                targets.push(bg[*p + 1 + j]);
14885            }
14886        }
14887        rows.sort_by_key(|r| r.0);
14888        if nll_cnt > 0 {
14889            let mean = nll_sum / nll_cnt as f64;
14890            println!(
14891                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
14892                mean.exp()
14893            );
14894        }
14895        Ok((rows, bg))
14896    }
14897}
14898
14899#[cfg(test)]
14900mod vg_debt_tests {
14901    use super::dspark_vg_debt_projection;
14902
14903    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
14904    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
14905    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
14906    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
14907    /// impossible must zero the debt.
14908    #[test]
14909    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
14910        const MIB: usize = 1 << 20;
14911        let d = dspark_vg_debt_projection;
14912        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
14913        assert_eq!(d(0, 256, 0, None), 0);
14914        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
14915        assert_eq!(d(10, 0, 500 * MIB, None), 0);
14916        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
14917        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
14918        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
14919
14920        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
14921        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
14922        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
14923
14924        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
14925        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
14926        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
14927        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
14928
14929        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
14930        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
14931        assert_eq!(debt, 250 * (40 * MIB));
14932        assert!(
14933            debt > 3 * (1536 * MIB),
14934            "real growth must dwarf SPEC_SHRINK_RESERVE"
14935        );
14936
14937        // a shrinking/recycled reading never becomes a negative charge.
14938        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
14939        // a stale observation at the same capture count falls back to bootstrap.
14940        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
14941    }
14942}
14943
14944#[cfg(test)]
14945mod mtp_chain_tests {
14946    use super::mtp_chain_head_index;
14947
14948    #[test]
14949    fn embedded_step_heads_cycle_in_declared_order() {
14950        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
14951        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
14952    }
14953
14954    #[test]
14955    fn standalone_draft_remains_single_head() {
14956        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
14957    }
14958}
14959
14960#[cfg(test)]
14961mod tp_verified_prefix_tests {
14962    use super::rewind_tp_kv_verified_prefix;
14963    use crate::tp::ResidentTpKvCache;
14964
14965    fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
14966        let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
14967        let transaction = cache.begin_transaction().unwrap();
14968        let target = cache.append_target(transaction, committed).unwrap();
14969        cache.publish_append(transaction, target).unwrap();
14970        let target = cache.commit_target(transaction, committed).unwrap();
14971        cache.publish_finalize(transaction, target).unwrap();
14972        cache
14973    }
14974
14975    #[test]
14976    fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
14977        let mut layers = vec![Some(cache_with_committed_len(5)), None];
14978        rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
14979        let cache = layers[0].as_ref().unwrap();
14980        assert_eq!(cache.committed_len(), 3);
14981        assert_eq!(cache.staged_len(), 3);
14982    }
14983
14984    #[test]
14985    fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
14986        let mut layers = vec![Some(cache_with_committed_len(1))];
14987        let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
14988            .unwrap_err()
14989            .to_string();
14990        assert!(error.contains("changed shape"), "unexpected error: {error}");
14991    }
14992}
14993
14994#[cfg(test)]
14995mod dspark_sparse_tests {
14996    use super::dspark_sparse_softmax_topk;
14997
14998    #[test]
14999    fn topk_keeps_full_softmax_mass_and_stable_ties() {
15000        let logits = [1.0f32, 3.0, 3.0, -2.0];
15001        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
15002        assert_eq!(ids, vec![1, 2]);
15003        assert_eq!(top_logits, vec![3.0, 3.0]);
15004        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
15005        let expected = 1.0 / denominator;
15006        assert!((probs[0] - expected).abs() < 1.0e-6);
15007        assert!((probs[1] - expected).abs() < 1.0e-6);
15008        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
15009        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
15010    }
15011}
15012
15013#[cfg(test)]
15014mod spec_replay_env_tests {
15015    use super::spec_replay_env_on;
15016
15017    #[test]
15018    fn replay_requires_literal_one() {
15019        assert!(!spec_replay_env_on(None));
15020        assert!(!spec_replay_env_on(Some("")));
15021        assert!(!spec_replay_env_on(Some("0")));
15022        assert!(!spec_replay_env_on(Some("true")));
15023        assert!(!spec_replay_env_on(Some("2")));
15024        assert!(spec_replay_env_on(Some("1")));
15025    }
15026}
15027
15028#[cfg(test)]
15029mod telem_tests {
15030    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
15031
15032    #[test]
15033    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
15034        let counters = SpecTelemetryCounters::default();
15035        for mask in [
15036            [true, true, true],
15037            [true, true, false],
15038            [true, false, false],
15039            [false, false, false],
15040        ] {
15041            let accepted = mask.iter().take_while(|&&value| value).count();
15042            counters.record_round(mask.len(), accepted);
15043        }
15044
15045        let snapshot = counters.snapshot();
15046        assert_eq!(
15047            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
15048            (4, 12, 6)
15049        );
15050        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
15051        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
15052        assert_eq!(snapshot.tau(), 1.5);
15053        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15054        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
15055    }
15056
15057    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
15058    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
15059    #[test]
15060    fn delta_isolates_burst_contribution() {
15061        let mut t = SpecTelemetry::default();
15062        // "previous request": 2 rounds of k=3, accepts 3 then 1.
15063        for (kr, na) in [(3usize, 3usize), (3, 1)] {
15064            t.rounds += 1;
15065            t.drafted += kr as u64;
15066            t.accepted += na as u64;
15067            for j in 0..kr {
15068                t.pos_drafted[j] += 1;
15069            }
15070            for j in 0..na {
15071                t.pos_accepted[j] += 1;
15072            }
15073        }
15074        let before = t;
15075        // "this burst": 1 round k=3, accepts 2.
15076        t.rounds += 1;
15077        t.drafted += 3;
15078        t.accepted += 2;
15079        for j in 0..3 {
15080            t.pos_drafted[j] += 1;
15081        }
15082        for j in 0..2 {
15083            t.pos_accepted[j] += 1;
15084        }
15085        let d = t.delta_since(&before);
15086        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
15087        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
15088        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
15089        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15090    }
15091
15092    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
15093    /// aggregation invariant.
15094    #[test]
15095    fn merge_accumulates_fieldwise() {
15096        let mut agg = SpecTelemetry::default();
15097        let mut d1 = SpecTelemetry {
15098            rounds: 2,
15099            drafted: 6,
15100            accepted: 4,
15101            ..Default::default()
15102        };
15103        d1.pos_drafted[0] = 2;
15104        d1.pos_accepted[0] = 2;
15105        let mut d2 = SpecTelemetry {
15106            rounds: 1,
15107            drafted: 3,
15108            accepted: 1,
15109            ..Default::default()
15110        };
15111        d2.pos_drafted[0] = 1;
15112        d2.pos_accepted[0] = 1;
15113        d2.pos_drafted[1] = 1;
15114        agg.merge(&d1);
15115        agg.merge(&d2);
15116        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
15117        assert_eq!(agg.pos_drafted[0], 3);
15118        assert_eq!(agg.pos_accepted[0], 3);
15119        assert_eq!(agg.pos_drafted[1], 1);
15120        assert_eq!(agg.pos_accepted[1], 0);
15121    }
15122
15123    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
15124    /// public metrics surface and must never publish a u64-wrapped garbage value.
15125    #[test]
15126    fn delta_saturates_never_wraps() {
15127        let small = SpecTelemetry {
15128            rounds: 1,
15129            drafted: 2,
15130            accepted: 1,
15131            ..Default::default()
15132        };
15133        let big = SpecTelemetry {
15134            rounds: 5,
15135            drafted: 15,
15136            accepted: 9,
15137            ..Default::default()
15138        };
15139        let d = small.delta_since(&big);
15140        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
15141    }
15142}
15143
15144#[cfg(test)]
15145mod opti_fork_tests {
15146    use super::{
15147        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
15148    };
15149
15150    #[test]
15151    fn controller_threshold_and_three_miss_breaker_are_exact() {
15152        let mut policy = OptiControllerPolicy {
15153            threshold: 0.7,
15154            consecutive_misses: 0,
15155            breaker_tripped: false,
15156        };
15157        assert!(!policy.admit(0.699_999));
15158        assert!(policy.admit(0.7));
15159        assert!(!policy.resolve(false));
15160        assert!(!policy.resolve(false));
15161        assert!(policy.resolve(false));
15162        assert!(policy.breaker_tripped);
15163        assert!(!policy.admit(1.0));
15164        assert!(
15165            !policy.resolve(true),
15166            "a resolved hit cannot re-arm a tripped request"
15167        );
15168        assert!(policy.breaker_tripped);
15169    }
15170
15171    #[test]
15172    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
15173        let mut policy = OptiControllerPolicy {
15174            threshold: 0.0,
15175            consecutive_misses: 0,
15176            breaker_tripped: false,
15177        };
15178        for _ in 0..16 {
15179            assert!(policy.admit(0.0));
15180            assert!(!policy.resolve(false));
15181        }
15182        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
15183            assert!(
15184                !policy.admit(invalid),
15185                "invalid q proxy must fail closed: {invalid}"
15186            );
15187        }
15188        assert!(!policy.breaker_tripped);
15189        assert_eq!(policy.consecutive_misses, 0);
15190    }
15191
15192    #[test]
15193    fn alternating_mode_flips_by_generation_not_round_parity() {
15194        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
15195        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
15196        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
15197        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
15198    }
15199
15200    #[test]
15201    fn live_generation_cannot_be_overwritten() {
15202        let mut tracker = OptiForkGenerationTracker::default();
15203        let g0 = tracker.reserve().unwrap();
15204        let g1 = tracker.reserve().unwrap();
15205        let err = tracker.reserve().unwrap_err().to_string();
15206        assert!(
15207            err.contains("still owns generation 0"),
15208            "unexpected error: {err}"
15209        );
15210        tracker.retire(g0).unwrap();
15211        let g2 = tracker.reserve().unwrap();
15212        assert_eq!((g2.id, g2.slot), (2, 0));
15213        tracker.retire(g1).unwrap();
15214        tracker.retire(g2).unwrap();
15215    }
15216
15217    #[test]
15218    fn teardown_rejects_a_stale_generation_tag() {
15219        let mut tracker = OptiForkGenerationTracker::default();
15220        let g0 = tracker.reserve().unwrap();
15221        tracker.retire(g0).unwrap();
15222        let err = tracker.retire(g0).unwrap_err().to_string();
15223        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
15224    }
15225}
15226
15227#[cfg(test)]
15228mod draft_graph_fallback_tests {
15229    use super::DraftGraphFallback;
15230
15231    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
15232    #[test]
15233    fn flip_is_loud_once_and_memoized_after() {
15234        let mut f = DraftGraphFallback::default();
15235        let line = f
15236            .mark_greedy("out of memory")
15237            .expect("first flip must return the warn line");
15238        assert!(
15239            line.contains("WARN"),
15240            "flip line must be warn-level: {line}"
15241        );
15242        assert!(
15243            line.contains("out of memory"),
15244            "flip line must carry the reason: {line}"
15245        );
15246        assert!(f.greedy_failed());
15247        // re-marking an already-failed graph is the memoization: quiet, still failed.
15248        assert!(f.mark_greedy("out of memory").is_none());
15249        assert!(f.greedy_failed());
15250        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
15251        assert!(!f.sampled_failed());
15252        let line_s = f
15253            .mark_sampled("capture unsupported")
15254            .expect("sampled flip is its own flip");
15255        assert!(
15256            line_s.contains("sampled"),
15257            "sampled flip names itself: {line_s}"
15258        );
15259        assert!(f.mark_sampled("capture unsupported").is_none());
15260    }
15261
15262    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
15263    /// and says so exactly when there was something to reset.
15264    #[test]
15265    fn reset_on_resume_clears_flags_and_logs_once() {
15266        let mut f = DraftGraphFallback::default();
15267        // clean session: resume is silent, nothing to reset.
15268        assert!(f.reset_on_resume().is_none());
15269        f.mark_greedy("oom").unwrap();
15270        f.mark_sampled("oom").unwrap();
15271        let note = f
15272            .reset_on_resume()
15273            .expect("a set flag must produce the reset note");
15274        assert!(
15275            note.contains("greedy+sampled"),
15276            "note names what was reset: {note}"
15277        );
15278        assert!(
15279            !f.greedy_failed() && !f.sampled_failed(),
15280            "both flags cleared"
15281        );
15282        // and the NEXT failure after a reset is a fresh flip — loud again.
15283        assert!(f.mark_greedy("oom again").is_some());
15284        let note2 = f.reset_on_resume().expect("greedy-only reset");
15285        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
15286    }
15287
15288    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
15289    /// they precede a fresh capture attempt whose own failure re-flips loudly.
15290    #[test]
15291    fn shape_change_clears_are_silent() {
15292        let mut f = DraftGraphFallback::default();
15293        f.mark_greedy("oom").unwrap();
15294        f.clear_greedy();
15295        assert!(!f.greedy_failed());
15296        f.mark_sampled("oom").unwrap();
15297        f.clear_sampled();
15298        assert!(!f.sampled_failed());
15299        // after a silent clear there is nothing left for resume to report.
15300        assert!(f.reset_on_resume().is_none());
15301    }
15302}
15303
15304/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
15305///
15306/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
15307/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
15308/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
15309/// than remembered.
15310#[cfg(test)]
15311mod sampled_graph_key_tests {
15312    use super::{SampledGraphKey, debug_t_pred0};
15313
15314    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
15315    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
15316        (k.seed, k.temp_bits, k.k)
15317    }
15318
15319    fn pure_temp_key() -> SampledGraphKey {
15320        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
15321        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
15322    }
15323
15324    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
15325    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
15326    #[test]
15327    fn vendor_filters_change_the_key() {
15328        let parked = pure_temp_key();
15329        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
15330        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
15331        assert_eq!(
15332            legacy_key(&parked),
15333            legacy_key(&vendor),
15334            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
15335        );
15336        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
15337        assert!(parked.pure_temp());
15338        assert!(!vendor.pure_temp());
15339    }
15340
15341    /// Each distribution-shaping field alone is enough to drop the parked graph.
15342    #[test]
15343    fn every_filter_field_is_keyed() {
15344        let base = pure_temp_key();
15345        for (what, other) in [
15346            (
15347                "top_k",
15348                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
15349            ),
15350            (
15351                "top_p",
15352                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
15353            ),
15354            (
15355                "min_p",
15356                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
15357            ),
15358            (
15359                "penalties",
15360                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
15361            ),
15362        ] {
15363            assert_ne!(base, other, "{what} must be part of the key");
15364            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
15365            assert_eq!(
15366                legacy_key(&base),
15367                legacy_key(&other),
15368                "{what} was invisible to the pre-fix key",
15369            );
15370        }
15371    }
15372
15373    /// The baked constants stay keyed (this half was always right — regression cover for it).
15374    #[test]
15375    fn baked_constants_stay_keyed() {
15376        let base = pure_temp_key();
15377        assert_ne!(
15378            base,
15379            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
15380            "seed"
15381        );
15382        assert_ne!(
15383            base,
15384            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
15385            "temp"
15386        );
15387        assert_ne!(
15388            base,
15389            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
15390            "k"
15391        );
15392        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
15393        assert_eq!(
15394            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
15395            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
15396        );
15397    }
15398
15399    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
15400    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
15401    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
15402    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
15403    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
15404    ///
15405    /// This test is the other end of that argument, asserted here rather than remembered in a
15406    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
15407    /// would silently become the unsound thing it is documented not to be.
15408    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
15409    #[test]
15410    fn seed_alone_still_rekeys_the_draft_graph() {
15411        let parked = pure_temp_key();
15412        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
15413        assert_ne!(
15414            parked, reseeded,
15415            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
15416             decision not to compare seed rests on exactly this",
15417        );
15418        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
15419        // because of a filter difference.
15420        assert!(parked.pure_temp() && reseeded.pure_temp());
15421    }
15422
15423    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
15424    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
15425    /// agree on the regime, so a graph that survives the drop is legal to launch.
15426    #[test]
15427    fn equal_keys_agree_on_the_regime() {
15428        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
15429        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
15430        assert_eq!(a, b);
15431        assert_eq!(a.pure_temp(), b.pure_temp());
15432        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
15433        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
15434        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
15435        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
15436    }
15437
15438    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
15439    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
15440    #[test]
15441    fn debug_print_survives_the_sampled_arm() {
15442        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
15443        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
15444        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
15445        // round 0 without a pending bonus still reports last_pred, in both arms.
15446        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
15447        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
15448        // greedy keeps the real prediction it always printed.
15449        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
15450        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
15451    }
15452}