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). ON routes the step35 MTP
32/// block's draft attention through the WINDOWED device-counter family
33/// (`append_kv_quantized_dcw` + `fa_decode_dcw`, the step TP graph arc's kernels), which
34/// derives the SWA view entirely from device state (len_d, base_d, window): exactly the view
35/// offset the old capture refusal said `fa_decode_dc` could not express. BOTH draft modes
36/// switch together: eager and captured run the ONE launcher at the ONE bucket
37/// (min(cap, window)), so graph-vs-eager draft parity holds by construction (the
38/// `mtp_full_attn_dc` precedent).
39///
40/// DEFAULT ON since lane/step37-draft-graph-serving-20260830: the 20260829 lane shipped it
41/// OFF because it enabled nothing at the shipping head count (capture was structurally
42/// unreachable at heads=3); with the multi-head chain capture and the in-graph filtered
43/// sampler landed, this door is the kernel prerequisite for the captured chain on the
44/// QUALIFIED serving shape, and the exactness battery (greedy K=1..8 identity, per-K
45/// acceptance identity, seeded sampled twins) banks on the ON arm. Rollback seam:
46/// MEMRA_STEP35_DRAFT_DCW=0 restores the host-len eager arm (`mtp_step35_attn`) plus the
47/// named capture refusal, byte-for-byte the pre-lane serving; no state survives restart.
48fn step35_draft_dcw_on() -> bool {
49    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
50    *ON.get_or_init(|| std::env::var("MEMRA_STEP35_DRAFT_DCW").as_deref() != Ok("0"))
51}
52
53/// Multi-head MTP draft-chain capture door (lane/step37-draft-graph-serving-20260830,
54/// default ON — receipts in the lane RESULTS). ON lets the step-modulo prefix-replay chain
55/// (`mtp_extra` non-empty, the step37 3-head shipping shape) capture per-head single-row
56/// CUDA graphs and replay them in the exact eager launch order; the chain POLICY (head
57/// selection, prefix length, seed history) stays host-side, so graph-vs-eager drafts are
58/// bit-identical by construction. A failed capture degrades LOUDLY to the eager chain (the
59/// draft-graph WARN contract). OFF (=0) keeps the eager chain as the only multi-head path —
60/// the pre-lane serving byte-for-byte. Single-head capture is untouched by this door.
61fn mtp_chain_graph_on() -> bool {
62    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
63    *ON.get_or_init(|| std::env::var("MEMRA_MTP_CHAIN_GRAPH").as_deref() != Ok("0"))
64}
65
66/// In-graph FILTERED sampled draft door (lane/step37-draft-graph-serving-20260830, default
67/// ON — receipts in the lane RESULTS). ON widens the sampled draft-graph capture from the
68/// pure-temp regime to every truncation-filtered regime (top_k / top_p / min_p): the capture
69/// body runs `filter_stats` + `gumbel_perturb_filtered_ctr` IN-GRAPH, so the draft draws
70/// from the SAME filtered distribution the verify's accept test reconstructs (the
71/// graph-s-key exactness law, now satisfied inside the graph instead of by refusing it).
72/// Penalties stay eager either way (the history varies per round and cannot be baked).
73/// The pure-temp capture body is UNTOUCHED by this door (byte-identical to the pre-lane
74/// graph). OFF (=0) restores the pure-temp-only capture guard: filtered requests draft
75/// eager, byte-for-byte the pre-lane behavior.
76fn spec_graph_filtered_on() -> bool {
77    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_GRAPH_FILTERED").as_deref() != Ok("0"))
79}
80
81fn parse_prime_trows_width(value: Option<&str>) -> Result<usize, String> {
82    let Some(raw) = value else {
83        return Ok(8);
84    };
85    let width = raw
86        .parse::<usize>()
87        .map_err(|_| format!("MEMRA_PRIME_TROWS_T must be an integer in 2..=8, got {raw:?}"))?;
88    if !(2..=8).contains(&width) {
89        return Err(format!("MEMRA_PRIME_TROWS_T must be in 2..=8, got {width}"));
90    }
91    Ok(width)
92}
93
94#[cfg(test)]
95mod prime_trows_width_tests {
96    #[test]
97    fn width_defaults_to_eight_and_refuses_invalid_operator_values() {
98        assert_eq!(super::parse_prime_trows_width(None), Ok(8));
99        assert_eq!(super::parse_prime_trows_width(Some("2")), Ok(2));
100        assert_eq!(super::parse_prime_trows_width(Some("8")), Ok(8));
101        for invalid in ["", "1", "9", "32", "wide"] {
102            let err = super::parse_prime_trows_width(Some(invalid)).unwrap_err();
103            assert!(err.contains("MEMRA_PRIME_TROWS_T"), "{err}");
104            assert!(err.contains("2..=8"), "{err}");
105        }
106    }
107}
108
109/// One compact, anchor-bounded DSpark supervision record. `tokens[0]` is the anchor at p and
110/// `hidden` is its predecessor carrier h[p-1], matching the live NextN/DSpark pairing. Target
111/// rows p..p+gamma-1 score tokens p+1..p+gamma. They are the full-target softmax's top-k
112/// entries; `target_tail_probs[j]` is the probability mass outside those rows. All flattened
113/// target arrays are `[gamma, top_k]` in row-major order.
114pub struct DsparkAnchorRecord {
115    pub position: usize,
116    pub hidden: Vec<f32>,
117    pub tokens: Vec<u32>,
118    pub target_top_ids: Vec<u32>,
119    pub target_top_logits: Vec<f32>,
120    pub target_top_probs: Vec<f32>,
121    pub target_tail_probs: Vec<f32>,
122}
123
124fn dspark_sparse_softmax_topk(
125    logits: &[f32],
126    top_k: usize,
127    temperature: f32,
128) -> Result<(Vec<u32>, Vec<f32>, Vec<f32>, f32), Box<dyn std::error::Error>> {
129    if logits.is_empty() || top_k == 0 || top_k > logits.len() || temperature <= 0.0 {
130        return Err("invalid DSpark sparse-softmax shape or temperature".into());
131    }
132    if logits.iter().any(|value| !value.is_finite()) {
133        return Err("DSpark target logits contain a non-finite value".into());
134    }
135    let mut ranked: Vec<(u32, f32)> = logits
136        .iter()
137        .copied()
138        .enumerate()
139        .map(|(index, value)| (index as u32, value))
140        .collect();
141    let compare = |left: &(u32, f32), right: &(u32, f32)| {
142        right.1.total_cmp(&left.1).then(left.0.cmp(&right.0))
143    };
144    ranked.select_nth_unstable_by(top_k - 1, compare);
145    ranked[..top_k].sort_unstable_by(compare);
146
147    let max_logit = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
148    let inv_temperature = 1.0f64 / temperature as f64;
149    let denominator: f64 = logits
150        .iter()
151        .map(|value| (((*value - max_logit) as f64) * inv_temperature).exp())
152        .sum();
153    let ids: Vec<u32> = ranked[..top_k].iter().map(|(index, _)| *index).collect();
154    let top_logits: Vec<f32> = ranked[..top_k].iter().map(|(_, value)| *value).collect();
155    let top_probs: Vec<f32> = top_logits
156        .iter()
157        .map(|value| ((((value - max_logit) as f64) * inv_temperature).exp() / denominator) as f32)
158        .collect();
159    let top_mass: f64 = top_probs.iter().map(|value| *value as f64).sum();
160    let tail = (1.0f64 - top_mass).clamp(0.0, 1.0) as f32;
161    Ok((ids, top_logits, top_probs, tail))
162}
163
164fn flatten_dspark_rows<T>(
165    rows: Vec<Option<Vec<T>>>,
166    position: usize,
167    label: &str,
168) -> Result<Vec<T>, Box<dyn std::error::Error>> {
169    let mut flattened = Vec::new();
170    for (slot, row) in rows.into_iter().enumerate() {
171        flattened.extend(
172            row.ok_or_else(|| format!("missing DSpark {label} at {position} slot {slot}"))?,
173        );
174    }
175    Ok(flattened)
176}
177
178/// H-SEED CONVENTION (MEMRA_SPEC_HPOST=1): feed the MTP head the POST-norm hidden — trunk rows
179/// hand over `output_norm(x)` and the draft chain recurrence hands over `shared_head_norm(h_nextn)`
180/// (= final_h) — matching the reference engines: llama.cpp #24025 ("qwen35: use post-norm hidden
181/// state for MTP", t_h_nextn is taken AFTER the final norm in both trunk and MTP graphs) and
182/// SGLang's qwen3_5_mtp (spec_info.hidden_states = the target model's post-norm output). memra's
183/// historical convention (default, MTP-PLAN §A) is PRE-norm x. Draft-quality-only: exactness is
184/// the verify's job either way; acceptance arbitrates. OnceLock: read once, hot-loop safe.
185/// `MEMRA_SPEC_HEAD_ROWS=1` — batch the verify tail's LM head over its t columns instead of running
186/// it at m=1 once per column. See the call site in `decode_step_t_core_stream` for why the batched
187/// form is the same per-row arithmetic (the bf16/q8 rows twins, not cuBLASLt) and what it costs
188/// today: the head is re-streamed t times per verify pass. Default off until the byte tape says so.
189pub(crate) fn head_rows_on() -> bool {
190    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
191    crate::step37_door(&ENV, "MEMRA_SPEC_HEAD_ROWS")
192}
193
194/// The serving walk's own doors, tri-stated the same way (owner flip 2026-08-27): env forces,
195/// unset takes the step37 family default. Call sites are the t-row verify walk itself.
196pub(crate) fn spec_verify_eager_on() -> bool {
197    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
198    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_EAGER")
199}
200
201pub(crate) fn spec_verify_tcol_on() -> bool {
202    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
203    crate::step37_door(&ENV, "MEMRA_SPEC_VERIFY_TCOL")
204}
205
206/// NOT family-armed (2026-08-27): the walk's prime leaves its sub-32 TAIL chunk out of the
207/// DISTRIBUTED kv, so the server refuses before decode with "cache lengths diverged
208/// local=N distributed=floor(N/32)*32" for every prompt whose token count is not a multiple of
209/// 32 — i.e. nearly all real traffic. Isolated on the server route: defaults ERR (local=445
210/// distributed=416), MEMRA_PRIME_TROWS=0 OK. It was default-OFF before the 2026-08-27 flip and
211/// goes back to opt-in until the tail append is fixed and gated ON THE SERVER ROUTE, not just
212/// run-gen (run-gen calls decode_step_t on the whole prompt and never exercises this path — the
213/// reason a run-gen-only receipt could not see it). The GEMM prime supersedes it on this route.
214pub(crate) fn prime_trows_on() -> bool {
215    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
216    *ON.get_or_init(|| std::env::var("MEMRA_PRIME_TROWS").as_deref() == Ok("1"))
217}
218
219pub(crate) fn tcol_ffn_on() -> bool {
220    static ENV: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
221    crate::step37_door(&ENV, "MEMRA_TCOL_FFN")
222}
223
224pub(crate) fn spec_hpost() -> bool {
225    static H: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
226    *H.get_or_init(|| {
227        std::env::var("MEMRA_SPEC_HPOST")
228            .map(|v| v != "0")
229            .unwrap_or(false)
230    })
231}
232
233/// LEAN VERIFY (default ON since 2026-07-08; MEMRA_SPEC_LEAN=0 reverts — close35 lane): the verify m-scaling
234/// probe + nsys diff showed the verify t-path pays ~1.0ms/call at m=1 over eager decode on the
235/// 35B, and the kernels are NOT the cause (dev-MoE identical, kernel-time delta only +179us).
236/// The overhead is (a) ~250 extra cuMemsetD8Async/call from `e.zeros()` on buffers every kernel
237/// fully overwrites (~0.9ms host issue + ~0.35ms GPU) and (b) the t=1 FA rows dispatch (rows_v2 +
238/// combine_rows, +50us vs the eager fa_decode pair). This flag switches (a) fully-overwritten
239/// verify buffers to `e.uninit` (identical bytes: every element is written before read) and
240/// (b) t==1 verify FA to the eager `fa_decode` entry (byte-identical: kernel-check pins the
241/// rows-vs-loop identity and the per-row loop at t=1 IS fa_decode on the same q). Gates arbitrate.
242pub(crate) fn spec_lean() -> bool {
243    static L: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
244    // DEFAULT ON since 2026-07-08 (MEMRA_SPEC_LEAN=0 reverts): bit-identical (buffers fully
245    // overwritten; gates green incl maxdiff-identical run-gen) and measured +2.4% e2e p3 /
246    // +1.5% p2 at the daily 35B config. m=1 verify now costs eager-decode parity.
247    *L.get_or_init(|| {
248        std::env::var("MEMRA_SPEC_LEAN")
249            .map(|v| v != "0")
250            .unwrap_or(true)
251    })
252}
253
254/// SMALL-M BATCHED VERIFY (default ON since 2026-07-09; MEMRA_SPEC_M2=0 reverts — lane/spec-m2): extend the
255/// batched linear-attn verify arm down to t=2 and batch the MoE dev token loop over a
256/// grid.z=token axis at every verify t. The close35 m-scaling probe put the m=2 verify tier at
257/// x1.54 of m=1 (llama x1.14); the per-column linear chain (t<3) and the serial MoE dev token
258/// loop are the two launch-structure causes. Both changes are LAUNCH-STRUCTURE ONLY:
259/// (a) the batched conv's t<pad ring update is pure copies (ssm_conv_ring_rebuild from a cloned
260///     ring — the ring stores raw input columns); every arithmetic kernel is the same one the
261///     t>=3 arm already runs (matmul_decode_exact bit-identical at m=2-4, gdn_scan's internal
262///     t-loop == chained T=1 steps);
263/// (b) the MoE dev-rows twins run the serial loop's per-token warp program with tok-offset
264///     pointers (same sel/w/aq/ad bytes, same dot order, same slot-ordered FMA chain).
265/// Gates arbitrate: run-spec K=1..8 self-consistency (35B+9B), kernel-check, run-gen argmax.
266pub(crate) fn spec_m2() -> bool {
267    static M: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
268    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_M2=0 reverts): launch-structure only — t=2
269    // batched linear arm (ring-roll copies, zero new FP order) + MoE dev-rows kernels
270    // (grid.z=token, 4 launches/layer at any verify t). Acceptance bit-identical at every K;
271    // 35B p2 +3.4% / p3 +3.6%; the profitable-K plateau widens (new optimum K=3 at 223).
272    *M.get_or_init(|| {
273        std::env::var("MEMRA_SPEC_M2")
274            .map(|v| v != "0")
275            .unwrap_or(true)
276    })
277}
278pub(crate) fn spec_stream() -> bool {
279    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
280    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_STREAM").as_deref() == Ok("1"))
281}
282pub(crate) fn spec_stream_m() -> usize {
283    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
284    *M.get_or_init(|| {
285        std::env::var("MEMRA_SPEC_STREAM_M")
286            .ok()
287            .and_then(|v| v.parse().ok())
288            .unwrap_or(4)
289    })
290}
291pub(crate) fn spec_devacc() -> bool {
292    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
293    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_DEVACC").as_deref() == Ok("1"))
294}
295/// Engine-bundle slice 2 (DSF-ROUNDCOST-20260820 §1.1 host/device round trips + §2 rows 2-3),
296/// DEFAULT ON (`MEMRA_DSPARK_DEFER_READBACK=0` reverts): the dspark round's draft-chain DtoH
297/// is DEFERRED past verify dispatch and merged with the verify-argmax readback into ONE host
298/// sync (2 blocking DtoH/round -> 1). Verify embeds DEVICE tokens (`chain_d`) through the
299/// resident embed table — `embed_gather_u32_t`, bit-identical rows to the host gather by its
300/// own pinned contract. The host therefore dispatches snap + the whole verify while the DRAFT
301/// is still executing, instead of blocking ~1.7 ms on the chain and letting the device drain.
302/// Ladder arm only: the confidence policies size vt from a pre-verify head readback (their
303/// chain readback merges into that same sync instead). Exactness unchanged BY CONSTRUCTION —
304/// same tokens, same kernels, same order; E2E + accept-bank gates arbitrate.
305pub(crate) fn dspark_defer_readback_on() -> bool {
306    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
307    *ON.get_or_init(|| {
308        std::env::var("MEMRA_DSPARK_DEFER_READBACK")
309            .map(|v| v != "0")
310            .unwrap_or(true)
311    })
312}
313/// Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1, lane/dspark-engine-bundle-20260820),
314/// DEFAULT ON (`MEMRA_STATE_COPY_BATCH=0` reverts): batch the dspark round's GDN state
315/// snapshot and partial-accept restore into single `copy_batch_uniform_f32` launches
316/// instead of ~2 memcpy dispatches (+2 alloc_zeros on the snap side) per linear layer per
317/// round — measured 0.67 ms/round snap + 0.25 ms/round commit of pure dispatch on the q38
318/// route. Launch-structure only: bytes, buffers and stream order are unchanged, so
319/// acceptance and streams stay bit-identical (E2E-gated on the B1 packs).
320pub(crate) fn state_copy_batch_on() -> bool {
321    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
322    *ON.get_or_init(|| {
323        std::env::var("MEMRA_STATE_COPY_BATCH")
324            .map(|v| v != "0")
325            .unwrap_or(true)
326    })
327}
328/// Engine-bundle slice 3 + fa-execupdate slice 4c (DSF-ROUNDCOST-20260820 §5 rank 1),
329/// DEFAULT OFF — `MEMRA_DSPARK_VERIFY_GRAPH=1` opts in: per-(segment, vt) CUDA graphs
330/// for the LINEAR-layer runs, plus the full-verify single graph per (vt, rung) when a
331/// round's rows all ride one seqs rung — see [`DsparkVerifyGraphs`]. Requires the
332/// slice-2 deferred path (device tokens); the eager walk is the byte-identical fallback.
333///
334/// MEASURED disposition (box6 card0, agentic pack, 2026-08-20, both slices): exactness
335/// holds everywhere (ALL EXACT, accept lines byte-match the banks, ckpt-gate oracle
336/// green over the graph + slab-commit paths). Slice-3's AUTO_FREE launch-scan limiter
337/// (25.6 us x 16 launches ≈ 0.41 ms/round) is FIXED — the captured bodies' alloc nodes
338/// are balanced by in-graph frees (census 84/84 per segment, 1776/1776 full) so graphs
339/// instantiate USE_NODE_PRIORITY and the scan is gone. What remains at gate scale:
340/// segment graphs +0.1 tok/s over the batched-rows default (114.4 vs 114.3 x5
341/// interleaved — the linear launch overhead was only ~0.1 ms); the FULL-verify graph is
342/// NET NEGATIVE at gate scale (110.6 vs 114.2: ~14-21 (vt, rung) captures/process at
343/// 2 full-walk executions + ~2.9k-node instantiate each eat far more than the ~0.2-0.3
344/// ms/round of remaining launch overhead). The orchestration ceiling of §1.3 is spent —
345/// the fa/append recovery landed DEFAULT-ON as the batched rows arm
346/// (`dspark_fa_rows_on`), not as a graph. The serve-lifetime cell (DSF-ROUNDCOST §9,
347/// nj-ws-solo) measured the amortization: crossover K≈33 requests, steady −0.246
348/// ms/round, −1.25% session wall over 240 requests — and the graphs-serve lane wired
349/// the door into the session arm (`dspark_spec_session_burst`) as a model-owned
350/// capture pool shared across sessions. Stays opt-in pending the owner's default-ON
351/// ratification on the serve-surface battery.
352pub(crate) fn dspark_verify_graph_on() -> bool {
353    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
354    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() == Ok("1"))
355}
356/// MTP-ROUTE verify graphs, DEFAULT ON for the GDN+MoE family since 2026-08-23
357/// (`MEMRA_SPEC_VERIFY_GRAPH=0` is the kill switch, `=1` opts other families in).
358///
359/// The slice-4c capture already lived inside `qwen35_verify_tparallel` and said so in its own
360/// comment — "stream rides the qwen35moe burst, graphs ride the dspark route" — with no caller
361/// on this route. The MTP spec round is that caller.
362///
363/// WHY it is worth a default (receipts: `research/orndecode-20260822/VGRAPH.md`). With
364/// `MEMRA_SPEC_PHASE=1` this route's round reads verify-ISSUE 44-58% and verify-WAIT **0.0%**:
365/// the host is never waiting for the device, it is spending its own time launching the trunk.
366/// Replay collapses that into one graph launch and the phase all but disappears (55-62 ms ->
367/// 8-10 ms per burst).
368///
369/// MEASURED, two host generations, forced ON/OFF, balanced 4+4 boots in both orders:
370///   * current-generation host (9950X, the serving class): OFF 266.0-266.5, ON 318.8-319.5
371///     tok/s — **+19.7%**, no overlap, sub-1% spread per arm; per-round 6.9 -> 5.7 ms.
372///   * Zen 3 host: +3-9% (that rig's own clock drift is wider than the effect, so the ratio
373///     comes from per-round phase totals, which are internal to each boot).
374/// The ON arm lands at ~320 tok/s on BOTH hosts while OFF tracks host speed — the arm moves
375/// the round off the host and onto the device, which is the whole point.
376///
377/// EXACTNESS is structural (same kernels, same order) and gated anyway: a fixed-seed SAMPLED
378/// completion hashes identically ON vs OFF **and across both hosts** (`08941d5bb9762b21`),
379/// greedy seed-pinned likewise, `run-spec` K=1..8 PASS on both arms with identical acceptance
380/// at every K, kernel-check ALL GREEN.
381///
382/// SCOPE, deliberately narrow: default ON only where it was measured — the GatedDeltaNet +
383/// MoE family (`vgraph_family_default`). Qwen3.8-27B is GDN + DENSE mlp and would otherwise
384/// inherit this default unmeasured, which is the family-by-family law this repo keeps; it can
385/// opt in with `=1` once it has its own interleave. Also never armed together with
386/// ROUND-STREAM, and a round wider than the pool declines it for the eager walk.
387pub(crate) fn spec_verify_graph_env() -> Option<bool> {
388    static ON: std::sync::OnceLock<Option<bool>> = std::sync::OnceLock::new();
389    *ON.get_or_init(
390        || match std::env::var("MEMRA_SPEC_VERIFY_GRAPH").as_deref() {
391            Ok("1") => Some(true),
392            Ok("0") => Some(false),
393            _ => None,
394        },
395    )
396}
397/// SERVE-ROUTE twin of [`dspark_verify_graph_on`], DEFAULT ON — owner-ratified
398/// 2026-08-22 on the §10 serve-lifetime battery (DSF-ROUNDCOST-20260820 §10.3:
399/// crossover K=36–43, steady −0.357 ms/round, session wall −1.55..−1.65%, byte-exact
400/// 240/240 ×3 pairs, pool bounded at 8,852 MiB under `MEMRA_DSPARK_VG_MAX`). The env
401/// stays as the kill-switch: `MEMRA_DSPARK_VERIFY_GRAPH=0` restores the eager walk
402/// (byte-identical body); `MEMRA_DSPARK_VG_MAX=0` is the finer freeze valve. The BIN
403/// arm keeps its own opt-in default (`dspark_verify_graph_on`): at gate scale the
404/// capture toll is never repaid (§8 measured disposition — 14–21 captures over a
405/// 256-token run vs the serve session's thousands of rounds), and the two
406/// instruments must keep their own measured dispositions rather than share one flag.
407pub(crate) fn dspark_verify_graph_serve_on() -> bool {
408    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
409    *ON.get_or_init(|| std::env::var("MEMRA_DSPARK_VERIFY_GRAPH").as_deref() != Ok("0"))
410}
411/// Capture-count ceiling for the dspark verify-graph pool (graphs-serve lane) — the
412/// pool's memory policy STATED instead of silently unbounded. The keyspace is
413/// intrinsically finite — segment keys (run_start, vt) ≤ 16 runs x 7 windows, full
414/// keys (vt, rung, hi) ≤ 7 windows x the split-rung ladder (8 rungs at 32k ctx), ~168
415/// on the q38 export — so the default (256) never engages there; the knob is the
416/// safety valve for a future export with a wider ladder. At the ceiling the pool
417/// FREEZES: existing keys keep replaying, rounds needing a new capture run the eager
418/// walk byte-identically (round-atomic — a partial refusal would mix slab- and
419/// cols-stashed layers inside one commit). No eviction by design: destroying a live
420/// exec graph re-opens the stale-address class the indirect tables exist to close,
421/// and the bounded keyspace makes reclaim worthless.
422pub(crate) fn dspark_vg_cap() -> usize {
423    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
424    *CAP.get_or_init(|| {
425        std::env::var("MEMRA_DSPARK_VG_MAX")
426            .ok()
427            .and_then(|v| v.parse().ok())
428            .unwrap_or(256)
429    })
430}
431
432/// PROJECTED REMAINING GROWTH of the verify-graph pool, in bytes (lane/hermes-perf-fixes,
433/// 2026-08-23 — the admission accounting the "pool dwarfs spec admission reserve" finding
434/// asks for). The pool was measured at 8,852 MiB at storm-complete on the q38 export while
435/// admission's transient floor (`SPEC_SHRINK_RESERVE`) is 1.5 GiB and never charged for it:
436/// sessions admitted while the pool is cold overcommit VRAM the pool WILL hold, because the
437/// pool grows monotonically (no eviction by design) and is model-owned across sessions.
438///
439/// SELF-MEASURING, no per-model constant (generic-model law — the 8,852 MiB is a q38 number
440/// and proves nothing about another export): the debt is remaining capture slots x the
441/// MARGINAL bytes a capture adds to this device's graph mem pool.
442///
443/// MARGINAL, NOT MEAN — measured correction (box9 on-box receipt, 2026-08-23). The first
444/// version of this used the mean (`reserved / captures`) and the live serve log showed why
445/// that is wrong: with the pool's reservation flat at ~33.6 MiB across captures 1..3, the
446/// mean-based debt printed **8,556 MB, then 4,261, then 2,830** — it extrapolated capture
447/// #1's ONE-TIME shared allocation (staging buffers, stash slabs, pointer tables: sized
448/// once per pool, shared by every key) across all 256 slots. An 8.5 GB phantom reserve at
449/// boot can refuse admissions that would have fit, which is a worse defect than the
450/// under-charge this accounting exists to remove. The marginal reading prices what an
451/// ADDITIONAL key actually costs: two observations `(captures, reserved)` give
452/// `(r1 - r0) / (c1 - c0)`, which is ~0 on an export whose pool does not grow per key and
453/// tracks real growth on one that does.
454///
455/// BOOTSTRAP (only one observation so far, so growth is unmeasurable): reserve one more
456/// pool's worth — `min(remaining x mean, reserved)`. "We have measured `reserved` bytes for
457/// `captures` keys; until growth is measurable, assume at most a doubling" is fail-safe in
458/// the same direction as the old rule without the 255x extrapolation.
459///
460/// Before the FIRST capture the debt is 0 (a single capture lands well inside the existing
461/// 1.5 GiB floor). `cap` is the intrinsic freeze ceiling (`MEMRA_DSPARK_VG_MAX`; =0 freeze
462/// valve => the pool cannot grow => debt 0); at or past the cap the pool FREEZES, so the
463/// debt is 0 there too.
464pub fn dspark_vg_debt_projection(
465    captures: usize,
466    cap: usize,
467    reserved_bytes: usize,
468    prev: Option<(usize, usize)>,
469) -> usize {
470    if captures == 0 || cap == 0 {
471        return 0;
472    }
473    let remaining = cap.saturating_sub(captures);
474    if remaining == 0 {
475        return 0;
476    }
477    match prev {
478        // marginal growth between two observations of the same pool
479        Some((c0, r0)) if captures > c0 => {
480            let marginal = reserved_bytes.saturating_sub(r0) / (captures - c0);
481            remaining.saturating_mul(marginal)
482        }
483        // bootstrap: at most one more pool's worth
484        _ => remaining
485            .saturating_mul(reserved_bytes / captures)
486            .min(reserved_bytes),
487    }
488}
489/// Engine-bundle slice 4 (fa-execupdate lane, DSF-ROUNDCOST-20260820 §6 close: "the
490/// residual gap lives in the FULL-ATTENTION per-row section"), DEFAULT ON —
491/// `MEMRA_DSPARK_FA_ROWS=0` reverts to the per-row loop: when every row of a verify
492/// round takes the v4-seqs arm on ONE `fa_split_keys` rung (the straddle law, evaluated
493/// at the round's first and last t_kv — both eligibility gates are intervals in t_kv),
494/// the qwen35 t-parallel verify's per-row KV-append + fa-decode loop collapses into the
495/// z-batched serving twins: ONE `append_quantize_kv_q8_0_q5_1_seqs` + ONE
496/// `fa_decode_vec_q_seqs_v4` + ONE combine per full-attention layer, replacing
497/// T x (4 dtod row copies + append + 3 memsets + main + combine) launches. Bytes are
498/// pinned by the batched-tick increment-2 kernel-check (seqs-vs-per-seq-loop bit
499/// identity: per-row T_kv derives in-kernel from pos_seq[z]; splits >= ns_eff write the
500/// empty partial the combine never reads, so the shared n_splits_max stride changes no
501/// bytes) and re-gated e2e by this lane's battery.
502pub(crate) fn dspark_fa_rows_on() -> bool {
503    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
504    *ON.get_or_init(|| {
505        std::env::var("MEMRA_DSPARK_FA_ROWS")
506            .map(|v| v != "0")
507            .unwrap_or(true)
508    })
509}
510
511/// `t_pred0` for the `MEMRA_DEBUG_SPEC` per-round print, sampled-safe.
512///
513/// `generate_spec_inner2` fills its `preds` vector ONLY on the greedy path (`if !sampled`), and
514/// the per-round debug print was the sole consumer in the sampled arm: `t_pred(0)` survives round
515/// 0 (`base == 0` returns `last_pred`) and from round 1 (`base == 1`, a pending bonus) indexes an
516/// EMPTY vector — `index out of bounds: the len is 0 but the index is 0`, in the GPU worker
517/// thread, which then respawns and reloads weights while the request dies. So any sampled spec
518/// request longer than one round used to kill the worker whenever `MEMRA_DEBUG_SPEC` was set:
519/// the flag crashed precisely the regime it exists to investigate.
520///
521/// Fixed at the print site, not inside the closure, so the greedy accept walk keeps its strict
522/// indexing (an out-of-range pred there is a real bug and must still be loud).
523fn debug_t_pred0(sampled: bool, base: usize, last_pred: u32, preds: &[u32]) -> String {
524    if base == 0 {
525        return last_pred.to_string();
526    }
527    match preds.get(base - 1) {
528        Some(p) => p.to_string(),
529        // sampled: the greedy per-column argmax was never run for this round.
530        None => {
531            debug_assert!(
532                sampled,
533                "greedy spec: preds[{}] missing at base {base}",
534                base - 1
535            );
536            "n/a".to_string()
537        }
538    }
539}
540
541/// `MEMRA_SKEY_PROBE=1` — sampled-draft-graph key probe (lane/graph-s-key-exactness-20260819).
542///
543/// Reports, per burst and per round, which draft chain the sampled arm chose and under which
544/// filter regime, plus the ONE observable that separates a legal filtered draft from a stale
545/// pure-temp graph replayed under filters: an accept test whose gathered `q` is exactly 0.
546/// A draft token sampled from the FILTERED softmax can never gather q=0 (it was drawn from the
547/// kept set), so `q=0` in the verify means the draft came from a distribution the verify does
548/// not believe in — and `u * 0 < p` then accepts it unconditionally.
549///
550/// Its own env var, deliberately NOT `MEMRA_DEBUG_SPEC`: that flag panicked the GPU worker on
551/// any sampled spec request past round 0 until this lane fixed it (§2 of the bank note).
552pub(crate) fn skey_probe() -> bool {
553    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
554    *ON.get_or_init(|| std::env::var("MEMRA_SKEY_PROBE").as_deref() == Ok("1"))
555}
556
557/// GRAMMAR HOOK for constrained spec decode (lane/constrained-full, 2026-08-03). The engine
558/// stays llguidance-agnostic: the server adapts its per-session grammar state behind this
559/// trait. CONTRACT (the verify-side truncation rule — token-identical to constrained plain
560/// greedy decode): the exactness walk runs UNMASKED first; the hook then (a) truncates
561/// acceptance at the first grammar-illegal accepted token, and (b) when the truncation fired
562/// or the bonus is illegal, the engine recomputes that slot as the MASKED argmax of the
563/// target's own verify column (an unmasked argmax that is grammar-legal IS the masked argmax
564/// — masking only removes tokens — so the common case pays nothing). `consume` advances the
565/// state with each EMITTED token in order; EOS handling is the implementor's job (skip).
566pub trait SpecConstraint {
567    /// -inf the current state's banned ids on a HOST logits row (prompt-tail / init-feed
568    /// masked argmax).
569    fn mask_logits(&mut self, logits: &mut [f32]) -> Result<(), String>;
570    /// Packed 32-bit bitset words of the CURRENT state's allowed set (device-mask form).
571    fn mask_words(&mut self) -> Result<Vec<u32>, String>;
572    /// Is `tok` consumable in the CURRENT state?
573    fn is_allowed(&mut self, tok: u32) -> Result<bool, String>;
574    /// Advance the state with an emitted token.
575    fn consume(&mut self, tok: u32) -> Result<(), String>;
576
577    // --- DRAFT-SIDE MASKING (lane/draft-mask, 2026-08-04) ---
578    // The drafter proposed grammar-illegal tokens under tight schemas, so verify-side
579    // truncation cut nearly every round (measured acceptance 0.467-0.513 tight vs 0.62-0.82
580    // loose, research/constrained-full-20260803). These three methods let the engine mask the
581    // DRAFT model's own sampling with the grammar's legal set, so proposals are legal by
582    // construction. The state they walk is a SPECULATIVE CLONE of the session matcher — the
583    // real state is advanced only by `consume` (emitted tokens), so verify-side truncation
584    // stays the correctness backstop and the emitted stream is unchanged by construction
585    // (an accepted draft is the target's unmasked argmax AND grammar-legal, hence the masked
586    // argmax; a cut slot is recomputed as the masked argmax either way).
587    // Default impls = feature OFF (pre-lane behaviour: unmasked drafts).
588
589    /// Is draft-side masking available on this hook? Probed ONCE per burst, before the draft
590    /// graph is captured (the mask is an in-graph node — its presence is a capture-time shape).
591    fn draft_mask_enabled(&self) -> bool {
592        false
593    }
594    /// Start a draft chain: clone the CURRENT (committed) grammar state into the speculative
595    /// slot. Called once per spec round, before the first draft position.
596    fn draft_begin(&mut self) -> Result<(), String> {
597        Ok(())
598    }
599    /// Packed 32-bit bitset words of the SPECULATIVE state's allowed set (target-vocab ids),
600    /// for the draft position about to be sampled. `None` = draft masking off (no-op).
601    fn draft_mask_words(&mut self) -> Result<Option<Vec<u32>>, String> {
602        Ok(None)
603    }
604    /// Advance the SPECULATIVE state with a PROPOSED draft token. `false` = the chain cannot
605    /// continue (EOS proposed, or an unmasked position proposed something illegal) — the
606    /// engine stops drafting; the token already pushed still goes through verify.
607    fn draft_advance(&mut self, _tok: u32) -> Result<bool, String> {
608        Ok(false)
609    }
610}
611
612/// DRAFT-MASK UPLOAD (lane/draft-mask): pull the speculative state's allowed set (TARGET-id
613/// space) from the hook, project it into the DRAFT head's vocab space, and upload it into the
614/// stable device buffer the draft chain reads. Returns false when the chain must stop drafting:
615/// the hook handed out no mask, or NO draft-vocab row is grammar-legal at this position (a
616/// trimmed FR-Spec head genuinely cannot propose a legal token there — masking it would leave
617/// a fully-banned row whose argmax is meaningless, so the round drafts fewer tokens and the
618/// verify emits the masked argmax as usual).
619fn upload_draft_mask(
620    e: &Engine,
621    c: &mut dyn SpecConstraint,
622    dst: &mut CudaSlice<u32>,
623    d2t: Option<&Vec<u32>>,
624    d_vocab: usize,
625    words: usize,
626) -> Result<bool, Box<dyn std::error::Error>> {
627    let Some(tw) = c
628        .draft_mask_words()
629        .map_err(|e2| format!("constraint: {e2}"))?
630    else {
631        return Ok(false);
632    };
633    let bit = |t: usize| -> bool {
634        let w = t >> 5;
635        w < tw.len() && (tw[w] >> (t & 31)) & 1 == 1
636    };
637    let mut buf = vec![0u32; words];
638    match d2t {
639        // TRIMMED draft head: row i proposes target id d2t[i] — permute the mask accordingly.
640        Some(map) => {
641            for (i, &t) in map.iter().enumerate().take(d_vocab) {
642                if bit(t as usize) {
643                    buf[i >> 5] |= 1u32 << (i & 31);
644                }
645            }
646        }
647        // UNTRIMMED: draft ids ARE target ids; the packed words transfer verbatim (a short
648        // mask leaves the padded tail zeroed == banned, same rule as constrained::apply_mask).
649        None => {
650            let n = tw.len().min(words);
651            buf[..n].copy_from_slice(&tw[..n]);
652        }
653    }
654    if buf.iter().all(|w| *w == 0) {
655        return Ok(false);
656    }
657    e.htod_u32_into(dst, &buf)?;
658    Ok(true)
659}
660
661/// Keep the full token-embedding table in host memory and upload only the rows needed by each
662/// MTP/verify step. This is an exact memory-capacity seam for very large BF16 vocab tables: host
663/// gather expands the same source bits to f32, and only O(T*n_embd) bytes cross PCIe per step.
664/// CUDA-graph/round-stream draft paths require device token ids and therefore stay disabled.
665pub(crate) fn spec_host_embd() -> bool {
666    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
667    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_HOST_EMBD").as_deref() == Ok("1"))
668}
669
670/// VERIFY-TIER TRUNK LAUNCH-FUSION (default ON since 2026-07-09; MEMRA_SPEC_FUSED_T=0 reverts — lane/close35b): extend
671/// the t=1 fused2/fused3 Q8_0 trunk launches to the batched verify tier (t=2-4, the K=1..3
672/// verify shapes). At t>1 the trunk pairs/triples (35B wqkv+wqkv_gate, wq/wk/wv,
673/// gate_shexp+up_shexp) each run a separate `matmul_decode_exact` — one q8_1 re-quantize of the
674/// SAME activation plus one _b2/_b4 launch per tensor. The fused twins share ONE quantize and
675/// ONE launch per group; per (tensor,token,row) the kernel body is q8_0_mmvq_batched verbatim
676/// with the identical row mapping -> BIT-IDENTICAL by construction (kernel-check pins it,
677/// run-spec K=1..8 + acceptance identity arbitrate e2e).
678pub(crate) fn spec_fused_t() -> bool {
679    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
680    // DEFAULT ON since 2026-07-09 (MEMRA_SPEC_FUSED_T=0 reverts): verify t=2-4 trunk launch-fusion
681    // (fused2/fused3 Q8_0 batched twins, bit-identical by construction — m=1 block-offset split on
682    // the batched body). m=2 marginal token 2117->1762us; 35B daily: p3 +3.7% (crosses llama), p2 +5%.
683    *F.get_or_init(|| {
684        std::env::var("MEMRA_SPEC_FUSED_T")
685            .map(|v| v != "0")
686            .unwrap_or(true)
687    })
688}
689
690/// zeros/uninit switch for verify-path buffers that are FULLY OVERWRITTEN before any read.
691/// Only call this on such buffers — the lean contract is "identical bytes by construction".
692/// TOKEN-ID GUARD for every id that reaches an embed gather (#87 family).
693///
694/// A device argmax seeds its running index with 0x7FFFFFFF and replaces it only through
695/// comparisons, all of which are FALSE against NaN. An all-NaN logits row therefore returns
696/// the sentinel, and the next thing done with a token id is `embed_row(id)` — table +
697/// ~4.6 TB, never mapped, an MMU fault that kills the CUDA context for the whole process
698/// (research/pp2spec-crash-20260807). The draft chain and the GREEDY verify walk already
699/// trap this; the SAMPLED verify bonus, the boundary sampler and the replay arm's last_pred
700/// did not, which is why the recoverable fault on the greedy instrument is a TERMINAL one on
701/// the vendor-default sampled shape we actually serve.
702pub(crate) fn guard_vocab_token(
703    tok: u32,
704    n_vocab: usize,
705    what: &str,
706) -> Result<u32, Box<dyn std::error::Error>> {
707    if (tok as usize) >= n_vocab {
708        return Err(format!(
709            "{what}: token id 0x{tok:08x} >= n_vocab {n_vocab} — an all-NaN logits row left \
710             the device argmax's init sentinel in place; refusing to dereference the embed \
711             row (#87 trap)"
712        )
713        .into());
714    }
715    Ok(tok)
716}
717
718/// SPEC NaN-ORIGIN SCAN (`MEMRA_SPEC_NAN_SCAN=1`, DEFAULT OFF, diagnostic only).
719///
720/// The `#87` trap reports an all-NaN VERIFY logits column, which says the poison reached the
721/// head but not where it entered. With the scan armed the verify walk syncs and reads back
722/// every layer's output, so the FIRST layer whose residual carries a NaN names itself with the
723/// round's row and position. Off by default and never on a serving path: it costs one host
724/// sync + one `t*n_embd` D2H per layer, and the syncs change scheduling (so a run that stops
725/// reproducing under the scan is itself a datum, not an all-clear).
726///
727/// Rollback seam: unset `MEMRA_SPEC_NAN_SCAN` (or set it to 0). Every call site is behind
728/// `spec_nan_scan()`, so the default path keeps the exact launch sequence it had.
729pub(crate) fn spec_nan_scan() -> bool {
730    spec_nan_scan_level() > 0
731}
732
733/// `MEMRA_SPEC_NAN_SCAN` as a LEVEL, not a boolean. `1` scans each layer's residual, which
734/// names the layer. `2` also scans INSIDE the t-column layer body — the per-column attention
735/// output, the deferred-column o-proj/fa2 join, the post-attention norm and the routed-MoE
736/// output — because "layer 20 poisons row 0" does not say whether the attention or the routed
737/// MoE produced it, and those are different bugs with different fixes.
738pub(crate) fn spec_nan_scan_level() -> u8 {
739    static LVL: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
740    *LVL.get_or_init(|| match std::env::var("MEMRA_SPEC_NAN_SCAN").as_deref() {
741        Ok("1") => 1,
742        Ok("2") => 2,
743        _ => 0,
744    })
745}
746
747/// Read back `[rows, cols]` and fail with the first NaN's coordinates. `what` names the
748/// producer (layer index, walk arm) so the error line is the localization.
749/// VERIFY-ARM RECEIPT (rides `MEMRA_SPEC_NAN_SCAN>=1`, bounded to 200 lines).
750///
751/// Names, per trunk layer, WHICH attention arm the t-column walk actually took. This exists
752/// because the level-1 residual scan below sat only on the non-fused tail: the fused
753/// rope+append+fa arm ends in `continue`, so every layer that fused was NEVER SCANNED and
754/// silently read as "clean". A poisoned residual therefore first reported at the next
755/// non-fused layer, which is how "layer 20 creates the poison" could be true of the scan and
756/// false of the engine. Also carries the row-table lookup counter, so "the fused path never
757/// ran" is distinguishable from "it ran and was innocent".
758/// KV-PLANE SCAN (`MEMRA_KV_PLANE_SCAN=1`, DEFAULT OFF, diagnostic only).
759///
760/// Reads back the STAGED rows of a layer's distributed K/V planes and reports the first row
761/// whose quantization scale is not finite. No kernel required: q8_0 blocks are
762/// `[half d][32 x i8]` and q5_1 blocks carry `half d` then `half m`, so the fp16 scale at the
763/// head of each block is host-checkable straight out of the byte plane.
764///
765/// It exists because the level-2 bad-row bitmap says EVERY verify row is non-finite at a
766/// global-attention layer's join, and row r attends a strict superset of row r-1's keys: that
767/// implicates the shared KV history those rows walk, not per-column staging. "The attention
768/// output is NaN" and "the KV history it attends is already NaN" are different bugs with
769/// different owners, and nothing measured so far separates them. A first-corrupt-row index
770/// also dates the corruption against the prime/decode boundary.
771///
772/// Bounded hard: only layers whose geometry has NO window (the global planes), only the first
773/// `MEMRA_KV_PLANE_SCAN_ROUNDS` verify rounds of a process (default 2), and it copies only
774/// `[0, staged_len)`, which is ~1.6 MB at the 1480-token repro rather than the 262144-row
775/// provision. It still syncs per layer, so it is never a serving or a measured-perf arm.
776pub(crate) fn kv_plane_scan_on() -> bool {
777    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
778    *ON.get_or_init(|| std::env::var("MEMRA_KV_PLANE_SCAN").as_deref() == Ok("1"))
779}
780
781fn kv_plane_scan_rounds() -> usize {
782    static R: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
783    *R.get_or_init(|| {
784        std::env::var("MEMRA_KV_PLANE_SCAN_ROUNDS")
785            .ok()
786            .and_then(|v| v.parse().ok())
787            .unwrap_or(2)
788    })
789}
790
791/// First non-finite fp16 block scale in `bytes`, as (block index, raw u16), scanning one
792/// scale every `stride` bytes. Returns None when every block scale is finite.
793fn first_bad_scale(bytes: &[u8], stride: usize) -> Option<(usize, u16)> {
794    if stride == 0 {
795        return None;
796    }
797    for (i, blk) in bytes.chunks_exact(stride).enumerate() {
798        let raw = u16::from_le_bytes([blk[0], blk[1]]);
799        if half_is_non_finite(raw) {
800            return Some((i, raw));
801        }
802    }
803    None
804}
805
806/// IEEE binary16: exponent all ones is Inf or NaN, whatever the mantissa says.
807fn half_is_non_finite(raw: u16) -> bool {
808    (raw & 0x7C00) == 0x7C00
809}
810
811/// Scan one layer's staged K/V planes for a non-finite quantization scale. Returns the
812/// receipt line, or None when the layer is out of scope or every scale is finite.
813pub(crate) fn scan_kv_plane(
814    e: &crate::Engine,
815    distributed: &memra_kv::ResidentTpKvCache,
816    il: usize,
817    pos0: usize,
818) -> Result<(), Box<dyn std::error::Error>> {
819    // One "round" is one pos0, not one layer: the walk visits 45 layers per verify. The
820    // default of 2 rounds is for a fault that shows up immediately; the step37 repro does not
821    // fire until rep 3 or later, i.e. round ~60 of the process, so that arm MUST raise
822    // MEMRA_KV_PLANE_SCAN_ROUNDS or it will scan only the two rounds that were never going to
823    // be poisoned and report a clean history it never looked at.
824    static ROUNDS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
825    static LAST_POS: std::sync::atomic::AtomicUsize =
826        std::sync::atomic::AtomicUsize::new(usize::MAX);
827    if LAST_POS.swap(pos0, std::sync::atomic::Ordering::Relaxed) != pos0 {
828        ROUNDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
829    }
830    if ROUNDS.load(std::sync::atomic::Ordering::Relaxed) > kv_plane_scan_rounds() {
831        return Ok(());
832    }
833    let staged = distributed.staged_len();
834    if staged == 0 {
835        return Ok(());
836    }
837    // ENGAGEMENT RECEIPT. This scan prints only on corruption, so `kvbad=0` in a cell would
838    // read the same whether the history was clean or the scan never ran once. Bounded so a
839    // 45-layer walk cannot flood the log.
840    static SEEN: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
841    let seen = SEEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
842    let (ktb, vtb) = (distributed.k_tok_bytes(), distributed.v_tok_bytes());
843    if seen < 4 {
844        eprintln!(
845            "[kv-plane] engaged #{seen} layer {il} pos0={pos0} staged={staged} \
846             ktok={ktb} vtok={vtb} (scan armed; a corrupt plane prints its own line)"
847        );
848    }
849    for rank in 0..distributed.ranks().len() {
850        let Some(rc) = distributed.rank(rank) else {
851            continue;
852        };
853        // q8_0 K blocks are [half d][32 x i8] = 34B; q5_1 V blocks lead with half d then half m.
854        let kbytes = e.dtoh_u8_view(&rc.k().slice(0..staged * ktb))?;
855        let vbytes = e.dtoh_u8_view(&rc.v().slice(0..staged * vtb))?;
856        let kbad = first_bad_scale(&kbytes, 34);
857        let vbad = first_bad_scale(&vbytes, 24);
858        if kbad.is_some() || vbad.is_some() {
859            let row = |b: Option<(usize, u16)>, tok: usize| {
860                b.map(|(i, raw)| format!("blk {i} (row {}) raw={raw:#06x}", i * 34 / tok.max(1)))
861                    .unwrap_or_else(|| "clean".into())
862            };
863            eprintln!(
864                "[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",
865                row(kbad, ktb),
866                row(vbad, vtb)
867            );
868            return Ok(());
869        }
870    }
871    Ok(())
872}
873
874pub(crate) fn verify_arm_receipt(
875    arm: &str,
876    il: usize,
877    pos0: usize,
878    t: usize,
879    staged: Option<usize>,
880) {
881    static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
882    if N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) >= 200 {
883        return;
884    }
885    eprintln!(
886        "[verify-arm] layer {il} arm={arm} pos0={pos0} t={t} staged_len={} rows_tab_lookups={}",
887        staged.map(|v| v as i64).unwrap_or(-1),
888        crate::tp::ROWS_TAB_ENGAGED.load(std::sync::atomic::Ordering::Relaxed)
889    );
890}
891
892pub(crate) fn nan_scan_rows(
893    e: &Engine,
894    buf: &CudaSlice<f32>,
895    rows: usize,
896    cols: usize,
897    what: &str,
898) -> Result<(), Box<dyn std::error::Error>> {
899    // The readback is also the ATTRIBUTION point for an asynchronous fault: a
900    // CUDA_ERROR_ILLEGAL_ADDRESS raised by any launch since the previous scan surfaces on this
901    // sync, and the bare DriverError names nothing. Wrapping it with `what` turns "the process
902    // died somewhere" into "it died at or before this layer, on this row, at this position".
903    let host = e.dtoh(buf).map_err(|err| -> Box<dyn std::error::Error> {
904        format!(
905            "spec nan-scan: sync at {what} FAILED: {err} — the fault is at or before \
906                     this point in the walk"
907        )
908        .into()
909    })?;
910    if host.len() < rows * cols {
911        return Err(format!(
912            "nan-scan {what}: buffer holds {} < {rows}x{cols}",
913            host.len()
914        )
915        .into());
916    }
917    // SCAN EVERY ROW BEFORE REPORTING. A first-hit return says "row 0 is bad" and leaves the
918    // other rows UNEXAMINED, which is exactly the bit that discriminates the two mechanisms: in
919    // the t-column verify, row 0 attends keys [0..p+1) and row 1 attends [0..p+2), a strict
920    // superset, so poison in the SHARED KV history must appear in BOTH rows, while poison in
921    // per-column staging can appear in one. Report the whole map.
922    let mut per_row: Vec<usize> = Vec::with_capacity(rows);
923    let mut first_bad: Option<(usize, usize)> = None;
924    for r in 0..rows {
925        let row = &host[r * cols..(r + 1) * cols];
926        let bad = row.iter().filter(|v| !v.is_finite()).count();
927        per_row.push(bad);
928        if bad > 0 && first_bad.is_none() {
929            first_bad = Some((r, row.iter().position(|v| !v.is_finite()).unwrap_or(0)));
930        }
931    }
932    if let Some((r0, c0)) = first_bad {
933        let map: String = per_row
934            .iter()
935            .map(|&b| if b == 0 { '.' } else { 'X' })
936            .collect();
937        return Err(format!(
938            "spec nan-scan: {what} produced non-finite values — rows[{rows}] map={map} \
939             counts={per_row:?} of {cols} each; first at row {r0} element {c0}. Both rows bad \
940             implicates shared state (the KV history this layer reads); one row bad implicates \
941             per-column staging."
942        )
943        .into());
944    }
945    Ok(())
946}
947
948fn vbuf(e: &Engine, n: usize) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
949    if spec_lean() { e.uninit(n) } else { e.zeros(n) }
950}
951
952/// Scratch KV for the MTP block (one full-attn layer).
953///
954/// PERSISTENT MODE (default, 2026-07-03 — the acceptance lever): sized cap = max_ctx and kept in
955/// sync with the COMMITTED sequence — slot p holds the MTP block's K/V for committed token p
956/// (roped p+1, the chain's rope convention), so the draft chain's self-attention sees the FULL
957/// committed history instead of only the current round's 1..K+1 chain tokens (the reference
958/// engine's "mtp_update" design). Entries come from two sources:
959///   - chain appends: accepted positions KEEP their chain-computed entries (embedding exact,
960///     hidden chain-approximate — the reference engine accepts the same);
961///   - `mtp_kv_fill` batches: prompt positions + the last-draft position on full accept, computed
962///     from EXACT trunk hiddens (K/V-only MTP-block pass, no attention/FFN/lm_head).
963/// Rejected drafts / p-min extras / pseudo-seed appends are all discarded by the round-start
964/// `set_len` truncation (the KvLayer len mechanism — §C rollback for the draft side).
965/// Multi-turn spec-decode session (2026-07-05): trunk Cache + persistent MTP draft scratch +
966/// the committed token list, alive across generate_spec_session calls. Turn N+1 primes ONLY its
967/// suffix (chunked continuation prime over the quantized past) and mtp_kv_fill's its suffix rows,
968/// then the round loop runs unchanged. `last_h` carries the pre-output_norm hidden of the last
969/// committed row across turns (the predecessor-pairing seed + fill anchor).
970/// Per-request sampling config for the sampled-spec serve path.
971#[derive(Clone, Copy, Debug)]
972pub struct SpecSampling {
973    pub temp: f32,
974    pub seed: u64,
975    pub top_k: i32,            // 0 = off
976    pub top_p: f32,            // 1.0 = off
977    pub min_p: f32,            // 0.0 = off
978    pub penalty_last_n: usize, // 0 = penalties off
979    pub penalty_repeat: f32,
980    pub penalty_freq: f32,
981    pub penalty_present: f32,
982}
983
984impl SpecSampling {
985    /// Non-identity penalties requested — THE `pen_on` predicate (one definition; the
986    /// same group-off rule `SamplerIdentity::of` canonicalizes: a window with neutral
987    /// coefficients is penalties-absent). Both spec routes and the dspark accept walk
988    /// key their penalty arms off this.
989    pub fn pen_on(&self) -> bool {
990        self.penalty_last_n > 0
991            && (self.penalty_repeat != 1.0
992                || self.penalty_freq != 0.0
993                || self.penalty_present != 0.0)
994    }
995}
996
997/// Host Philox4x32-10 uniform in (0,1) — mirrors spec_sample.cu's `philox4`/`u01` with the
998/// ctr_lo tag 0xFFFF_FFFE, so the host accept-test stream never collides with any device
999/// sampling event (device Gumbel uses (i>>2, stream_pos); device residual uses 0xFFFF_FFFD).
1000/// One value per (seed, ctr) EVENT; callers own the counter discipline. Extracted verbatim
1001/// from generate_spec_inner2's closure for the dspark sampled-admission walk (the two paths
1002/// MUST consume the identical stream construction — two ad-hoc Philox copies drifting apart
1003/// is a distributional bug, not a style problem).
1004pub(crate) fn host_u01(seed: u64, ctr: u32) -> f32 {
1005    let (m0, m1) = (0xD2511F53u32, 0xCD9E8D57u32);
1006    let (mut c0, mut c1, mut c2, mut c3) = (0xFFFF_FFFEu32, ctr, 0u32, 0u32);
1007    let (mut k0, mut k1) = ((seed & 0xFFFF_FFFF) as u32, (seed >> 32) as u32);
1008    for _ in 0..10 {
1009        let (h0, l0) = (((m0 as u64 * c0 as u64) >> 32) as u32, m0.wrapping_mul(c0));
1010        let (h1, l1) = (((m1 as u64 * c2 as u64) >> 32) as u32, m1.wrapping_mul(c2));
1011        let (n0, n1, n2, n3) = (h1 ^ c1 ^ k0, l1, h0 ^ c3 ^ k1, l0);
1012        c0 = n0;
1013        c1 = n1;
1014        c2 = n2;
1015        c3 = n3;
1016        k0 = k0.wrapping_add(0x9E3779B9);
1017        k1 = k1.wrapping_add(0xBB67AE85);
1018    }
1019    (c0 as f32 + 1.0) * (1.0 / 4294967296.0)
1020}
1021
1022/// Tracked draft positions for [`SpecTelemetry`] (serve K defaults to 3; the run-spec gate
1023/// sweeps K=1..8, and MEMRA_SPEC_CAPMAX defaults to 7 — 8 covers every tuned config).
1024pub const SPEC_TELEM_POS: usize = 8;
1025
1026/// Always-on per-draft-position acceptance telemetry (lane/accept-telemetry, 2026-08-05 —
1027/// the llama.cpp #26389 / vLLM spec-decode counter schema, upstream-sweeps 2026-08-05).
1028/// Lives on the [`SpecSession`] and accumulates across bursts; the serve worker diffs a
1029/// stashed copy per burst for its per-model /metrics aggregation and per-request usage.
1030/// Same normalization as the `[spec-stats]` line: p-min-discarded chain tokens are counted
1031/// in NEITHER drafted nor accepted.
1032#[derive(Clone, Copy, Default, Debug)]
1033pub struct SpecTelemetry {
1034    /// verify rounds completed (a round-stream burst counts each of its M rounds).
1035    pub rounds: u64,
1036    /// tokens drafted / accepted across all rounds.
1037    pub drafted: u64,
1038    pub accepted: u64,
1039    /// how often draft position j (0-based within a round's chain) was offered / accepted.
1040    /// Positions >= SPEC_TELEM_POS are untracked (totals still count them). The opt-in
1041    /// round-stream arm (MEMRA_SPEC_STREAM=1) reads back only totals, so under it these
1042    /// arrays cover the standard-path rounds only and their sums may undercount the totals.
1043    pub pos_drafted: [u64; SPEC_TELEM_POS],
1044    pub pos_accepted: [u64; SPEC_TELEM_POS],
1045}
1046
1047impl SpecTelemetry {
1048    /// Fieldwise `self - prev` — the worker's per-burst delta off a copy stashed before the
1049    /// burst call. Saturating: a caller diffing against the wrong snapshot gets zeros, not
1050    /// a wrapped counter.
1051    pub fn delta_since(&self, prev: &SpecTelemetry) -> SpecTelemetry {
1052        let mut d = SpecTelemetry {
1053            rounds: self.rounds.saturating_sub(prev.rounds),
1054            drafted: self.drafted.saturating_sub(prev.drafted),
1055            accepted: self.accepted.saturating_sub(prev.accepted),
1056            ..Default::default()
1057        };
1058        for j in 0..SPEC_TELEM_POS {
1059            d.pos_drafted[j] = self.pos_drafted[j].saturating_sub(prev.pos_drafted[j]);
1060            d.pos_accepted[j] = self.pos_accepted[j].saturating_sub(prev.pos_accepted[j]);
1061        }
1062        d
1063    }
1064    /// Fieldwise `self += d` — the worker's per-model aggregation.
1065    pub fn merge(&mut self, d: &SpecTelemetry) {
1066        self.rounds += d.rounds;
1067        self.drafted += d.drafted;
1068        self.accepted += d.accepted;
1069        for j in 0..SPEC_TELEM_POS {
1070            self.pos_drafted[j] += d.pos_drafted[j];
1071            self.pos_accepted[j] += d.pos_accepted[j];
1072        }
1073    }
1074
1075    /// Mean accepted draft-prefix length per verify round (tau).
1076    pub fn tau(&self) -> f64 {
1077        if self.rounds > 0 {
1078            self.accepted as f64 / self.rounds as f64
1079        } else {
1080            0.0
1081        }
1082    }
1083}
1084
1085/// Session-lifetime atomic acceptance counters. The verifier records only after the greedy or
1086/// rejection-sampling walk has resolved on the host, so these relaxed increments add no GPU
1087/// launch, synchronization, allocation, or ordering dependency to the numeric path.
1088struct SpecTelemetryCounters {
1089    rounds: AtomicU64,
1090    drafted: AtomicU64,
1091    accepted: AtomicU64,
1092    pos_drafted: [AtomicU64; SPEC_TELEM_POS],
1093    pos_accepted: [AtomicU64; SPEC_TELEM_POS],
1094}
1095
1096impl Default for SpecTelemetryCounters {
1097    fn default() -> Self {
1098        Self {
1099            rounds: AtomicU64::new(0),
1100            drafted: AtomicU64::new(0),
1101            accepted: AtomicU64::new(0),
1102            pos_drafted: std::array::from_fn(|_| AtomicU64::new(0)),
1103            pos_accepted: std::array::from_fn(|_| AtomicU64::new(0)),
1104        }
1105    }
1106}
1107
1108impl SpecTelemetryCounters {
1109    fn record_round(&self, drafted: usize, accepted: usize) {
1110        debug_assert!(accepted <= drafted);
1111        self.rounds.fetch_add(1, Ordering::Relaxed);
1112        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1113        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1114        for counter in self.pos_drafted.iter().take(drafted) {
1115            counter.fetch_add(1, Ordering::Relaxed);
1116        }
1117        for counter in self.pos_accepted.iter().take(accepted) {
1118            counter.fetch_add(1, Ordering::Relaxed);
1119        }
1120    }
1121
1122    /// Round-stream keeps each round's accept length on device; retain exact scalar totals while
1123    /// leaving the per-position arrays untouched, matching the pre-existing telemetry contract.
1124    fn record_totals(&self, rounds: usize, drafted: usize, accepted: usize) {
1125        self.rounds.fetch_add(rounds as u64, Ordering::Relaxed);
1126        self.drafted.fetch_add(drafted as u64, Ordering::Relaxed);
1127        self.accepted.fetch_add(accepted as u64, Ordering::Relaxed);
1128    }
1129
1130    fn snapshot(&self) -> SpecTelemetry {
1131        SpecTelemetry {
1132            rounds: self.rounds.load(Ordering::Relaxed),
1133            drafted: self.drafted.load(Ordering::Relaxed),
1134            accepted: self.accepted.load(Ordering::Relaxed),
1135            pos_drafted: std::array::from_fn(|j| self.pos_drafted[j].load(Ordering::Relaxed)),
1136            pos_accepted: std::array::from_fn(|j| self.pos_accepted[j].load(Ordering::Relaxed)),
1137        }
1138    }
1139}
1140
1141pub struct SpecSession {
1142    pub(crate) cache: Cache,
1143    pub(crate) scratch: MtpScratch,
1144    /// Every token whose state the caches hold, in order (prompt turns + generated), INCLUDING
1145    /// overshoot: spec commits accepted drafts past max_new; those rows are in the caches, so the
1146    /// session must count them. Callers render output from this, not from their own echo.
1147    pub committed: Vec<u32>,
1148    /// Pre-output_norm hidden of the LAST committed row (device). None before the first turn.
1149    pub(crate) last_h: Option<CudaSlice<f32>>,
1150    /// Greedy argmax predicting the token AFTER committed.last() (from the last turn's final
1151    /// logits). Fuels empty-suffix continuation bursts (serve): the next turn emits this token
1152    /// first, feeds it, and the round loop resumes without any prime. None before the first turn.
1153    pub next_pred: Option<u32>,
1154    /// SAMPLED-SPEC stream continuity across bursts: Philox event counters persist here so a
1155    /// session's randomness never repeats between generate_spec_session calls. (0,0) at admit.
1156    pub sctr: u32,
1157    pub uctr: u32,
1158    /// PERSISTENT DRAFT-GRAPH CONTEXT (2026-08-01, the serve-burst fixed-cost fix): the captured
1159    /// draft graph(s) + every device I/O buffer they bake, carried ACROSS generate_spec_session
1160    /// calls. Before this, every serve burst re-captured the draft graph (2 warmup forwards +
1161    /// instantiate) — measured ~16ms/burst on H100 q27 (MEMRA_SPEC_BURST sweep,
1162    /// research/spec-serving-20260801). None before the first turn; error paths drop it
1163    /// (next burst recaptures — serve retires errored sessions anyway).
1164    pub(crate) draft_ctx: Option<DraftGraphCtx>,
1165    /// PENDING-CARRY across bursts (2026-08-01, the serve burst-boundary fix): the bonus token
1166    /// emitted by the last round but NOT committed to the caches. The old tail committed it with
1167    /// a solo T=1 trunk pass (+ draft fill), and the next burst's setup fed the stashed next_pred
1168    /// with ANOTHER solo pass — 2x ~11.5ms/burst measured on H100 q27 ([spec-setup] trace).
1169    /// Carrying it lets the next empty-suffix greedy burst consume it as round-0 verify col 0,
1170    /// exactly like a mid-burst full-accept boundary (no solo passes). INVARIANT: when set,
1171    /// `committed` (== cache rows) EXCLUDES this token although it was already emitted in the
1172    /// last burst's output, and `last_h` holds the hidden of the last COMMITTED row (its
1173    /// predecessor — the chain-seed/fill anchor). `next_pred` is None (unknown without the
1174    /// commit pass). Non-empty-suffix or sampled turns must flush first (spec_flush_pending);
1175    /// generate_spec_session_sampled does this at entry, and serve parks only flushed sessions.
1176    pub pending_tok: Option<u32>,
1177    /// SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): the state at this
1178    /// turn's PROMPT-END boundary, retained so a later turn can REWIND here. See
1179    /// [`SpecCheckpoint`]. Refreshed by every non-empty prime; None until the first one, and on
1180    /// a rig too tight to hold it (a failed capture is silent — resume just isn't available).
1181    pub(crate) turn_ckpt: Option<SpecCheckpoint>,
1182    /// Session-lifetime acceptance telemetry. Relaxed atomics update at the host-side round
1183    /// accounting the loop already does — no syncs, no allocation. NOTE a
1184    /// pool-resumed session carries the PREVIOUS requests' counts; per-request consumers
1185    /// diff with [`SpecTelemetry::delta_since`] around each burst.
1186    telem: SpecTelemetryCounters,
1187    /// PREFIX-CACHE publication request (lane/spec-prefix-cache): worker sets this to the
1188    /// miss-LCP boundary before a cold burst; the prime captures at exactly that split (it must
1189    /// coincide with the burst's `prime_split` or no capture happens). One-shot: consumed by the
1190    /// prime, result lands in `boundary_captures`.
1191    pub capture_at: Option<usize>,
1192    /// The captures the last prime produced (see [`SpecBoundaryCapture`]). Worker drains them
1193    /// post-burst to assemble prefix entries. A failed capture is silent, like `turn_ckpt` —
1194    /// publication just isn't available for that request. Plural since
1195    /// lane/frspec-multiturn-cache (2026-08-21): a cold burst can capture BOTH the miss-LCP
1196    /// split (the shared-prefix class) and the stable pre-generation boundary (the
1197    /// next-turn re-render class) — one entry per stop, exactly the boundary set the plain
1198    /// prefill tick publishes/checkpoints.
1199    pub boundary_captures: Vec<SpecBoundaryCapture>,
1200    /// STABLE-BOUNDARY TURN CHECKPOINT REQUEST (lane/frspec-multiturn-cache, 2026-08-21): the
1201    /// ABSOLUTE committed-length position the next non-empty prime should capture `turn_ckpt`
1202    /// at, instead of prompt-end. The worker sets it to the STABLE PRE-GENERATION boundary
1203    /// (`plain_checkpoint_boundary` — before the live generation header the client rewrites),
1204    /// porting the 2026-08-09 plain-tier fix: a prompt-end spec checkpoint includes the
1205    /// template's live assistant-generation header (`<|im_start|>assistant\n<think>\n`), which
1206    /// the NEXT turn's re-render replaces, so `affinity_match` diverged a couple tokens below
1207    /// the checkpoint and the spec pool declined 100% of multi-turn agent traffic (measured:
1208    /// `spec-affinity: declined (history diverged at 6811 of checkpoint 6813)`,
1209    /// research/multiturn-cache-20260821 B4). One-shot, `capture_at` convention; None = legacy
1210    /// prompt-end capture.
1211    pub ckpt_at: Option<usize>,
1212}
1213impl SpecSession {
1214    /// Context capacity of the session's caches (the server's ContextFull guard).
1215    pub fn cache_max_ctx(&self) -> usize {
1216        self.cache.max_ctx
1217    }
1218    /// Read access to the live trunk cache (lane/spec-prefix-cache): the worker slices
1219    /// full-attn KV rows `[0..capture.pos)` out of it when publishing a boundary capture —
1220    /// those rows are append-only for the session's lifetime (rollbacks never truncate below
1221    /// the prime boundary), so no copy was taken at prime time.
1222    pub fn cache_ref(&self) -> &Cache {
1223        &self.cache
1224    }
1225    /// Read access to the persistent draft-scratch plane (lane/spec-on-cache-hit): the
1226    /// worker slices rows `[0..capture.pos)` when publishing a boundary capture, exactly
1227    /// like the trunk KV — draft rows below the prompt end are append-only for the
1228    /// session's lifetime (the prime fill wrote them once; rollbacks reset `len_d` to the
1229    /// committed length, never below the prime boundary, and the true-hidden refresh
1230    /// rewrites generated positions only). Returns `(k, v, k_tok_bytes, v_tok_bytes)`.
1231    /// None when the scratch is ring-backed (Step35 SWA — physical rows are not
1232    /// prefix-addressable; the prefix cache already refuses that class end to end).
1233    pub fn draft_plane_ref(&self) -> Option<(&CudaSlice<u8>, &CudaSlice<u8>, usize, usize)> {
1234        if self.scratch.kv.ring.is_some() {
1235            return None;
1236        }
1237        Some((
1238            &self.scratch.kv.k,
1239            &self.scratch.kv.v,
1240            self.scratch.kv.k_tok_bytes,
1241            self.scratch.kv.v_tok_bytes,
1242        ))
1243    }
1244    /// Snapshot the session's process-local acceptance counters for per-burst diffing.
1245    pub fn telemetry(&self) -> SpecTelemetry {
1246        self.telem.snapshot()
1247    }
1248    /// Committed position this session can REWIND to (its retained prompt-end boundary), if any.
1249    /// A request whose prompt matches `committed[..pos]` exactly can resume from here — see
1250    /// `spec_rewind_to_checkpoint`.
1251    pub fn rewind_pos(&self) -> Option<usize> {
1252        self.turn_ckpt.as_ref().map(|c| c.pos)
1253    }
1254    /// Whether every ring-backed trunk/draft row needed by the retained checkpoint is resident.
1255    pub fn rewind_is_resident(&self) -> bool {
1256        self.turn_ckpt.as_ref().is_some_and(|ckpt| {
1257            self.cache.can_rollback(&ckpt.snap, 0) && self.scratch.can_rewind_to(ckpt.pos)
1258        })
1259    }
1260    /// Is this session in the DEMOTION-READY shape (see [`SpecSession::into_demoted`])?
1261    /// `false` means a carried pending must be flushed first (`spec_flush_pending`), or the
1262    /// session has never run a turn and has no prediction to hand over.
1263    pub fn demote_ready(&self) -> bool {
1264        self.pending_tok.is_none() && self.next_pred.is_some()
1265    }
1266    /// Does this session hold a carried pending bonus (flush required before a handoff/park)?
1267    pub fn has_pending(&self) -> bool {
1268        self.pending_tok.is_some()
1269    }
1270    /// Committed row count == cache rows (the session invariant), for the caller's own
1271    /// `fed`-length cross-check at a handoff boundary.
1272    pub fn committed_len(&self) -> usize {
1273        self.committed.len()
1274    }
1275    /// DEMOTION HANDOFF (lane/spec-gate, 2026-08-07): consume this session and hand its trunk
1276    /// cache + next-token prediction to the plain batched-decode path.
1277    ///
1278    /// WHY THIS IS EXACT (greedy). The invariant at a burst boundary is `cache.pos ==
1279    /// committed.len()`: every committed row has trunk KV + recurrent state, exactly as a plain
1280    /// tokenwise prime of the same `committed` sequence would have left it (that is the
1281    /// session-tail contract, and the same property `spec_rewind_to_checkpoint` and the reuse
1282    /// pool already rely on). `next_pred` is the argmax of the verify's logits for the LAST
1283    /// committed row — and verify-column logits are bit-identical to plain decode's logits at
1284    /// that position, because `matmul_decode_exact` bit-identity IS the basis of the greedy
1285    /// accept walk. So handing (cache, next_pred) to the batched path continues the stream from
1286    /// a state indistinguishable from one the batched path produced itself: the batched tick
1287    /// emits `next_pred`, feeds it into this same cache, and decodes on.
1288    ///
1289    /// `None` when the session is not in the handoff shape — a carried pending (its bonus row is
1290    /// NOT in the cache, so `spec_flush_pending` must commit it first) or no `next_pred` yet
1291    /// (never bursted). Callers must not force it: a half-committed cache handed to the batched
1292    /// path would silently skip a token.
1293    ///
1294    /// The MTP draft scratch, the persistent draft-graph context and the turn checkpoint are
1295    /// DROPPED here (freeing their VRAM): the batched path never drafts, and this handoff is
1296    /// one-way by design — there is no cheap symmetric re-promotion (rebuilding the draft KV
1297    /// would mean an `mtp_kv_fill` over the whole committed history).
1298    pub fn into_demoted(self) -> Option<(Cache, u32)> {
1299        if self.pending_tok.is_some() {
1300            return None;
1301        }
1302        let np = self.next_pred?;
1303        debug_assert_eq!(
1304            self.cache.pos,
1305            self.committed.len(),
1306            "demotion handoff: cache rows != committed tokens"
1307        );
1308        Some((self.cache, np))
1309    }
1310    /// Pool-resume hook (audit Q2): clear the parked draft-graph failure memoization so a
1311    /// NEW request resuming this session gets one fresh capture chance — a transient-pressure
1312    /// capture failure must not persist for the pool's whole lifetime (the TRT #16072 class).
1313    /// Logs once iff a flag was actually set; a no-fallback resume is silent and free.
1314    pub fn reset_graph_fallback_on_resume(&mut self) {
1315        if let Some(line) = self
1316            .draft_ctx
1317            .as_mut()
1318            .and_then(|c| c.failed.reset_on_resume())
1319        {
1320            eprintln!("{line}");
1321        }
1322    }
1323}
1324
1325/// A session's PROMPT-END boundary state, the rewind target for session-affinity resume.
1326///
1327/// WHY THIS BOUNDARY, AND WHY IT IS THE ONLY ONE WORTH KEEPING. The rewrite class this lane
1328/// exists for (a client that strips `<think>` blocks out of prior assistant turns) mutates the
1329/// text the session GENERATED, never the prompt it was given. So turn N's prompt agrees with
1330/// turn N-1's committed tokens up to almost exactly where turn N-1's generation began — the
1331/// prompt-end boundary. Keeping a checkpoint there means the next turn re-primes only its own
1332/// delta (the rewritten answer + the new user turn) instead of the whole conversation.
1333///
1334/// WHAT IT MUST HOLD. Full-attn KV is append-only and position-addressed, so rewinding it is a
1335/// `len` truncation (no data). Linear-attn (GDN) conv/ssm state is mutated IN PLACE with no
1336/// position index, so it must be a real device COPY — that copy is the entire reason a spec
1337/// session could not previously rewind. The MTP draft scratch needs no copy either: its rows
1338/// below the boundary were written by this turn's fill and are never revisited (the per-round
1339/// true-hidden refresh only rewrites the CURRENT burst's committed positions), so rewinding it
1340/// is also just a `len` reset. `last_h` is the hidden of the last row below the boundary — the
1341/// predecessor-pairing anchor the next prime's fill reads for its first row.
1342///
1343/// COST: one `Cache::snapshot` per TURN, on a code path that already takes one per ROUND.
1344pub(crate) struct SpecCheckpoint {
1345    snap: crate::cache::CacheSnapshot,
1346    /// Committed length at the boundary (== cache.pos there, the session invariant).
1347    pos: usize,
1348    /// Pre-output_norm hidden of row `pos - 1`.
1349    last_h: CudaSlice<f32>,
1350}
1351
1352/// PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache, 2026-08-14): the state a spec session
1353/// records at its cold-prime split so the WORKER can publish a cross-request prefix entry —
1354/// the commit-gated-publication port (research/cache-spec-design-20260814/PORT-PLAN.md item 1).
1355/// Only the pieces that are DESTROYED by continuing the prime need copies here: the in-place
1356/// GDN conv/ssm states (via `Cache::snapshot`, same mechanism as [`SpecCheckpoint`]) and the
1357/// boundary logits. Full-attn KV rows `[0..pos)` and draft-scratch rows `[0..pos)` are
1358/// append-only for the session's lifetime (rollbacks never truncate below the prime boundary),
1359/// so the worker slices those from the live caches post-burst instead of copying at prime time.
1360pub struct SpecBoundaryCapture {
1361    pub snap: crate::cache::CacheSnapshot,
1362    /// Token boundary (== cache.pos at capture; == the worker's miss-LCP split).
1363    pub pos: usize,
1364    /// Full-vocab logits after the prefix prime — the entry's boundary logits.
1365    pub logits: Vec<f32>,
1366    /// Pre-output_norm trunk hidden of row `pos - 1` (lane/spec-on-cache-hit): the
1367    /// predecessor-pairing anchor a RESTORED spec session's first suffix-fill row reads
1368    /// (the `SpecSession::last_h` convention). Empty = unavailable (capture stays valid;
1369    /// the fill's zeros row-0 fallback covers it at a bounded acceptance cost).
1370    pub last_h: Vec<f32>,
1371}
1372
1373/// D2H one hidden row out of a `[T, n_embd]` prime hidden stack — the boundary anchor a
1374/// spec boundary capture carries for later restored-session fills. Failure is silent
1375/// (`turn_ckpt` convention): the capture publishes without an anchor.
1376fn capture_boundary_hidden(
1377    e: &Engine,
1378    h_rows: &CudaSlice<f32>,
1379    pos: usize,
1380    n_embd: usize,
1381) -> Vec<f32> {
1382    if pos == 0 || h_rows.len() < pos * n_embd {
1383        return Vec::new();
1384    }
1385    let Ok(mut row) = e.uninit(n_embd) else {
1386        return Vec::new();
1387    };
1388    if e.copy_view_into(
1389        &mut row,
1390        0,
1391        &h_rows.slice((pos - 1) * n_embd..pos * n_embd),
1392        n_embd,
1393    )
1394    .is_err()
1395    {
1396        return Vec::new();
1397    }
1398    e.dtoh(&row).unwrap_or_default()
1399}
1400
1401/// ROLLBACK DOOR for sampled BOUNDARY tokens (lane/sampled-spec-quality, 2026-08-19).
1402/// Default ON: the token a burst emits at its own boundary is drawn from the request's
1403/// sampler. `MEMRA_SPEC_SAMPLED_BOUNDARY=0` restores the pre-lane posture (an ARGMAX at
1404/// every boundary) without touching greedy, which is byte-unaffected either way.
1405pub fn spec_sampled_boundary_on() -> bool {
1406    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1407    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_SAMPLED_BOUNDARY").as_deref() != Ok("0"))
1408}
1409
1410/// ROLLBACK DOOR for SESSION-SPANNING penalty history (lane/sampled-spec-quality).
1411/// Default ON: `pen_hist` is seeded from the session's committed tail, so repetition /
1412/// frequency / presence penalties see the whole stream. `MEMRA_SPEC_PEN_SESSION=0`
1413/// restores the pre-lane posture (each burst restarts the window from its own prompt
1414/// slice, i.e. from NOTHING on a continuation burst) — and with the door shut the worker
1415/// must keep refusing penalized sampled prefix-cache restores, because the restored
1416/// session's continuation burst is handed no prompt slice at all.
1417pub fn spec_pen_session_on() -> bool {
1418    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1419    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_PEN_SESSION").as_deref() != Ok("0"))
1420}
1421
1422/// ROLLBACK DOOR for extended-entry publication from a RESTORED session
1423/// (lane/sampled-spec-quality, Item 3). Default ON: a converted prefix-cache hit that fed a
1424/// suffix captures its own prompt-end boundary so the NEXT turn can hit a longer prefix.
1425/// `MEMRA_SPEC_RESTORE_REPUBLISH=0` restores the pre-lane posture (a namespace learns exactly
1426/// one boundary and never advances it). Whole-entry semantics only — the boundary is the
1427/// restored session's own prompt end, so `entry_pos != fed_len` still refuses on the way in.
1428pub fn spec_restore_republish_on() -> bool {
1429    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1430    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_RESTORE_REPUBLISH").as_deref() != Ok("0"))
1431}
1432
1433/// Diagnostics: name every boundary token on stderr (`MEMRA_SPEC_BOUNDARY_TRACE=1`), with
1434/// the argmax the pre-lane code would have emitted from the same row. This is how the
1435/// lane MEASURES the boundary rate and the deviation rate instead of estimating them.
1436fn spec_boundary_trace() -> bool {
1437    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1438    *ON.get_or_init(|| std::env::var("MEMRA_SPEC_BOUNDARY_TRACE").as_deref() == Ok("1"))
1439}
1440
1441/// llama-parity floor for the penalty window when the request does not ask for a bigger
1442/// one (`repeat_last_n` default). The serve API arms `penalty_last_n = PEN_WINDOW_MAX` for any
1443/// non-identity penalty, so this floor only matters to explicit small windows and to the
1444/// CLI env path.
1445const PEN_WINDOW_FLOOR: usize = 64;
1446
1447/// CEILING on the penalty window, and it is a COST bound, not a semantic preference.
1448/// `penalize_logits_f32` (cu/spec_sample.cu) dedups on device by having thread `i` scan
1449/// `hist[0..i]`, so a pass is O(n_hist²) and it runs ~3x per verify round (the q rows, the
1450/// p column, the bonus column). The serve API uses this same bound for every non-identity
1451/// penalty so host/plain, sparse-device, and speculative sampling cannot change logits on
1452/// admission demotion. An uncapped 128k-token history would put ~1.7e10
1453/// comparisons per pass, tens of ms per round, i.e. penalties would silently destroy decode
1454/// throughput on exactly the long-context requests that most want them. 8192 keeps a pass
1455/// at ~7e7 comparisons (tens of microseconds) while still being **128x wider than the
1456/// pre-lane effective window** (64 prompt-tail tokens + whatever the current burst had
1457/// generated). A request that genuinely needs a window beyond this wants host-side dedup +
1458/// counts through a new kernel signature — a follow-up lane, named here rather than hidden.
1459/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route's accept walk and
1460/// the dspark_sample_gate binary trim their uploads with the SAME cap — a second constant
1461/// is a second thing to drift.
1462pub const PEN_WINDOW_MAX: usize = 8192;
1463
1464/// Seed a penalty window over the SESSION, not the burst (lane/sampled-spec-quality,
1465/// Item 2). The window is the last `max(penalty_last_n, 64)` tokens of
1466/// `session_committed ++ burst_prompt` — for a cold turn-1 burst (`session_committed`
1467/// empty, default `penalty_last_n`) that is byte-identically the pre-lane
1468/// `prompt.iter().rev().take(64).rev()`; for a continuation burst it is the stream the
1469/// client actually asked us to penalize, where the pre-lane code had NOTHING.
1470/// `pub` since lane/dspark-penalized-sampled-20260821: the dspark route seeds its session
1471/// window through the SAME function (one definition of "the window" across both spec
1472/// routes and the gate binary's trunk-only reference arm).
1473pub fn pen_window_seed(
1474    session_committed: &[u32],
1475    burst_prompt: &[u32],
1476    penalty_last_n: usize,
1477) -> Vec<u32> {
1478    let win = penalty_last_n.clamp(PEN_WINDOW_FLOOR, PEN_WINDOW_MAX);
1479    let take_prompt = burst_prompt.len().min(win);
1480    let take_sess = (win - take_prompt).min(session_committed.len());
1481    let mut hist = Vec::with_capacity(take_sess + take_prompt);
1482    hist.extend_from_slice(&session_committed[session_committed.len() - take_sess..]);
1483    hist.extend_from_slice(&burst_prompt[burst_prompt.len() - take_prompt..]);
1484    hist
1485}
1486
1487/// Draw a BOUNDARY token from the target distribution the request asked for
1488/// (lane/sampled-spec-quality, Item 1) — the fix for "sampled spec emits an ARGMAX token at
1489/// every burst boundary".
1490///
1491/// WHY THIS EXISTS. A spec burst's first emitted token is not produced by the accept walk:
1492/// it comes off a logits row that already exists (the prime's last row on a cold burst; the
1493/// row after the last committed token on a continuation burst; the prefix-cache entry's
1494/// boundary row on a restored one). Pre-lane that token was `argmax` in BOTH sampling
1495/// regimes, so a sampled stream took a greedy token once per burst — measured, not
1496/// estimated, in research/spec-cache-20260818/SAMPLED-QUALITY.md. At temperature > 0 the
1497/// customer asked for a sampled token, so this draws one.
1498///
1499/// THE PROGRAM IS THE FULL-ACCEPT BONUS'S PROGRAM, deliberately: penalize the row (over the
1500/// session's window), take this row's OWN filter stats (the sampfix-20260805 law — stats
1501/// from a neighbour row mis-scale every `e0` and can wipe the row to token 0), gumbel-perturb
1502/// with the session's Philox stream at `*sctr`, argmax the perturbed row. Reusing the bonus's
1503/// composition means `sample_check`'s distributional oracle covers this draw too, and the
1504/// boundary token is drawn from the same filtered/penalized `p` the accept walk targets.
1505///
1506/// THE STREAM IS THE SESSION'S, NOT A FRESH ONE. `sctr` is the caller's live counter and is
1507/// advanced by exactly one, so a boundary draw consumes the next value in the same Philox
1508/// stream the accept walk uses — never a second, independently seeded stream (which would be
1509/// a new distributional bug: two streams from one seed correlate wherever their counters
1510/// collide). That also makes a restored session's boundary draw at `sctr == 0` bit-identical
1511/// to the cold session's own first draw from the same logits row, which is what preserves the
1512/// sampled-hit lane's per-seed hit==cold byte identity.
1513#[allow(clippy::too_many_arguments)]
1514pub fn sample_boundary_token_dev(
1515    e: &Engine,
1516    logits: &CudaSlice<f32>,
1517    n_vocab: usize,
1518    sp: &SpecSampling,
1519    pen_hist: &[u32],
1520    sctr: &mut u32,
1521    site: &str,
1522) -> Result<u32, Box<dyn std::error::Error>> {
1523    debug_assert!(
1524        sp.temp > 0.0,
1525        "boundary sampling is the sampled regime only"
1526    );
1527    // Own copy: penalize_logits mutates in place and the caller's row is live state
1528    // (prime_logits back the constrained recompute; last_col_logits backs round 0's accept).
1529    let mut col = e.zeros(n_vocab)?;
1530    e.copy_into(&mut col, 0, logits, n_vocab)?;
1531    let pen_on = sp.penalty_last_n > 0
1532        && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
1533    if pen_on && !pen_hist.is_empty() {
1534        // window trim mirrors the round loop's own upload (`pen_hist[w0..]`), cap included.
1535        let w0 = pen_hist
1536            .len()
1537            .saturating_sub(sp.penalty_last_n.min(PEN_WINDOW_MAX));
1538        let hist = &pen_hist[w0..];
1539        let hd = e.htod_u32_v(hist)?;
1540        e.penalize_logits(
1541            &mut col,
1542            &hd,
1543            hist.len(),
1544            sp.penalty_repeat,
1545            sp.penalty_freq,
1546            sp.penalty_present,
1547            n_vocab,
1548        )?;
1549    }
1550    let rows0 = e.htod_i32(&[0])?;
1551    let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
1552    e.filter_stats(
1553        &col, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1, sp.temp, sp.top_k,
1554        sp.top_p, sp.min_p,
1555    )?;
1556    let (th, mx) = (e.dtoh(&th_d)?[0], e.dtoh(&mx_d)?[0]);
1557    let mut perturb = e.zeros(n_vocab)?;
1558    e.gumbel_perturb_filtered(&col, &mut perturb, n_vocab, sp.seed, *sctr, sp.temp, mx, th)?;
1559    *sctr = sctr.wrapping_add(1);
1560    let td = e.argmax_token_device(&perturb, n_vocab)?;
1561    let tok = guard_vocab_token(
1562        e.dtoh_u32_one(&td)?,
1563        n_vocab,
1564        &format!("sampled boundary token (site={site})"),
1565    )?;
1566    if spec_boundary_trace() {
1567        // the pre-lane token, from the SAME row, so the deviation rate is measurable.
1568        let raw = e.argmax_token_device(logits, n_vocab)?;
1569        let greedy = e.dtoh_u32_one(&raw)?;
1570        eprintln!(
1571            "[spec-boundary] site={site} sampled={tok} argmax={greedy} \
1572             deviates={} temp={} sctr={}",
1573            (tok != greedy) as u8,
1574            sp.temp,
1575            sctr.wrapping_sub(1),
1576        );
1577    }
1578    Ok(tok)
1579}
1580
1581/// Host-row twin of [`sample_boundary_token_dev`] (the prime / feed / entry rows arrive as
1582/// host `Vec<f32>`).
1583#[allow(clippy::too_many_arguments)]
1584pub fn sample_boundary_token(
1585    e: &Engine,
1586    logits: &[f32],
1587    sp: &SpecSampling,
1588    pen_hist: &[u32],
1589    sctr: &mut u32,
1590    site: &str,
1591) -> Result<u32, Box<dyn std::error::Error>> {
1592    let n_vocab = logits.len();
1593    let d = e.htod(logits)?;
1594    sample_boundary_token_dev(e, &d, n_vocab, sp, pen_hist, sctr, site)
1595}
1596
1597struct SpecPipeTraceClock {
1598    pair: usize,
1599    started: std::time::Instant,
1600}
1601
1602#[derive(Clone)]
1603struct SpecPipeTraceCtx {
1604    clock: std::sync::Arc<SpecPipeTraceClock>,
1605    round: usize,
1606    lane: usize,
1607}
1608
1609struct SpecPipeTraceMarker {
1610    trace: SpecPipeTraceCtx,
1611    phase: &'static str,
1612    edge: &'static str,
1613    slot: Option<usize>,
1614}
1615
1616unsafe extern "C" fn spec_pipe_trace_marker(raw: *mut std::ffi::c_void) {
1617    let marker = unsafe { Box::from_raw(raw.cast::<SpecPipeTraceMarker>()) };
1618    let lane = if marker.trace.lane == 0 { "A" } else { "B" };
1619    let slot = marker
1620        .slot
1621        .map(|v| v.to_string())
1622        .unwrap_or_else(|| "-".into());
1623    let t_ms = marker.trace.clock.started.elapsed().as_secs_f64() * 1e3;
1624    use std::io::Write as _;
1625    let stderr = std::io::stderr();
1626    let mut stderr = stderr.lock();
1627    let _ = writeln!(
1628        stderr,
1629        "[spec-pipe-timeline] pair={} round={} lane={lane} phase={} edge={} \
1630         slot={slot} t_ms={t_ms:.3}",
1631        marker.trace.clock.pair, marker.trace.round, marker.phase, marker.edge,
1632    );
1633}
1634
1635fn enqueue_spec_pipe_trace_marker(
1636    stream: &cudarc::driver::CudaStream,
1637    trace: Option<&SpecPipeTraceCtx>,
1638    phase: &'static str,
1639    edge: &'static str,
1640    slot: Option<usize>,
1641) -> Result<(), Box<dyn std::error::Error>> {
1642    let Some(trace) = trace else {
1643        return Ok(());
1644    };
1645    let marker = Box::new(SpecPipeTraceMarker {
1646        trace: trace.clone(),
1647        phase,
1648        edge,
1649        slot,
1650    });
1651    let raw = Box::into_raw(marker);
1652    let result = unsafe {
1653        cudarc::driver::result::stream::launch_host_function(
1654            stream.cu_stream(),
1655            spec_pipe_trace_marker,
1656            raw.cast(),
1657        )
1658    };
1659    if let Err(err) = result {
1660        unsafe {
1661            drop(Box::from_raw(raw));
1662        }
1663        return Err(err.into());
1664    }
1665    Ok(())
1666}
1667
1668#[derive(Default)]
1669struct SpecPipeProgress {
1670    setup_done: [bool; 2],
1671    draft_done: [usize; 2],
1672    stage0_done: [usize; 2],
1673    verify_done: [usize; 2],
1674    accept_done: [usize; 2],
1675    finished: [bool; 2],
1676    aborted: bool,
1677}
1678
1679/// Host-side issue coordinator for the reduced two-session speculative pipeline. Each session
1680/// keeps its existing call stack and round locals; this object only orders phase entry. The
1681/// primary mutex spans whole draft/accept/tail issue regions so Engine's single-stream scratch
1682/// cannot be interleaved by the two host threads.
1683struct SpecPipeSync {
1684    progress: std::sync::Mutex<SpecPipeProgress>,
1685    changed: std::sync::Condvar,
1686    primary: std::sync::Mutex<()>,
1687    trace: Option<std::sync::Arc<SpecPipeTraceClock>>,
1688}
1689
1690impl SpecPipeSync {
1691    fn new() -> Self {
1692        static TRACE_PAIR: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
1693        let trace = (std::env::var("MEMRA_SPEC_PIPE_TRACE").as_deref() == Ok("1")).then(|| {
1694            std::sync::Arc::new(SpecPipeTraceClock {
1695                pair: TRACE_PAIR.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1,
1696                started: std::time::Instant::now(),
1697            })
1698        });
1699        Self {
1700            progress: std::sync::Mutex::new(SpecPipeProgress::default()),
1701            changed: std::sync::Condvar::new(),
1702            primary: std::sync::Mutex::new(()),
1703            trace,
1704        }
1705    }
1706}
1707
1708#[derive(Clone)]
1709struct SpecPipeLane {
1710    sync: std::sync::Arc<SpecPipeSync>,
1711    lane: usize,
1712}
1713
1714impl SpecPipeLane {
1715    fn peer(&self) -> usize {
1716        1 - self.lane
1717    }
1718
1719    fn aborted() -> Box<dyn std::error::Error> {
1720        "paired speculative peer aborted".into()
1721    }
1722
1723    fn trace(&self, round: usize) -> Option<SpecPipeTraceCtx> {
1724        self.sync.trace.as_ref().map(|clock| SpecPipeTraceCtx {
1725            clock: clock.clone(),
1726            round,
1727            lane: self.lane,
1728        })
1729    }
1730
1731    fn setup_begin(&self) -> Result<(), Box<dyn std::error::Error>> {
1732        let mut p = self.sync.progress.lock().unwrap();
1733        while !p.aborted && self.lane == 1 && !p.setup_done[0] && !p.finished[0] {
1734            p = self.sync.changed.wait(p).unwrap();
1735        }
1736        if p.aborted {
1737            Err(Self::aborted())
1738        } else {
1739            Ok(())
1740        }
1741    }
1742
1743    fn setup_end(&self) {
1744        let mut p = self.sync.progress.lock().unwrap();
1745        p.setup_done[self.lane] = true;
1746        self.sync.changed.notify_all();
1747    }
1748
1749    fn draft_begin(
1750        &self,
1751        round: usize,
1752    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1753        let peer = self.peer();
1754        let mut p = self.sync.progress.lock().unwrap();
1755        loop {
1756            if p.aborted {
1757                return Err(Self::aborted());
1758            }
1759            let setup_ready =
1760                (p.setup_done[0] || p.finished[0]) && (p.setup_done[1] || p.finished[1]);
1761            let prior_ready = p.accept_done[self.lane] >= round
1762                && (p.accept_done[peer] >= round || p.finished[peer]);
1763            let turn_ready = if self.lane == 0 {
1764                true
1765            } else {
1766                p.draft_done[0] > round || p.finished[0]
1767            };
1768            if setup_ready && prior_ready && turn_ready {
1769                break;
1770            }
1771            p = self.sync.changed.wait(p).unwrap();
1772        }
1773        drop(p);
1774        Ok(self.sync.primary.lock().unwrap())
1775    }
1776
1777    fn draft_end(&self, round: usize) {
1778        let mut p = self.sync.progress.lock().unwrap();
1779        p.draft_done[self.lane] = round + 1;
1780        self.sync.changed.notify_all();
1781    }
1782
1783    /// Admit stage 0 and return whether this lane owns the interval's one reverse fence.
1784    /// Lane B releases as soon as lane A has issued its boundary TX, not after A's full body.
1785    fn stage0_begin(&self, round: usize) -> Result<bool, Box<dyn std::error::Error>> {
1786        let peer = self.peer();
1787        let mut p = self.sync.progress.lock().unwrap();
1788        loop {
1789            if p.aborted {
1790                return Err(Self::aborted());
1791            }
1792            let ready = if self.lane == 0 {
1793                p.draft_done[0] > round && (p.draft_done[1] > round || p.finished[1])
1794            } else {
1795                p.draft_done[1] > round && (p.stage0_done[0] > round || p.finished[0])
1796            };
1797            if ready {
1798                return Ok(self.lane == 0 || p.finished[peer]);
1799            }
1800            p = self.sync.changed.wait(p).unwrap();
1801        }
1802    }
1803
1804    fn stage0_end(&self, round: usize) {
1805        let mut p = self.sync.progress.lock().unwrap();
1806        p.stage0_done[self.lane] = round + 1;
1807        self.sync.changed.notify_all();
1808    }
1809
1810    /// Stage 1 is single-owner per engine. A proceeds immediately after its own ticket; B waits
1811    /// for A's full stage1/head issue so only A.S1 and B.S0 can overlap.
1812    fn stage1_begin(&self, round: usize) -> Result<(), Box<dyn std::error::Error>> {
1813        let mut p = self.sync.progress.lock().unwrap();
1814        while !p.aborted
1815            && !(p.stage0_done[self.lane] > round
1816                && (self.lane == 0 || p.verify_done[0] > round || p.finished[0]))
1817        {
1818            p = self.sync.changed.wait(p).unwrap();
1819        }
1820        if p.aborted {
1821            Err(Self::aborted())
1822        } else {
1823            Ok(())
1824        }
1825    }
1826
1827    fn verify_end(&self, round: usize) {
1828        let mut p = self.sync.progress.lock().unwrap();
1829        p.verify_done[self.lane] = round + 1;
1830        self.sync.changed.notify_all();
1831    }
1832
1833    fn accept_begin(
1834        &self,
1835        round: usize,
1836    ) -> Result<std::sync::MutexGuard<'_, ()>, Box<dyn std::error::Error>> {
1837        let mut p = self.sync.progress.lock().unwrap();
1838        loop {
1839            if p.aborted {
1840                return Err(Self::aborted());
1841            }
1842            let ready = if self.lane == 0 {
1843                p.verify_done[0] > round && (p.verify_done[1] > round || p.finished[1])
1844            } else {
1845                p.verify_done[1] > round && (p.accept_done[0] > round || p.finished[0])
1846            };
1847            if ready {
1848                break;
1849            }
1850            p = self.sync.changed.wait(p).unwrap();
1851        }
1852        drop(p);
1853        Ok(self.sync.primary.lock().unwrap())
1854    }
1855
1856    fn accept_end(&self, round: usize) {
1857        let mut p = self.sync.progress.lock().unwrap();
1858        p.accept_done[self.lane] = round + 1;
1859        self.sync.changed.notify_all();
1860    }
1861
1862    fn primary(&self) -> std::sync::MutexGuard<'_, ()> {
1863        self.sync.primary.lock().unwrap()
1864    }
1865
1866    fn finish(&self, failed: bool) {
1867        let mut p = self.sync.progress.lock().unwrap();
1868        p.finished[self.lane] = true;
1869        p.aborted |= failed;
1870        self.sync.changed.notify_all();
1871    }
1872}
1873
1874struct SpecPipeFinish<'a> {
1875    lane: &'a SpecPipeLane,
1876    closed: bool,
1877}
1878
1879impl<'a> SpecPipeFinish<'a> {
1880    fn new(lane: &'a SpecPipeLane) -> Self {
1881        Self {
1882            lane,
1883            closed: false,
1884        }
1885    }
1886
1887    fn close(&mut self, failed: bool) {
1888        self.lane.finish(failed);
1889        self.closed = true;
1890    }
1891}
1892
1893impl Drop for SpecPipeFinish<'_> {
1894    fn drop(&mut self) {
1895        if !self.closed {
1896            self.lane.finish(true);
1897        }
1898    }
1899}
1900
1901/// Scoped transfer of one exclusively-borrowed session to the second host issue thread.
1902/// `CudaGraph` is not marked Send by cudarc because its raw driver handles carry no automatic
1903/// trait. CUDA driver graph handles are context-scoped rather than OS-thread-affine; the caller
1904/// binds that context before touching the session, joins before returning, and never aliases the
1905/// pointer. Keep this exception local to the experimental pair call instead of marking the public
1906/// session type Send.
1907struct SpecPipeSessionPtr(*mut SpecSession);
1908
1909unsafe impl Send for SpecPipeSessionPtr {}
1910
1911impl SpecPipeSessionPtr {
1912    unsafe fn get_mut(&mut self) -> &mut SpecSession {
1913        unsafe { &mut *self.0 }
1914    }
1915}
1916
1917/// Per-session persistent draft-graph context: the captured CUDA graph(s) plus the device
1918/// buffers whose POINTERS the capture bakes. Reuse legality: the greedy capture bakes only
1919/// session-stable pointers (the session's own MtpScratch KV — allocated once, never realloc'd;
1920/// the model's resident embedding; the process-wide OnceLock p_min) and the g_* buffers held
1921/// HERE — so one capture serves the session's whole lifetime. The sampled capture additionally
1922/// bakes (seed, temp) as capture-time constants and needs k q-slots — keyed by `s_key`, dropped
1923/// and recaptured when a pool-resumed request changes them. `*_failed` memoizes a failed capture
1924/// so the eager fallback doesn't pay a doomed capture attempt every burst.
1925/// Capture identity of the parked SAMPLED draft graph (`DraftGraphCtx::graph_s`).
1926///
1927/// EXACTNESS, not perf (lane/graph-s-key-exactness-20260819; receipts
1928/// `research/spec-cache-20260818/GRAPH-S-KEY.md`). Two classes of field live here, both
1929/// load-bearing:
1930///
1931/// - **Baked constants.** `seed` and `temp` are capture-time constants INSIDE the graph and `k`
1932///   sizes the q slots its replays write. A resumed request changing any of them must recapture.
1933///   This is all the key used to carry.
1934/// - **Regime fields.** `top_k`/`top_p`/`min_p`/`pen_on` are not baked, but they decide whether
1935///   the captured graph is a legal draft chain AT ALL. The in-graph draw is one gumbel-max over
1936///   the RAW softmax (`gumbel_perturb_ctr`, unfiltered by construction), while the verify builds
1937///   the accept test's `q` from `filter_stats(q_slots, top_k, top_p, min_p)`. If those disagree
1938///   the accept test evaluates a distribution the draft was never sampled from: a draft token
1939///   below the filter threshold gathers `q = 0` (`softmax_gather_filtered_f32`,
1940///   `cu/spec_sample.cu`) and `u * 0 < p` accepts it UNCONDITIONALLY.
1941///
1942/// Omitting the regime fields was reachable — not through the prefix-cache spec restore (that
1943/// path is greedy-only, `memra-server` `spec_restore_convertible`), but through WHOLE-SESSION
1944/// spec reuse: a parked `SpecSession` carries this `DraftGraphCtx`, and the pool-resume probe
1945/// applies no sampler predicate at all. Turn 1 pure-temp parks a graph; turn 2 of the same
1946/// conversation, same explicit seed and temperature, adds `top_p`/`top_k` and inherits it.
1947#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1948pub(crate) struct SampledGraphKey {
1949    seed: u64,
1950    temp_bits: u32,
1951    k: usize,
1952    top_k: i32,
1953    top_p_bits: u32,
1954    min_p_bits: u32,
1955    pen_on: bool,
1956}
1957
1958impl SampledGraphKey {
1959    pub(crate) fn new(
1960        seed: u64,
1961        temp: f32,
1962        k: usize,
1963        top_k: i32,
1964        top_p: f32,
1965        min_p: f32,
1966        pen_on: bool,
1967    ) -> Self {
1968        SampledGraphKey {
1969            seed,
1970            temp_bits: temp.to_bits(),
1971            k,
1972            top_k,
1973            top_p_bits: top_p.to_bits(),
1974            min_p_bits: min_p.to_bits(),
1975            pen_on,
1976        }
1977    }
1978
1979    /// The one regime the PURE-TEMP in-graph sampled chain may stand in for the eager one:
1980    /// nothing but temperature shapes `q`. Computed FROM THE KEY so the capture guard, the
1981    /// launch guard and the key can never drift apart (they were three separate expressions
1982    /// before this lane, and the launch site simply forgot to ask).
1983    pub(crate) fn pure_temp(&self) -> bool {
1984        self.top_k == 0
1985            && f32::from_bits(self.top_p_bits) >= 1.0
1986            && f32::from_bits(self.min_p_bits) <= 0.0
1987            && !self.pen_on
1988    }
1989
1990    /// Truncation filters active — the capture body needs the IN-GRAPH filter nodes
1991    /// (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws from the same
1992    /// filtered distribution the accept test reconstructs. Meaningful only when
1993    /// `graph_capturable`; penalties never reach a capture body.
1994    pub(crate) fn filtered(&self) -> bool {
1995        !self.pure_temp()
1996    }
1997
1998    /// May the sampled draft graph be CAPTURED (and a parked one LAUNCHED) for this regime?
1999    /// Pure-temp always; filtered regimes when the filtered-capture door is on
2000    /// (lane/step37-draft-graph-serving-20260830); penalties never — the per-round history
2001    /// cannot be baked into a graph, and composing a raw-softmax (or stale-history) draw
2002    /// with a penalized accept test is the unconditional-accept exactness bug. Computed FROM
2003    /// THE KEY for the same no-drift reason as `pure_temp`.
2004    pub(crate) fn graph_capturable(&self) -> bool {
2005        !self.pen_on && (self.pure_temp() || spec_graph_filtered_on())
2006    }
2007}
2008
2009/// Per-head captured graphs for the MULTI-HEAD MTP draft chain (step-modulo prefix-replay,
2010/// lane/step37-draft-graph-serving-20260830). The chain POLICY — which head serves step j,
2011/// how long the replayed prefix is, which stored seed feeds row r — stays HOST-SIDE in the
2012/// launch loop, exactly `mtp_chain_forward_dev`'s order; the graphs capture ONE head-row
2013/// forward each, on the head's OWN scratch plane:
2014/// - `interior[i]`: head i, `with_head=false` — KV append + carrier only. Interior rows'
2015///   logits are dead in the eager chain too (`mtp_chain_forward_dev` keeps only the last
2016///   row), so skipping the head matmul changes no consumed byte and removes the eager
2017///   chain's per-replay-row full-vocab matmul.
2018/// - `last[i]`: head i, `with_head=true` + the mode's tail (greedy argmax, or the sampled
2019///   gumbel draw — filtered in-graph when the request carries filters).
2020/// One `DraftChainGraphs` per MODE (greedy vs sampled), owning its keeper: dropping the
2021/// sampled chain on an s_key change never invalidates the greedy one.
2022struct DraftChainGraphs {
2023    interior: Vec<cudarc::driver::CudaGraph>,
2024    last: Vec<cudarc::driver::CudaGraph>,
2025    keeper: Vec<Box<dyn std::any::Any + Send>>,
2026}
2027
2028/// Sampled-tail capture pack for `mtp_head_forward_cap`: the persistent buffers and baked
2029/// constants of the in-graph categorical draw. `filt: None` = the PURE-TEMP body (gumbel
2030/// over the raw softmax), byte-identical to the pre-lane capture; `Some` adds the in-graph
2031/// truncation filter (`filter_stats` + `gumbel_perturb_filtered_ctr`) so the draft draws
2032/// from the same filtered distribution the accept test reconstructs
2033/// (lane/step37-draft-graph-serving-20260830).
2034struct SampledCapArgs<'a> {
2035    ctr: &'a mut CudaSlice<u32>,
2036    perturb: &'a mut CudaSlice<f32>,
2037    q_out: &'a mut CudaSlice<f32>,
2038    seed: u64,
2039    temp: f32,
2040    filt: Option<SampledCapFilter<'a>>,
2041}
2042
2043/// In-graph truncation-filter nodes: the stat slots `filter_stats` fills and the perturb
2044/// reads, plus the filter constants baked into the capture (they live in `s_key`, so a
2045/// request whose filters differ drops the parked graph before this ever goes stale).
2046struct SampledCapFilter<'a> {
2047    rows0: &'a CudaSlice<i32>,
2048    th: &'a mut CudaSlice<f32>,
2049    z: &'a mut CudaSlice<f32>,
2050    mx: &'a mut CudaSlice<f32>,
2051    top_k: i32,
2052    top_p: f32,
2053    min_p: f32,
2054}
2055
2056pub(crate) struct DraftGraphCtx {
2057    g_tok: CudaSlice<u32>,
2058    g_pos: CudaSlice<i32>,
2059    g_seed: CudaSlice<f32>,
2060    g_p: CudaSlice<f32>,
2061    g_ctr: CudaSlice<u32>,
2062    g_q: CudaSlice<f32>,
2063    g_perturb: CudaSlice<f32>,
2064    /// IN-GRAPH filter-stat slots (filtered sampled capture): `filter_stats` writes
2065    /// (th, z, mx) here inside the graph; `gumbel_perturb_filtered_ctr` reads (mx, th) from
2066    /// the same slots. Persistent so the baked pointers survive replays. `g_rows0` is the
2067    /// constant row-index-0 the single-row `filter_stats` launch reads (a captured memcpy
2068    /// source must not be a host temporary).
2069    g_rows0: CudaSlice<i32>,
2070    g_th: CudaSlice<f32>,
2071    g_z: CudaSlice<f32>,
2072    g_mx: CudaSlice<f32>,
2073    q_slots: Vec<CudaSlice<f32>>,
2074    /// DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): packed allowed-set words over the DRAFT
2075    /// head's vocab, at a STABLE address so the captured draft graph's mask node reads the
2076    /// per-position contents the host re-uploads before each replay (the graph-promote
2077    /// pattern from decode.rs). Empty unless the session drafts under a grammar.
2078    g_dmask: CudaSlice<u32>,
2079    /// was `graph` captured WITH the mask node? A parked graph of the wrong shape is dropped.
2080    /// Covers the multi-head `chain` too (single-head and chain are mutually exclusive for a
2081    /// given model, so one flag serves whichever is active).
2082    graph_masked: bool,
2083    graph: Option<cudarc::driver::CudaGraph>,
2084    graph_s: Option<cudarc::driver::CudaGraph>,
2085    /// Multi-head chain graphs (see [`DraftChainGraphs`]): greedy and sampled chains, the
2086    /// chain twins of `graph` / `graph_s`. `chain_s`'s capture identity is `s_key` (shared
2087    /// with `graph_s` — a session is either single-head or chain, never both), and it obeys
2088    /// the same drop rules (key mismatch, penalty regime, mask-shape change).
2089    chain: Option<DraftChainGraphs>,
2090    chain_s: Option<DraftChainGraphs>,
2091    /// Failed-capture memoization for both graphs — LOUD on flip, cleared on pool resume
2092    /// (audit Q2, the TRT #16072 silent-permanent-coverage-loss class).
2093    failed: DraftGraphFallback,
2094    /// Capture identity of `graph_s` — see [`SampledGraphKey`]. `None` iff no sampled graph is
2095    /// parked; a request whose key differs drops the parked graph (and its q slots/keeper).
2096    s_key: Option<SampledGraphKey>,
2097    /// CAPTURE-RETAIN keepers (#68 root cause, 2026-08-04): the warmup-run transients whose
2098    /// pool addresses the captured graph(s) bake. Without these, the transients return to the
2099    /// pool at capture-body exit and later work (burst-boundary prime/fill/commit passes, or a
2100    /// co-served session in the worker) reuses those addresses — the persisted graph's replay
2101    /// then reads/writes live unrelated buffers (exactness corruption, first seen as the ST
2102    /// serve-spec 4B graph-arm corruption; one-shot CLI calls never re-shuffled the pool, which
2103    /// is why run-spec K=1..8 passed on the same checkpoint). Same fix class as
2104    /// capture_graph_retained's gemma/decode.rs sites — hold as long as the graph replays.
2105    keeper: Vec<Box<dyn std::any::Any + Send>>,
2106    keeper_s: Vec<Box<dyn std::any::Any + Send>>,
2107}
2108
2109/// Failed-capture memoization for the two draft graphs (audit Q2, 2026-08-05 — the
2110/// TRT #16072 trap class: pressure-triggered, silent, long-lived coverage loss).
2111///
2112/// Three contracts:
2113/// - LOUD FLIP: `mark_*` returns the warn line exactly on the false→true transition
2114///   (returned, not printed, so the once-per-flip contract is unit-testable); the caller
2115///   `eprintln!`s it UNCONDITIONALLY — a dropped draft graph is never silent. Re-marking
2116///   an already-failed graph returns None (the per-burst memoization that keeps the eager
2117///   fallback from paying a doomed capture attempt every burst).
2118/// - RESET ON RESUME: `reset_on_resume` clears both flags — a parked session resumed by a
2119///   NEW request gets one fresh capture chance instead of carrying a transient-pressure
2120///   failure for the pool's whole lifetime. Returns the note line only when a flag was
2121///   actually set (quiet on the common clean-resume path).
2122/// - Shape-change clears (`clear_*`) stay silent, exactly as before: they precede a fresh
2123///   capture attempt whose own failure would re-flip loudly.
2124#[derive(Default)]
2125pub(crate) struct DraftGraphFallback {
2126    greedy: bool,
2127    sampled: bool,
2128}
2129impl DraftGraphFallback {
2130    fn mark_greedy(&mut self, reason: &str) -> Option<String> {
2131        if self.greedy {
2132            return None;
2133        }
2134        self.greedy = true;
2135        Some(format!(
2136            "[spec] WARN: draft-graph capture failed ({reason}); eager fallback until session resume"
2137        ))
2138    }
2139    fn mark_sampled(&mut self, reason: &str) -> Option<String> {
2140        if self.sampled {
2141            return None;
2142        }
2143        self.sampled = true;
2144        Some(format!(
2145            "[spec] WARN: sampled draft-graph capture failed ({reason}); eager fallback until session resume"
2146        ))
2147    }
2148    fn greedy_failed(&self) -> bool {
2149        self.greedy
2150    }
2151    fn sampled_failed(&self) -> bool {
2152        self.sampled
2153    }
2154    fn clear_greedy(&mut self) {
2155        self.greedy = false;
2156    }
2157    fn clear_sampled(&mut self) {
2158        self.sampled = false;
2159    }
2160    /// Pool-resume reset: both graphs get a fresh capture chance. Some(note) iff any flag
2161    /// was set (so clean resumes stay quiet).
2162    pub(crate) fn reset_on_resume(&mut self) -> Option<String> {
2163        if !self.greedy && !self.sampled {
2164            return None;
2165        }
2166        let which = match (self.greedy, self.sampled) {
2167            (true, true) => "greedy+sampled",
2168            (true, false) => "greedy",
2169            _ => "sampled",
2170        };
2171        self.greedy = false;
2172        self.sampled = false;
2173        Some(format!(
2174            "[spec] draft-graph fallback reset on session resume ({which}); recapture eligible"
2175        ))
2176    }
2177}
2178
2179impl DraftGraphCtx {
2180    fn new(e: &Engine, n_embd: usize, qlen: usize) -> Result<Self, Box<dyn std::error::Error>> {
2181        Ok(DraftGraphCtx {
2182            g_tok: e.alloc_u32_zeroed(1)?,
2183            g_pos: e.htod_i32(&[0])?,
2184            g_seed: e.zeros(n_embd)?,
2185            g_p: e.zeros(1)?,
2186            g_ctr: e.alloc_u32_zeroed(1)?,
2187            g_q: e.zeros(qlen)?,
2188            g_perturb: e.zeros(qlen)?,
2189            g_rows0: e.htod_i32(&[0])?,
2190            g_th: e.zeros(1)?,
2191            g_z: e.zeros(1)?,
2192            g_mx: e.zeros(1)?,
2193            q_slots: Vec::new(),
2194            g_dmask: e.alloc_u32_zeroed(1)?,
2195            graph_masked: false,
2196            graph: None,
2197            graph_s: None,
2198            chain: None,
2199            chain_s: None,
2200            failed: DraftGraphFallback::default(),
2201            s_key: None,
2202            keeper: Vec::new(),
2203            keeper_s: Vec::new(),
2204        })
2205    }
2206}
2207
2208pub(crate) struct MtpScratch {
2209    kv: KvLayer,
2210    /// Logical row capacity. On the graph/DC draft path it also doubles as fa_decode_dc's
2211    /// bucket_max: n_splits is sized from it ONCE, so the graph captured at round 0 stays valid
2212    /// for every later t_kv. Step35 refuses that path and may back this logical extent with the
2213    /// smaller host-indexed SWA ring instead.
2214    cap: usize,
2215    extra: Vec<MtpScratchPlane>,
2216}
2217
2218struct MtpScratchPlane {
2219    kv: KvLayer,
2220    cap: usize,
2221}
2222
2223fn mtp_scratch_layout(
2224    cfg: &memra_gguf::config::ModelConfig,
2225    geom: Option<&crate::hybrid::DraftGeom>,
2226) -> (usize, usize, usize, usize) {
2227    // Student draft heads carry fewer KV heads (head_dim unchanged) -> smaller scratch rows.
2228    let n_head_kv = geom.map(|g| g.n_head_kv).unwrap_or(cfg.n_head_kv as usize);
2229    let head_dim_k = cfg.head_dim_k as usize;
2230    let head_dim_v = cfg.head_dim_v as usize;
2231    assert!(
2232        head_dim_k % 32 == 0 && head_dim_v % 32 == 0,
2233        "KVQUANT requires head_dim%32==0 (MTP scratch)"
2234    );
2235    let kv_dim_k = head_dim_k * n_head_kv;
2236    let kv_dim_v = head_dim_v * n_head_kv;
2237    // The fp8-KV arm deliberately does not reach the draft scratch; keep the exact format
2238    // policy shared with `MtpScratch::new` so admission scales the same allocation.
2239    let (kbb, vbb) = crate::kv_blk_bytes();
2240    let k_tok_bytes = (kv_dim_k / 32) * kbb;
2241    let v_tok_bytes = (kv_dim_v / 32) * vbb;
2242    (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes)
2243}
2244
2245fn mtp_chain_head_index(step: usize, head_count: usize) -> usize {
2246    assert!(head_count > 0, "MTP chain requires at least one head");
2247    step % head_count
2248}
2249
2250impl MtpScratch {
2251    fn alloc_plane(
2252        e: &Engine,
2253        cfg: &memra_gguf::config::ModelConfig,
2254        plan: &memra_gguf::model_plan::ModelPlan,
2255        cap: usize,
2256        geom: Option<&crate::hybrid::DraftGeom>,
2257    ) -> Result<MtpScratchPlane, Box<dyn std::error::Error>> {
2258        let (kv_dim_k, kv_dim_v, k_tok_bytes, v_tok_bytes) = mtp_scratch_layout(cfg, geom);
2259        let ring = if crate::cache::swa_ring_on()
2260            && crate::plan_backend::decode_batch_program(plan)
2261                == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
2262        {
2263            let window = plan
2264                .layers
2265                .iter()
2266                .find_map(|layer| match layer.attention {
2267                    memra_gguf::model_plan::AttentionPlan::SlidingWindow { window, .. } => {
2268                        Some(window as usize)
2269                    }
2270                    _ => None,
2271                })
2272                .ok_or("sliding-gated-MoE draft scratch has no sliding-window layer")?;
2273            Some(crate::cache::KvRing::new(
2274                crate::cache::swa_ring_rows(window, cap),
2275                window,
2276            ))
2277        } else {
2278            None
2279        };
2280        let alloc_rows = ring.as_ref().map(crate::cache::KvRing::rows).unwrap_or(cap);
2281        // Ring-backed planes arm the device base mirror for the dcw draft arm (see
2282        // KvLayer::base_d): the captured chain derives its physical rows from
2283        // (len_d, base_d, window) with zero per-token node updates.
2284        let base_d = match ring.as_ref() {
2285            Some(_) => Some(e.htod_i32(&[0])?),
2286            None => None,
2287        };
2288        Ok(MtpScratchPlane {
2289            kv: KvLayer {
2290                k: e.alloc_u8(alloc_rows * k_tok_bytes)?,
2291                v: e.alloc_u8(alloc_rows * v_tok_bytes)?,
2292                kv_dim_k,
2293                kv_dim_v,
2294                k_tok_bytes,
2295                v_tok_bytes,
2296                len: 0,
2297                ring,
2298                len_d: e.htod_i32(&[0])?,
2299                base_d,
2300            },
2301            cap,
2302        })
2303    }
2304
2305    fn new(
2306        e: &Engine,
2307        cfg: &memra_gguf::config::ModelConfig,
2308        plan: &memra_gguf::model_plan::ModelPlan,
2309        cap: usize,
2310        geom: Option<&crate::hybrid::DraftGeom>,
2311    ) -> Result<Self, Box<dyn std::error::Error>> {
2312        // env-selected KV formats (default 34/24). The fp8-KV arm (MEMRA_KV_FP8) deliberately
2313        // does NOT reach the draft scratch: fp8 drafts drifted acceptance 69-88% -> 46%
2314        // (2026-07-12 A/B); the scratch is tiny, so it keeps baseline q8_0/q5_1 numerics
2315        // while the TRUNK cache carries the fp8 depth win. Scratch append/fa pass g=false.
2316        let primary = Self::alloc_plane(e, cfg, plan, cap, geom)?;
2317        Ok(MtpScratch {
2318            kv: primary.kv,
2319            cap: primary.cap,
2320            extra: Vec::new(),
2321        })
2322    }
2323
2324    fn push_plane(
2325        &mut self,
2326        e: &Engine,
2327        cfg: &memra_gguf::config::ModelConfig,
2328        plan: &memra_gguf::model_plan::ModelPlan,
2329        geom: Option<&crate::hybrid::DraftGeom>,
2330    ) -> Result<(), Box<dyn std::error::Error>> {
2331        self.extra
2332            .push(Self::alloc_plane(e, cfg, plan, self.cap, geom)?);
2333        Ok(())
2334    }
2335
2336    fn plane_count(&self) -> usize {
2337        1 + self.extra.len()
2338    }
2339
2340    fn plane(&self, index: usize) -> (&KvLayer, usize) {
2341        if index == 0 {
2342            (&self.kv, self.cap)
2343        } else {
2344            let plane = &self.extra[index - 1];
2345            (&plane.kv, plane.cap)
2346        }
2347    }
2348
2349    fn plane_mut(&mut self, index: usize) -> (&mut KvLayer, usize) {
2350        if index == 0 {
2351            (&mut self.kv, self.cap)
2352        } else {
2353            let plane = &mut self.extra[index - 1];
2354            (&mut plane.kv, plane.cap)
2355        }
2356    }
2357
2358    // #[track_caller]: set_len/set_plane_len have eight call sites (checkpoint restore, spec
2359    // rollback, session grow, seed replay ...) and the lap failure needs to say WHICH one, not
2360    // just that a rewind was refused.
2361    #[track_caller]
2362    fn set_plane_len(
2363        &mut self,
2364        e: &Engine,
2365        index: usize,
2366        n: usize,
2367    ) -> Result<(), Box<dyn std::error::Error>> {
2368        let caller = std::panic::Location::caller();
2369        let (kv, cap) = self.plane_mut(index);
2370        if let Some(ring) = kv.ring.as_ref() {
2371            if !ring.can_rewind_to(n) {
2372                // NAME THE NUMBERS (2026-08-28). This error is a step37 serving blocker on the
2373                // vendor-default shape and it fires from more than one call path with more than
2374                // one trigger: a long generation walks the checkpoint out of the ring, but a
2375                // ~4.5k-token prompt also fails within 5 s of prime, which accumulation cannot
2376                // explain. A bare message forced two rounds of guessing; the operands make each
2377                // trigger name itself.
2378                let raw = n.saturating_sub(ring.window().saturating_sub(1));
2379                return Err(format!(
2380                    "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})",
2381                    ring.window(),
2382                    ring.base(),
2383                    ring.rows(),
2384                    raw & !31usize,
2385                )
2386                .into());
2387            }
2388        }
2389        kv.len = n;
2390        e.set_i32_one(&mut kv.len_d, n as i32)
2391    }
2392
2393    /// Set BOTH length counters: the host mirror AND the device len_d the captured append/fa read
2394    /// (a 4-byte in-place htod — the counter pointer is baked into the graph, never realloc'd).
2395    /// This is the ONLY truncation/rollback mechanism the persistent draft KV needs.
2396    #[track_caller]
2397    fn set_len(&mut self, e: &Engine, n: usize) -> Result<(), Box<dyn std::error::Error>> {
2398        let caller = std::panic::Location::caller();
2399        if !self.can_rewind_to(n) {
2400            // set_plane_len re-checks and reports the operands; call it so the failure carries
2401            // which plane refused and why, instead of this bare aggregate.
2402            for index in 0..self.plane_count() {
2403                self.set_plane_len(e, index, n)?;
2404            }
2405            return Err(format!(
2406                "SWA ring MTP checkpoint has been lapped; full re-prime required (aggregate rewind_to={n}, no single plane reported, called from {caller})"
2407            )
2408            .into());
2409        }
2410        for index in 0..self.plane_count() {
2411            self.set_plane_len(e, index, n)?;
2412        }
2413        Ok(())
2414    }
2415
2416    fn can_rewind_to(&self, n: usize) -> bool {
2417        (0..self.plane_count()).all(|index| {
2418            self.plane(index)
2419                .0
2420                .ring
2421                .as_ref()
2422                .is_none_or(|ring| ring.can_rewind_to(n))
2423        })
2424    }
2425
2426    /// Pre-arm ring headroom for `rows` upcoming DEVICE-COUNTER appends (the dcw draft arm):
2427    /// a captured chain cannot rebase mid-replay, so any rebase the coming appends could need
2428    /// happens HERE, host-side, before the capture warmups or the round's replays (the rebase
2429    /// arm of `prepare_kv_append` also refreshes the plane's `base_d` device mirror). No-op on
2430    /// flat planes and when the ring already has room; `len` is untouched either way.
2431    fn ensure_dcw_headroom(
2432        &mut self,
2433        e: &Engine,
2434        rows: usize,
2435    ) -> Result<(), Box<dyn std::error::Error>> {
2436        for index in 0..self.plane_count() {
2437            let (kv, _) = self.plane_mut(index);
2438            let Some(ring) = kv.ring.as_ref() else {
2439                continue;
2440            };
2441            let retain = memra_kv::swa_retain_from(kv.len, ring.window(), ring.base());
2442            e.prepare_kv_append(kv, retain, rows)?;
2443        }
2444        Ok(())
2445    }
2446}
2447
2448/// Retained verify intermediates for the REPLAY-FREE partial accept (2026-07-03, the profiled
2449/// #1 spec cost at long ctx: the partial-accept replay was a DUPLICATE trunk pass — ~0.54 extra
2450/// full weight reads per round — recomputing columns the verify had already produced
2451/// bit-identically). Holds, per linear layer, everything needed to rebuild its recurrent state
2452/// to "after the first j verify columns" WITHOUT re-running the trunk:
2453/// - BATCHED-path layers (`gdn`): the exact token-major inputs the round's ONE gdn_scan
2454///   consumed. A prefix re-run of the SAME kernel (t=j) from the snapshot state is bit-identical
2455///   to the first j iterations of the verify's scan — the kernel's t-loop carries state in
2456///   registers and iteration t never depends on T. `qkv_mixed` (the conv input) feeds the
2457///   pure-copy ring rebuild.
2458/// - PER-COLUMN-path layers (`cols`): dtod clones of (conv_state, ssm_state) taken after each
2459///   column 0..t-2 — pure copies of the actual chain states (the last column is never a rebuild
2460///   target: j <= t-1).
2461/// Full-attn layers need nothing: their verify KV rows are bit-identical to eager's (the
2462/// decode-exact contract; verify-probe pins it), so rollback = len truncation.
2463struct GdnStash {
2464    qkv_mixed: CudaSlice<f32>, // [t, conv_dim] token-major (conv input)
2465    q_l2: CudaSlice<f32>,
2466    k_l2: CudaSlice<f32>,
2467    v_g: CudaSlice<f32>, // [t, num_v, d_state]
2468    g_log: CudaSlice<f32>,
2469    beta: CudaSlice<f32>, // [t, num_v]
2470}
2471pub(crate) struct VerifyCkpt {
2472    gdn: Vec<Option<GdnStash>>, // [n_layer], Some iff batched linear path ran
2473    cols: Vec<Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>>>, // [n_layer][col] = (conv, ssm) after col
2474}
2475/// Opaque handle for the dspark round (dflash.rs) — VerifyCkpt stays spec-private.
2476pub(crate) struct DsparkVerifyCkpt(VerifyCkpt);
2477
2478/// Engine-bundle slice 3 (DSF-ROUNDCOST-20260820 §2 row 4 / §5 rank 1): bucketed CUDA
2479/// graphs for the dspark verify's LINEAR-layer segments. The measured verify is ~2,800
2480/// eager launches whose residual cost is DEVICE-side per-launch overhead (slice 2 proved
2481/// host dispatch is not the binder: fully-deferred dispatch bought ~0 wall). The 48 GDN
2482/// layers between full-attention layers are shape-static given vt — no positions, no
2483/// t_kv, state addressed through pointer tables — so runs of them capture per
2484/// (segment, vt) and replay as ONE graph launch each. Full-attention layers stay eager
2485/// (their per-row append/fa arm picks are t_kv-driven — the exec-update extension).
2486///
2487/// Per round out-of-graph: one pointer-table refresh (gdn ping-pong moves the canonical
2488/// handles), one input-staging copy per segment, host parity bookkeeping. Captured via
2489/// `capture_graph_retained` (2 warmups + capture, keeper retains warmup transients so
2490/// pool addresses stay stable); the warmups EXECUTE, so segment conv/ssm state is saved
2491/// before and restored after — the graph's first real launch starts from the exact
2492/// pre-round state. The ckpt column stash rides persistent slabs (written inside the
2493/// graph as memcpy nodes); commit reads them via `dspark_commit_prefix_slab`.
2494/// `MEMRA_DSPARK_VERIFY_GRAPH=0` reverts to the eager walk (byte-identical body).
2495pub(crate) struct DsparkVerifyGraphs {
2496    /// Linear-attention layer indices ascending; `lin_pos[il]` = index into the vecs.
2497    lin: Vec<usize>,
2498    lin_pos: std::collections::HashMap<usize, usize>,
2499    /// [n_lin x 6] pointer table (conv, s0, s1, conv, s1, s0 per layer), refreshed per
2500    /// verify from the live handles; layer il's slice starts at lin_pos[il]*6.
2501    table_all: CudaSlice<u64>,
2502    host_table: Vec<u64>,
2503    /// Persistent per-layer ckpt stash slabs: row r of the verify at slab offset
2504    /// r*words. Shared by every (segment, vt) bucket — one verify runs at a time.
2505    stash_conv: Vec<CudaSlice<f32>>,
2506    stash_ssm: Vec<CudaSlice<f32>>,
2507    conv_words: usize,
2508    ssm_words: usize,
2509    /// Per-vt input/output staging (stable addresses the graphs bake).
2510    stage: std::collections::HashMap<usize, (CudaSlice<f32>, CudaSlice<f32>)>,
2511    /// Per-vt dflash tap-sink buffers — the captured segments bake the tap dst address,
2512    /// so the sink buffer must live (and persist) with the graphs, not with the round.
2513    pub(crate) tap_bufs: std::collections::HashMap<usize, CudaSlice<f32>>,
2514    graphs: std::collections::HashMap<(usize, usize), DsparkSegGraph>,
2515    /// Warmup-corruption guard scratch: pre-capture conv/ssm of every linear layer
2516    /// (sized n_lin — the slice-4c full-verify warmups execute the whole walk).
2517    save_conv: CudaSlice<f32>,
2518    save_ssm: CudaSlice<f32>,
2519    max_run: usize,
2520    n_embd: usize,
2521    /// Set by the verify walk: this round's linear ckpt lives in the slabs (the caller
2522    /// commits through `dspark_commit_prefix_slab` instead of the cols arm).
2523    pub(crate) round_slab: bool,
2524    // ---- slice 4c: full-verify single graph per (vt, rung) ----
2525    /// Full-attention layer indices ascending; `fa_pos[il]` = index into the vec.
2526    fa: Vec<usize>,
2527    fa_pos: std::collections::HashMap<usize, usize>,
2528    /// [n_fa x 2 x t_cap] interleaved (k,v) base-pointer pairs, refreshed per verify;
2529    /// layer il's slice starts at `fa_pos[il] * 2 * t_cap` (the seqs twins read pairs
2530    /// [2z], z < t <= t_cap, so one t_cap-sized table serves every vt).
2531    fa_table: CudaSlice<u64>,
2532    fa_host_table: Vec<u64>,
2533    t_cap: usize,
2534    /// Per-vt position staging for the captured bodies — contents refreshed per round
2535    /// (rope reads row r; the seqs twins derive append slot and T_kv per z from it).
2536    pos_stage: std::collections::HashMap<usize, CudaSlice<i32>>,
2537    /// Full-verify graphs keyed (vt, rung_end, hi).
2538    full: std::collections::HashMap<(usize, usize, usize), DsparkSegGraph>,
2539    /// Largest n with every layer in [0, n) linear or full-attention (walk coverage).
2540    covered: usize,
2541    /// Every layer in [0, n) is linear or full-attention (no MLA/unknown mixers) — the
2542    /// full-verify capture walks all of them.
2543    walk_uniform: bool,
2544    /// Last `(captures, device graph-mem reserved bytes)` reading taken by
2545    /// `HybridModel::dspark_vg_admission_debt` — the two-point base of the MARGINAL debt
2546    /// projection (see `dspark_vg_debt_projection`; a mean-based reading extrapolated the
2547    /// pool's one-time shared allocation and reserved 8.5 GB of phantom VRAM).
2548    debt_obs: Option<(usize, usize)>,
2549}
2550
2551struct DsparkSegGraph {
2552    graph: cudarc::driver::CudaGraph,
2553    _keeper: Vec<Box<dyn std::any::Any + Send>>,
2554}
2555
2556/// Per-call arguments of [`HybridModel::qwen35_tparallel_fa_layer`] — one struct so the
2557/// eager walk and the slice-4c captured full-verify graphs hand the SAME body its two
2558/// modes without a second copy of the math.
2559pub(crate) struct FaLayerArgs<'a> {
2560    /// [T] per-row positions (device): rope reads them row-indexed; the seqs twins read
2561    /// them per-z (append slot = pos, T_kv = pos + 1).
2562    pub pos_d: &'a CudaSlice<i32>,
2563    /// Verify-level lazy per-row 1-element position buffers — only the per-row fallback
2564    /// arm builds/uses them (graph mode refuses that arm).
2565    pub pos_rows: &'a mut Option<Vec<CudaSlice<i32>>>,
2566    pub pos0: usize,
2567    pub seqs_append: bool,
2568    pub batch_fa_on: bool,
2569    /// Some((kv pointer table, offset-in-u64s, rung_end)) = captured-graph mode.
2570    pub graph_cap: Option<(&'a CudaSlice<u64>, usize, usize)>,
2571    /// ROUND-STREAM (lane/draftcost-moe, v0.100 train merge): Some((token stream, device
2572    /// round counter)) routes the FA attend through the dc rows kernels and the Linear
2573    /// mixer through `linear_attn_verify_t` (the stream arms the old inline body carried).
2574    /// Never armed together with `graph_cap` (the verify-level merge guard refuses).
2575    pub stream: Option<(&'a CudaSlice<u32>, &'a CudaSlice<i32>)>,
2576    /// VerifyCkpt for the stream-Linear arm's GdnStash install; None in graph mode and
2577    /// for FA layers that never touch it.
2578    pub ckpt: Option<&'a mut VerifyCkpt>,
2579}
2580
2581// SAFETY: `CudaGraph` is not marked Send by cudarc because its raw driver handles carry
2582// no automatic trait; CUDA driver graph handles are context-scoped rather than
2583// OS-thread-affine (the SpecPipeSessionPtr precedent above). The ctx lives in
2584// `HybridModel::dspark_vgraphs` behind a Mutex and every touch happens on the engine's
2585// single decode-stream thread.
2586unsafe impl Send for DsparkVerifyGraphs {}
2587
2588impl DsparkVerifyGraphs {
2589    /// Live capture count (segment + full graphs) — the denominator of
2590    /// [`dspark_vg_debt_projection`]'s observed bytes/capture mean.
2591    pub(crate) fn captures(&self) -> usize {
2592        self.graphs.len() + self.full.len()
2593    }
2594
2595    /// Take the marginal-growth debt reading and record this observation for the next one.
2596    /// Called under the pool mutex by `HybridModel::dspark_vg_admission_debt`.
2597    pub(crate) fn admission_debt(&mut self, reserved_bytes: usize) -> usize {
2598        let captures = self.captures();
2599        let debt =
2600            dspark_vg_debt_projection(captures, dspark_vg_cap(), reserved_bytes, self.debt_obs);
2601        if captures > 0 {
2602            match self.debt_obs {
2603                Some((c0, _)) if captures <= c0 => {}
2604                _ => self.debt_obs = Some((captures, reserved_bytes)),
2605            }
2606        }
2607        debt
2608    }
2609
2610    /// Build for this cache's shape. None when there are no linear layers, sizes are
2611    /// non-uniform, or the trunk keeps a gemma4 config (never on the qwen35 family).
2612    pub(crate) fn new(
2613        e: &Engine,
2614        cache: &Cache,
2615        t_max: usize,
2616        n_embd: usize,
2617    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
2618        let lin: Vec<usize> = (0..cache.recur.len())
2619            .filter(|&il| cache.recur[il].is_some())
2620            .collect();
2621        if lin.is_empty() || t_max < 2 {
2622            return Ok(None);
2623        }
2624        let first = cache.recur[lin[0]].as_ref().unwrap();
2625        let (conv_words, ssm_words) = (first.conv_state.len(), first.ssm_state.len());
2626        for &il in &lin {
2627            let rl = cache.recur[il].as_ref().unwrap();
2628            if rl.conv_state.len() != conv_words || rl.ssm_state.len() != ssm_words {
2629                return Ok(None);
2630            }
2631        }
2632        let n = lin.len();
2633        let mut lin_pos = std::collections::HashMap::with_capacity(n);
2634        for (k, &il) in lin.iter().enumerate() {
2635            lin_pos.insert(il, k);
2636        }
2637        // longest run of consecutive linear layers (save-scratch sizing)
2638        let mut max_run = 1usize;
2639        let mut run = 1usize;
2640        for w in lin.windows(2) {
2641            if w[1] == w[0] + 1 {
2642                run += 1;
2643                max_run = max_run.max(run);
2644            } else {
2645                run = 1;
2646            }
2647        }
2648        let rows = t_max - 1;
2649        let mut stash_conv = Vec::with_capacity(n);
2650        let mut stash_ssm = Vec::with_capacity(n);
2651        for _ in 0..n {
2652            stash_conv.push(e.uninit(rows * conv_words)?);
2653            stash_ssm.push(e.uninit(rows * ssm_words)?);
2654        }
2655        let host_table = vec![0u64; n * 6];
2656        let table_all = e.htod_u64(&host_table)?;
2657        // slice 4c: full-attention census for the full-verify graphs.
2658        let fa: Vec<usize> = (0..cache.kv.len())
2659            .filter(|&il| cache.kv[il].is_some())
2660            .collect();
2661        let mut fa_pos = std::collections::HashMap::with_capacity(fa.len());
2662        for (k, &il) in fa.iter().enumerate() {
2663            fa_pos.insert(il, k);
2664        }
2665        let n_layers = cache.kv.len().max(cache.recur.len());
2666        // exactly one of (linear state, kv cache) per layer — no MLA/unknown mixers.
2667        let walk_uniform = (0..n_layers).all(|il| {
2668            cache.recur.get(il).is_some_and(|r| r.is_some())
2669                != cache.kv.get(il).is_some_and(|k| k.is_some())
2670        });
2671        // Contiguous covered prefix: the largest n such that every layer in [0, n) is
2672        // linear or full-attention. The TRUNK walk is [0, layers.len()) and the cache
2673        // vecs can carry EXTRA state slots past it (the q38 export keeps the MTP head
2674        // layer's kv at the tail — hi == lin+fa never held, the s4c battery's zero
2675        // 'full' captures). The full-graph guard is walk coverage, not slot arithmetic.
2676        let covered = (0..n_layers)
2677            .take_while(|il| lin_pos.contains_key(il) || fa_pos.contains_key(il))
2678            .count();
2679        let t_cap = t_max;
2680        let fa_host_table = vec![0u64; fa.len() * 2 * t_cap];
2681        let fa_table = e.htod_u64(&fa_host_table)?;
2682        Ok(Some(Self {
2683            lin,
2684            lin_pos,
2685            table_all,
2686            host_table,
2687            stash_conv,
2688            stash_ssm,
2689            conv_words,
2690            ssm_words,
2691            stage: std::collections::HashMap::new(),
2692            tap_bufs: std::collections::HashMap::new(),
2693            graphs: std::collections::HashMap::new(),
2694            save_conv: e.uninit(n * conv_words)?,
2695            save_ssm: e.uninit(n * ssm_words)?,
2696            max_run,
2697            n_embd,
2698            round_slab: false,
2699            fa,
2700            fa_pos,
2701            fa_table,
2702            fa_host_table,
2703            t_cap,
2704            pos_stage: std::collections::HashMap::new(),
2705            full: std::collections::HashMap::new(),
2706            covered,
2707            walk_uniform,
2708            debt_obs: None,
2709        }))
2710    }
2711
2712    /// Rebuild the pointer tables from the live handles (once per verify — the gdn
2713    /// ping-pong swaps the canonical/alt handles between rounds; a fresh generation's
2714    /// cache buffers land at new addresses; a stale table would read the wrong state).
2715    pub(crate) fn refresh_tables(
2716        &mut self,
2717        e: &Engine,
2718        cache: &Cache,
2719    ) -> Result<(), Box<dyn std::error::Error>> {
2720        use cudarc::driver::DevicePtr;
2721        {
2722            let s = &e.gpu.stream();
2723            for (k, &il) in self.lin.iter().enumerate() {
2724                let rl = cache.recur[il].as_ref().unwrap();
2725                let (pc, _g0) = rl.conv_state.device_ptr(s);
2726                let (p0, _g1) = rl.ssm_state.device_ptr(s);
2727                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
2728                let o = k * 6;
2729                self.host_table[o] = pc as u64;
2730                self.host_table[o + 1] = p0 as u64;
2731                self.host_table[o + 2] = p1 as u64;
2732                self.host_table[o + 3] = pc as u64;
2733                self.host_table[o + 4] = p1 as u64;
2734                self.host_table[o + 5] = p0 as u64;
2735            }
2736            for (k, &il) in self.fa.iter().enumerate() {
2737                let kvl = cache.kv[il].as_ref().unwrap();
2738                let (pk, _g0) = kvl.k.device_ptr(s);
2739                let (pv, _g1) = kvl.v.device_ptr(s);
2740                let o = k * 2 * self.t_cap;
2741                for z in 0..self.t_cap {
2742                    self.fa_host_table[o + 2 * z] = pk as u64;
2743                    self.fa_host_table[o + 2 * z + 1] = pv as u64;
2744                }
2745            }
2746        }
2747        e.htod_u64_into(&self.host_table, &mut self.table_all)?;
2748        if !self.fa_host_table.is_empty() {
2749            e.htod_u64_into(&self.fa_host_table, &mut self.fa_table)?;
2750        }
2751        Ok(())
2752    }
2753
2754    /// Slice 4c eligibility: Some(rung_end) when this round can replay (or capture) a
2755    /// full-verify graph — the whole walk [lo, hi) is covered, every layer is linear or
2756    /// full-attention, and ALL of the round's per-row t_kv values take the v4-seqs arm
2757    /// on ONE `fa_split_keys` ladder step that the rung also sits on (the straddle law;
2758    /// both gates are t_kv intervals, so ends-inside means all-inside). The rung is the
2759    /// round's next power of two — grid/partial sizing only (`n_splits_max` is pure
2760    /// stride; splits >= ns_eff write the empty partial the combine never reads), so one
2761    /// captured graph is bit-identical for every round the rung covers.
2762    #[allow(clippy::too_many_arguments)]
2763    pub(crate) fn full_rung(
2764        &self,
2765        model: &crate::hybrid::HybridModel,
2766        cache: &Cache,
2767        lo: usize,
2768        hi: usize,
2769        t: usize,
2770        seqs_arms_on: bool,
2771    ) -> Option<usize> {
2772        if std::env::var("MEMRA_DSPARK_FULLG_DEBUG").as_deref() == Ok("1") {
2773            static ONCE: std::sync::Once = std::sync::Once::new();
2774            let len0 = self
2775                .fa
2776                .first()
2777                .and_then(|&il| cache.kv[il].as_ref())
2778                .map(|k| k.len);
2779            ONCE.call_once(|| {
2780                eprintln!(
2781                    "[fullg-debug] walk_uniform={} covered={} seqs_arms_on={} fa_rows_on={} t={} lo={} hi={} lin={} fa={} t_cap={} len0={:?}",
2782                    self.walk_uniform, self.covered, seqs_arms_on, dspark_fa_rows_on(), t, lo, hi,
2783                    self.lin.len(), self.fa.len(), self.t_cap, len0
2784                );
2785            });
2786        }
2787        if !self.walk_uniform
2788            || !seqs_arms_on
2789            || !dspark_fa_rows_on()
2790            || t < 2
2791            || lo != 0
2792            || hi > self.covered
2793            || t > self.t_cap
2794            || self.fa.is_empty()
2795        {
2796            return None;
2797        }
2798        let cfg = &model.cfg;
2799        let head_dim_global = cfg.head_dim_k as usize;
2800        let nkv = cfg.n_head_kv as usize;
2801        let kvl0 = cache.kv[self.fa[0]].as_ref().unwrap();
2802        // the z-batched twins read stacked rows at the cache's kv dims — must equal the
2803        // projection stride (the body's guard, hoisted so ineligible models fall back
2804        // instead of refusing mid-capture).
2805        let geom = cfg.full_attention_geometry_at(self.fa[0] as u32);
2806        let kv_dim = geom.n_head_kv as usize * geom.head_dim_k as usize;
2807        if kvl0.kv_dim_k != kv_dim || kvl0.kv_dim_v != kv_dim {
2808            return None;
2809        }
2810        let len0 = kvl0.len;
2811        let (t_kv_first, t_kv_last) = (len0 + 1, len0 + t);
2812        if !crate::fa_seqs_eligible(t_kv_first, head_dim_global)
2813            || !crate::fa_seqs_eligible(t_kv_last, head_dim_global)
2814            || crate::fa_split_keys(t_kv_first, nkv) != crate::fa_split_keys(t_kv_last, nkv)
2815        {
2816            return None;
2817        }
2818        let rung = t_kv_last.next_power_of_two().max(256);
2819        if crate::fa_split_keys(rung, nkv) != crate::fa_split_keys(t_kv_last, nkv) {
2820            return None;
2821        }
2822        Some(rung)
2823    }
2824
2825    /// Run the WHOLE verify walk [lo, hi) as one captured graph at (vt=t, rung): stage
2826    /// the residual + refresh the per-vt position staging, capture on first encounter
2827    /// (2 executing warmups bracketed by a full linear-state save/restore; KV warmup
2828    /// appends write the exact slots the replay writes — idempotent), launch, then apply
2829    /// the host bookkeeping the captured body skipped (per-linear-layer parity swap for
2830    /// odd t, per-fa-layer len bump). Returns the fresh residual.
2831    #[allow(clippy::too_many_arguments)]
2832    pub(crate) fn run_full(
2833        &mut self,
2834        model: &crate::hybrid::HybridModel,
2835        e: &Engine,
2836        lo: usize,
2837        hi: usize,
2838        x: &CudaSlice<f32>,
2839        t: usize,
2840        pos0: usize,
2841        rung: usize,
2842        cache: &mut Cache,
2843    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2844        let n_embd = self.n_embd;
2845        if !self.stage.contains_key(&t) {
2846            let xin = e.uninit(t * n_embd)?;
2847            let xout = e.uninit(t * n_embd)?;
2848            self.stage.insert(t, (xin, xout));
2849        }
2850        if !self.pos_stage.contains_key(&t) {
2851            self.pos_stage.insert(t, e.htod_i32(&vec![0i32; t])?);
2852        }
2853        // Per-round refresh: position contents + input staging (both addresses are baked
2854        // by the captured bodies; only their CONTENTS change round to round).
2855        {
2856            let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
2857            let pb = self.pos_stage.get_mut(&t).unwrap();
2858            e.htod_i32_into(pb, &pos_host)?;
2859            let (xin, _) = self.stage.get_mut(&t).unwrap();
2860            e.copy_into(xin, 0, x, t * n_embd)?;
2861        }
2862        let key = (t, rung, hi);
2863        if !self.full.contains_key(&key) {
2864            // The warmups EXECUTE the whole walk on live state — save every linear
2865            // layer's conv + canonical ssm first, restore after (KV needs no restore:
2866            // graph mode never bumps host lens and the appends write this round's own
2867            // slots).
2868            for (k, &il) in self.lin.iter().enumerate() {
2869                let rl = cache.recur[il].as_ref().unwrap();
2870                e.copy_into(
2871                    &mut self.save_conv,
2872                    k * self.conv_words,
2873                    &rl.conv_state,
2874                    self.conv_words,
2875                )?;
2876                e.copy_into(
2877                    &mut self.save_ssm,
2878                    k * self.ssm_words,
2879                    &rl.ssm_state,
2880                    self.ssm_words,
2881                )?;
2882            }
2883            let (graph, keeper) = {
2884                let table_all = &self.table_all;
2885                let lin_pos = &self.lin_pos;
2886                let fa_pos = &self.fa_pos;
2887                let fa_table = &self.fa_table;
2888                let t_cap = self.t_cap;
2889                let stash_conv = &mut self.stash_conv;
2890                let stash_ssm = &mut self.stash_ssm;
2891                let pos_d: &CudaSlice<i32> = &self.pos_stage[&t];
2892                let (xin, xout) = self
2893                    .stage
2894                    .get_mut(&t)
2895                    .map(|(a, b)| (&*a, b))
2896                    .expect("stage bucket created above");
2897                let cache_ref: &mut Cache = cache;
2898                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
2899                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
2900                } else {
2901                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
2902                };
2903                e.capture_graph_retained_flags(iflag, move |e| {
2904                    let mut xc: Option<CudaSlice<f32>> = None;
2905                    for il in lo..hi {
2906                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
2907                        let nx = if let Some(&k) = lin_pos.get(&il) {
2908                            model.qwen35_tparallel_linear_layer(
2909                                e,
2910                                il,
2911                                xr,
2912                                t,
2913                                cache_ref,
2914                                None,
2915                                Some((&mut stash_conv[k], &mut stash_ssm[k])),
2916                                Some((table_all, k * 6)),
2917                            )?
2918                        } else if let Some(&kf) = fa_pos.get(&il) {
2919                            let mut no_rows: Option<Vec<CudaSlice<i32>>> = None;
2920                            model.qwen35_tparallel_fa_layer(
2921                                e,
2922                                il,
2923                                xr,
2924                                t,
2925                                cache_ref,
2926                                FaLayerArgs {
2927                                    pos_d,
2928                                    pos_rows: &mut no_rows,
2929                                    pos0,
2930                                    seqs_append: true,
2931                                    batch_fa_on: true,
2932                                    graph_cap: Some((fa_table, kf * 2 * t_cap, rung)),
2933                                    stream: None,
2934                                    ckpt: None,
2935                                },
2936                            )?
2937                        } else {
2938                            return Err(format!(
2939                                "run_full: layer {il} is neither linear nor full-attention"
2940                            )
2941                            .into());
2942                        };
2943                        xc = Some(nx);
2944                    }
2945                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
2946                    Ok(())
2947                })?
2948            };
2949            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
2950            // is odd -> 3 runs = net one swap), then restore the device state the
2951            // warmups consumed (walk scope only — layers past hi never executed). The
2952            // launch below then behaves exactly like one run.
2953            if t % 2 == 1 {
2954                for &il in &self.lin {
2955                    if il < lo || il >= hi {
2956                        continue;
2957                    }
2958                    let rl = cache.recur[il].as_mut().unwrap();
2959                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
2960                }
2961            }
2962            for (k, &il) in self.lin.iter().enumerate() {
2963                if il < lo || il >= hi {
2964                    continue;
2965                }
2966                let rl = cache.recur[il].as_mut().unwrap();
2967                let (cw, sw) = (self.conv_words, self.ssm_words);
2968                {
2969                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
2970                    let win = sv.slice(k * cw..(k + 1) * cw);
2971                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
2972                }
2973                {
2974                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
2975                    let win = sv.slice(k * sw..(k + 1) * sw);
2976                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
2977                }
2978            }
2979            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
2980                if let Ok(c) = crate::graph_update::node_census(&graph) {
2981                    eprintln!("[dspark-vg-census] full vt={t} rung={rung} {c:?}");
2982                }
2983            }
2984            self.full.insert(
2985                key,
2986                DsparkSegGraph {
2987                    graph,
2988                    _keeper: keeper,
2989                },
2990            );
2991        }
2992        self.full[&key].graph.launch()?;
2993        // Host bookkeeping for the replayed body (captured host code does not re-run):
2994        // gdn parity swap per linear layer (t odd), kv len bump per fa layer — scoped
2995        // to the WALK [lo, hi): the cache can carry extra state slots past it (the MTP
2996        // head layer's kv) that the walk never touches.
2997        if t % 2 == 1 {
2998            for &il in &self.lin {
2999                if il < lo || il >= hi {
3000                    continue;
3001                }
3002                let rl = cache.recur[il].as_mut().unwrap();
3003                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3004            }
3005        }
3006        for &il in &self.fa {
3007            if il < lo || il >= hi {
3008                continue;
3009            }
3010            cache.kv[il].as_mut().unwrap().len += t;
3011        }
3012        let (_, xout) = self.stage.get(&t).unwrap();
3013        let mut out = e.uninit(t * n_embd)?;
3014        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3015        Ok(out)
3016    }
3017
3018    /// Run layers [start, end) (all linear) as one captured graph at this vt: stage the
3019    /// residual into the bucket's x_in, capture on first encounter (2 executing warmups
3020    /// bracketed by a segment state save/restore), launch, then apply the host parity
3021    /// bookkeeping the captured body would have done. Returns the fresh residual.
3022    #[allow(clippy::too_many_arguments)]
3023    fn run_segment(
3024        &mut self,
3025        model: &crate::hybrid::HybridModel,
3026        e: &Engine,
3027        start: usize,
3028        end: usize,
3029        x: &CudaSlice<f32>,
3030        t: usize,
3031        cache: &mut Cache,
3032    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3033        let n_embd = self.n_embd;
3034        debug_assert!(end - start <= self.max_run);
3035        if !self.stage.contains_key(&t) {
3036            let xin = e.uninit(t * n_embd)?;
3037            let xout = e.uninit(t * n_embd)?;
3038            self.stage.insert(t, (xin, xout));
3039        }
3040        // Stage the residual at the bucket's baked input address.
3041        {
3042            let (xin, _) = self.stage.get_mut(&t).unwrap();
3043            e.copy_into(xin, 0, x, t * n_embd)?;
3044        }
3045        let key = (start, t);
3046        if !self.graphs.contains_key(&key) {
3047            // The 2 warmups EXECUTE the segment on live state — save conv + the canonical
3048            // ssm of every segment layer first, restore after, so the graph's first real
3049            // launch starts from the exact pre-round state (bytes gated e2e).
3050            for (k, il) in (start..end).enumerate() {
3051                let rl = cache.recur[il].as_ref().unwrap();
3052                e.copy_into(
3053                    &mut self.save_conv,
3054                    k * self.conv_words,
3055                    &rl.conv_state,
3056                    self.conv_words,
3057                )?;
3058                e.copy_into(
3059                    &mut self.save_ssm,
3060                    k * self.ssm_words,
3061                    &rl.ssm_state,
3062                    self.ssm_words,
3063                )?;
3064            }
3065            let (graph, keeper) = {
3066                let table_all = &self.table_all;
3067                let lin_pos = &self.lin_pos;
3068                let stash_conv = &mut self.stash_conv;
3069                let stash_ssm = &mut self.stash_ssm;
3070                let (xin, xout) = self
3071                    .stage
3072                    .get_mut(&t)
3073                    .map(|(a, b)| (&*a, b))
3074                    .expect("stage bucket created above");
3075                let cache_ref: &mut Cache = cache;
3076                // Slice 4 (fa-execupdate lane): USE_NODE_PRIORITY instead of
3077                // AUTO_FREE_ON_LAUNCH. The slice-3 measured limiter was AUTO_FREE's
3078                // launch-time mem-pool scan — 25.6 us per cuGraphLaunch x 16 segments
3079                // = ~0.41 ms/round, most of the eager-launch savings. The captured
3080                // body's cuMemAllocAsync transients are BALANCED by in-graph frees
3081                // (every transient drops inside the capture region — the generic
3082                // capture path's census precedent, 1589/1589), so AUTO_FREE has
3083                // nothing to reclaim and the graph is legal to instantiate without
3084                // it; PRIORITY is the flag the gemma slotted door ships for exactly
3085                // this reason (both alternatives drop the scan; UPLOAD via
3086                // cuGraphInstantiateWithFlags is WithParams-only and refused).
3087                // MEMRA_DSPARK_VG_AUTOFREE=1 reverts; MEMRA_GRAPH_CENSUS=1 prints
3088                // the node census at capture (the ALLOC==FREE receipt).
3089                let iflag = if std::env::var("MEMRA_DSPARK_VG_AUTOFREE").as_deref() == Ok("1") {
3090                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH
3091                } else {
3092                    cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
3093                };
3094                e.capture_graph_retained_flags(iflag, move |e| {
3095                    let mut xc: Option<CudaSlice<f32>> = None;
3096                    for il in start..end {
3097                        let k = lin_pos[&il];
3098                        let xr: &CudaSlice<f32> = xc.as_ref().unwrap_or(xin);
3099                        let nx = model.qwen35_tparallel_linear_layer(
3100                            e,
3101                            il,
3102                            xr,
3103                            t,
3104                            cache_ref,
3105                            None,
3106                            Some((&mut stash_conv[k], &mut stash_ssm[k])),
3107                            Some((table_all, k * 6)),
3108                        )?;
3109                        xc = Some(nx);
3110                    }
3111                    e.copy_into(xout, 0, xc.as_ref().unwrap(), t * n_embd)?;
3112                    Ok(())
3113                })?
3114            };
3115            // Undo the net host parity motion of the 3 body runs (each run swaps iff t
3116            // is odd -> 3 runs = net one swap), then restore the device state the
3117            // warmups consumed. The launch below then behaves exactly like one run.
3118            if t % 2 == 1 {
3119                for il in start..end {
3120                    let rl = cache.recur[il].as_mut().unwrap();
3121                    std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3122                }
3123            }
3124            for (k, il) in (start..end).enumerate() {
3125                let rl = cache.recur[il].as_mut().unwrap();
3126                let (cw, sw) = (self.conv_words, self.ssm_words);
3127                {
3128                    let sv = e.view(&self.save_conv, self.lin.len() * cw);
3129                    let win = sv.slice(k * cw..(k + 1) * cw);
3130                    e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
3131                }
3132                {
3133                    let sv = e.view(&self.save_ssm, self.lin.len() * sw);
3134                    let win = sv.slice(k * sw..(k + 1) * sw);
3135                    e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
3136                }
3137            }
3138            if std::env::var("MEMRA_GRAPH_CENSUS").as_deref() == Ok("1") {
3139                if let Ok(c) = crate::graph_update::node_census(&graph) {
3140                    eprintln!("[dspark-vg-census] seg={start}..{end} vt={t} {c:?}");
3141                }
3142            }
3143            self.graphs.insert(
3144                key,
3145                DsparkSegGraph {
3146                    graph,
3147                    _keeper: keeper,
3148                },
3149            );
3150        }
3151        self.graphs[&key].graph.launch()?;
3152        // Host parity bookkeeping for the replayed body (the captured host swaps do not
3153        // re-run at replay).
3154        if t % 2 == 1 {
3155            for il in start..end {
3156                let rl = cache.recur[il].as_mut().unwrap();
3157                std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3158            }
3159        }
3160        let (_, xout) = self.stage.get(&t).unwrap();
3161        let mut out = e.uninit(t * n_embd)?;
3162        e.copy_into(&mut out, 0, xout, t * n_embd)?;
3163        Ok(out)
3164    }
3165
3166    /// Pool freeze check (`dspark_vg_cap`): below the ceiling new keys may capture.
3167    fn can_capture(&self) -> bool {
3168        self.graphs.len() + self.full.len() < dspark_vg_cap()
3169    }
3170
3171    /// Round-atomic segment-door readiness: TRUE when this round's walk can ride the
3172    /// per-(segment, vt) graphs without a NEW capture past the pool ceiling — every
3173    /// linear run in [lo, hi) already has its (run_start, t) key, or capture is still
3174    /// allowed. FALSE sends the WHOLE round down the eager cols-ckpt walk: a partial
3175    /// refusal would stash some layers in the ctx slabs and others in the round's cols
3176    /// while one commit reads only one of them.
3177    pub(crate) fn segments_ready(
3178        &self,
3179        model: &crate::hybrid::HybridModel,
3180        lo: usize,
3181        hi: usize,
3182        t: usize,
3183    ) -> bool {
3184        if self.can_capture() {
3185            return true;
3186        }
3187        let mut il = lo;
3188        while il < hi {
3189            if matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3190                let start = il;
3191                while il < hi && matches!(model.layers[il].mixer, Mixer::Linear(_)) {
3192                    il += 1;
3193                }
3194                if !self.graphs.contains_key(&(start, t)) {
3195                    return false;
3196                }
3197            } else {
3198                il += 1;
3199            }
3200        }
3201        true
3202    }
3203
3204    /// Widest verify window this pool was built for. A caller whose round exceeds it must
3205    /// take the eager walk: the stash slabs hold `t_capacity() - 1` column rows, and slicing
3206    /// past them is a panic rather than a refusal.
3207    pub(crate) fn t_capacity(&self) -> usize {
3208        self.t_cap
3209    }
3210
3211    /// Slab row (conv, ssm) device pointers + lengths for the commit restore of column
3212    /// `row` (0-based) of layer `il`. None for non-linear layers.
3213    pub(crate) fn slab_row(
3214        &self,
3215        e: &Engine,
3216        il: usize,
3217        row: usize,
3218    ) -> Option<(u64, u64, usize, usize)> {
3219        use cudarc::driver::DevicePtr;
3220        let k = *self.lin_pos.get(&il)?;
3221        let s = &e.gpu.stream();
3222        let (pc, _g0) = self.stash_conv[k].device_ptr(s);
3223        let (ps, _g1) = self.stash_ssm[k].device_ptr(s);
3224        Some((
3225            pc as u64 + (row * self.conv_words * 4) as u64,
3226            ps as u64 + (row * self.ssm_words * 4) as u64,
3227            self.conv_words,
3228            self.ssm_words,
3229        ))
3230    }
3231}
3232
3233impl VerifyCkpt {
3234    fn new(n_layer: usize) -> Self {
3235        VerifyCkpt {
3236            gdn: (0..n_layer).map(|_| None).collect(),
3237            cols: (0..n_layer).map(|_| None).collect(),
3238        }
3239    }
3240}
3241
3242/// The stage-0/TX half of one PP verify. The boundary slot is the ownership token: stage 1
3243/// consumes exactly the slot selected by `tx()` / `tx_pipelined()`, never a slot inferred from
3244/// a logical round number.
3245struct VerifyBoundaryTicket {
3246    rt: &'static crate::pp::PpNRt,
3247    caller_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3248    slot: usize,
3249    pos0: usize,
3250    t: usize,
3251    payload: usize,
3252    n_st: usize,
3253    pipelined: bool,
3254    pp_anatomy: bool,
3255    pp_started: std::time::Instant,
3256    reverse_ms: f64,
3257    stage0_ms: f64,
3258    tx_ms: f64,
3259    trace: Option<SpecPipeTraceCtx>,
3260}
3261
3262/// Explicit OPTIPIPE diagnostic control. Forced modes are set only by `optipipe-gate`; the
3263/// increment-2 controller can also be armed by the server's fresh-process research door.
3264#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3265pub enum OptiForkGateMode {
3266    Disabled,
3267    Hit,
3268    Miss,
3269    Alternate,
3270    Abort,
3271    Controller,
3272}
3273
3274static OPTI_FORK_GATE_MODE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
3275static OPTI_CONTROLLER_THRESHOLD: std::sync::atomic::AtomicU32 =
3276    std::sync::atomic::AtomicU32::new(0);
3277static OPTI_FORK_ATTEMPTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3278static OPTI_FORK_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3279static OPTI_FORK_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3280static OPTI_FORK_ABORT_DRAINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3281static OPTI_FORK_REFUSALS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3282static OPTI_GATE_CHECKS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3283static OPTI_GATE_ADMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3284static OPTI_GATE_REJECTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3285static OPTI_RECONCILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3286static OPTI_WASTED_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3287    std::sync::atomic::AtomicU64::new(0);
3288static OPTI_SHADOW_DRAFT_TOKENS: std::sync::atomic::AtomicU64 =
3289    std::sync::atomic::AtomicU64::new(0);
3290static OPTI_BREAKER_TRIPS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
3291
3292impl OptiForkGateMode {
3293    fn code(self) -> u8 {
3294        match self {
3295            Self::Disabled => 0,
3296            Self::Hit => 1,
3297            Self::Miss => 2,
3298            Self::Alternate => 3,
3299            Self::Abort => 4,
3300            Self::Controller => 5,
3301        }
3302    }
3303
3304    fn configured() -> Self {
3305        match OPTI_FORK_GATE_MODE.load(std::sync::atomic::Ordering::Relaxed) {
3306            1 => Self::Hit,
3307            2 => Self::Miss,
3308            3 => Self::Alternate,
3309            4 => Self::Abort,
3310            5 => Self::Controller,
3311            _ => Self::Disabled,
3312        }
3313    }
3314
3315    fn action(self, generation: u64) -> OptiForkAction {
3316        match self {
3317            Self::Hit => OptiForkAction::Hit,
3318            Self::Miss => OptiForkAction::Miss,
3319            Self::Alternate if generation & 1 == 0 => OptiForkAction::Hit,
3320            Self::Alternate => OptiForkAction::Miss,
3321            Self::Abort => OptiForkAction::Abort,
3322            Self::Disabled | Self::Controller => {
3323                unreachable!("non-forced mode cannot choose a forced fork action")
3324            }
3325        }
3326    }
3327
3328    fn is_forced(self) -> bool {
3329        matches!(self, Self::Hit | Self::Miss | Self::Alternate | Self::Abort)
3330    }
3331}
3332
3333/// Arm or disarm the forced harness. Serving uses only `set_optipipe_controller_threshold`.
3334pub fn set_optipipe_gate_mode(mode: OptiForkGateMode) {
3335    OPTI_FORK_GATE_MODE.store(mode.code(), std::sync::atomic::Ordering::Relaxed);
3336}
3337
3338/// Arm the increment-2 diagnostic controller. The threshold applies to the uncalibrated
3339/// two-token draft-probability product. Serving can call this only through its explicit
3340/// fresh-process research door; the absent-door default remains byte-for-byte disabled.
3341pub fn set_optipipe_controller_threshold(threshold: f32) {
3342    assert!(threshold.is_finite() && (0.0..=1.0).contains(&threshold));
3343    OPTI_CONTROLLER_THRESHOLD.store(threshold.to_bits(), std::sync::atomic::Ordering::Relaxed);
3344    set_optipipe_gate_mode(OptiForkGateMode::Controller);
3345}
3346
3347#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3348pub struct OptiForkGateStats {
3349    pub attempts: u64,
3350    pub hits: u64,
3351    pub misses: u64,
3352    pub abort_drains: u64,
3353    pub refusals: u64,
3354    pub gate_checks: u64,
3355    pub gate_admits: u64,
3356    pub gate_rejects: u64,
3357    pub reconciles: u64,
3358    pub wasted_draft_tokens: u64,
3359    pub shadow_draft_tokens: u64,
3360    pub breaker_trips: u64,
3361}
3362
3363#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
3364pub struct OptiForkStateIdentity {
3365    pub trunk_kv_bytes: usize,
3366    pub recurrent_bytes: usize,
3367    pub scratch_kv_bytes: usize,
3368    pub hidden_bytes: usize,
3369}
3370
3371pub fn reset_optipipe_gate_stats() {
3372    for counter in [
3373        &OPTI_FORK_ATTEMPTS,
3374        &OPTI_FORK_HITS,
3375        &OPTI_FORK_MISSES,
3376        &OPTI_FORK_ABORT_DRAINS,
3377        &OPTI_FORK_REFUSALS,
3378        &OPTI_GATE_CHECKS,
3379        &OPTI_GATE_ADMITS,
3380        &OPTI_GATE_REJECTS,
3381        &OPTI_RECONCILES,
3382        &OPTI_WASTED_DRAFT_TOKENS,
3383        &OPTI_SHADOW_DRAFT_TOKENS,
3384        &OPTI_BREAKER_TRIPS,
3385    ] {
3386        counter.store(0, std::sync::atomic::Ordering::Relaxed);
3387    }
3388}
3389
3390pub fn optipipe_gate_stats() -> OptiForkGateStats {
3391    let load = |v: &std::sync::atomic::AtomicU64| v.load(std::sync::atomic::Ordering::Relaxed);
3392    OptiForkGateStats {
3393        attempts: load(&OPTI_FORK_ATTEMPTS),
3394        hits: load(&OPTI_FORK_HITS),
3395        misses: load(&OPTI_FORK_MISSES),
3396        abort_drains: load(&OPTI_FORK_ABORT_DRAINS),
3397        refusals: load(&OPTI_FORK_REFUSALS),
3398        gate_checks: load(&OPTI_GATE_CHECKS),
3399        gate_admits: load(&OPTI_GATE_ADMITS),
3400        gate_rejects: load(&OPTI_GATE_REJECTS),
3401        reconciles: load(&OPTI_RECONCILES),
3402        wasted_draft_tokens: load(&OPTI_WASTED_DRAFT_TOKENS),
3403        shadow_draft_tokens: load(&OPTI_SHADOW_DRAFT_TOKENS),
3404        breaker_trips: load(&OPTI_BREAKER_TRIPS),
3405    }
3406}
3407
3408#[derive(Clone, Copy, Debug)]
3409struct OptiControllerPolicy {
3410    threshold: f32,
3411    consecutive_misses: u8,
3412    breaker_tripped: bool,
3413}
3414
3415impl OptiControllerPolicy {
3416    fn configured() -> Self {
3417        Self {
3418            threshold: f32::from_bits(
3419                OPTI_CONTROLLER_THRESHOLD.load(std::sync::atomic::Ordering::Relaxed),
3420            ),
3421            consecutive_misses: 0,
3422            breaker_tripped: false,
3423        }
3424    }
3425
3426    fn admit(&self, q_proxy: f32) -> bool {
3427        q_proxy.is_finite()
3428            && (0.0..=1.0).contains(&q_proxy)
3429            && (self.threshold == 0.0 || (!self.breaker_tripped && q_proxy >= self.threshold))
3430    }
3431
3432    /// Returns true exactly when this resolution newly trips the three-miss breaker.
3433    fn resolve(&mut self, hit: bool) -> bool {
3434        // q*=0 is the lane's explicit unconditional measurement arm. Its purpose is to price
3435        // every optimistic opportunity, so the safety breaker is measured separately and must
3436        // not silently turn this arm into "three attempts then serial".
3437        if self.threshold == 0.0 {
3438            self.consecutive_misses = 0;
3439            return false;
3440        }
3441        if hit {
3442            self.consecutive_misses = 0;
3443            return false;
3444        }
3445        self.consecutive_misses = self.consecutive_misses.saturating_add(1);
3446        if !self.breaker_tripped && self.consecutive_misses >= 3 {
3447            self.breaker_tripped = true;
3448            return true;
3449        }
3450        false
3451    }
3452}
3453
3454#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3455enum OptiForkAction {
3456    Hit,
3457    Miss,
3458    Abort,
3459}
3460
3461#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3462struct OptiForkGeneration {
3463    id: u64,
3464    slot: usize,
3465}
3466
3467#[derive(Default)]
3468struct OptiForkGenerationTracker {
3469    next: u64,
3470    live: [Option<u64>; 2],
3471}
3472
3473impl OptiForkGenerationTracker {
3474    fn reserve(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3475        let generation = OptiForkGeneration {
3476            id: self.next,
3477            slot: (self.next & 1) as usize,
3478        };
3479        if let Some(live) = self.live[generation.slot] {
3480            return Err(format!(
3481                "optipipe snapshot slot {} still owns generation {live}; refusing to overwrite it",
3482                generation.slot,
3483            )
3484            .into());
3485        }
3486        self.next += 1;
3487        self.live[generation.slot] = Some(generation.id);
3488        Ok(generation)
3489    }
3490
3491    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
3492        match self.live[generation.slot] {
3493            Some(id) if id == generation.id => {
3494                self.live[generation.slot] = None;
3495                Ok(())
3496            }
3497            other => Err(format!(
3498                "optipipe generation teardown mismatch: ticket={} slot={} live={other:?}",
3499                generation.id, generation.slot,
3500            )
3501            .into()),
3502        }
3503    }
3504}
3505
3506struct OptiForkSeedGeneration {
3507    h_seed: CudaSlice<f32>,
3508    fill_prev: CudaSlice<f32>,
3509    scratch_len: usize,
3510}
3511
3512/// Allocate or refresh one full checkpoint through the engine that owns each PP stage. The
3513/// generic cache helper accepts one device and therefore cannot copy GDN state split across
3514/// devices. KV lengths and position stay host metadata; only recurrent buffers need stage-local
3515/// device ownership.
3516fn opti_snapshot_stage_owned(
3517    e: &Engine,
3518    cache: &Cache,
3519    rt: &'static crate::pp::PpNRt,
3520    fence: &[usize],
3521) -> Result<crate::cache::CacheSnapshot, Box<dyn std::error::Error>> {
3522    let n = cache.kv.len();
3523    let mut snapshot = crate::cache::CacheSnapshot {
3524        kv_len: vec![None; n],
3525        tp_kv_len: vec![None; n],
3526        conv: (0..n).map(|_| None).collect(),
3527        ssm: (0..n).map(|_| None).collect(),
3528        pos: cache.pos,
3529    };
3530    opti_snapshot_stage_owned_into(e, cache, rt, fence, &mut snapshot)?;
3531    Ok(snapshot)
3532}
3533
3534fn opti_snapshot_stage_owned_into(
3535    e: &Engine,
3536    cache: &Cache,
3537    rt: &'static crate::pp::PpNRt,
3538    fence: &[usize],
3539    snapshot: &mut crate::cache::CacheSnapshot,
3540) -> Result<(), Box<dyn std::error::Error>> {
3541    if fence.len() != rt.n_stages() + 1
3542        || snapshot.kv_len.len() != cache.kv.len()
3543        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3544    {
3545        return Err("optipipe stage-owned snapshot shape mismatch".into());
3546    }
3547    for stage in 0..rt.n_stages() {
3548        opti_snapshot_one_stage_owned_into(e, cache, rt, fence, stage, snapshot)?;
3549    }
3550    snapshot.pos = cache.pos;
3551    Ok(())
3552}
3553
3554/// Refresh one PP stage of a checkpoint. Increment 2 uses this split form so stage 0's
3555/// optimistic post-N state is captured before N+1 stage 0 is queued, while stage 1's matching
3556/// post-N state is captured only after N stage 1 is enqueued. Calling the all-stage helper at
3557/// either point would capture one side of the fork at the wrong generation.
3558fn opti_snapshot_one_stage_owned_into(
3559    e: &Engine,
3560    cache: &Cache,
3561    rt: &'static crate::pp::PpNRt,
3562    fence: &[usize],
3563    stage: usize,
3564    snapshot: &mut crate::cache::CacheSnapshot,
3565) -> Result<(), Box<dyn std::error::Error>> {
3566    if fence.len() != rt.n_stages() + 1
3567        || snapshot.kv_len.len() != cache.kv.len()
3568        || snapshot.tp_kv_len.len() != cache.tp_kv.len()
3569        || stage >= rt.n_stages()
3570    {
3571        return Err("optipipe single-stage snapshot shape mismatch".into());
3572    }
3573    let _scope = rt.enter(stage);
3574    let owner = rt.engine(stage, e);
3575    for il in fence[stage]..fence[stage + 1] {
3576        snapshot.kv_len[il] = cache.kv[il].as_ref().map(|kv| kv.len);
3577        snapshot.tp_kv_len[il] = cache.tp_kv[il]
3578            .as_ref()
3579            .map(crate::tp::ResidentTpKvCache::committed_len);
3580        match &cache.recur[il] {
3581            Some(recur) => {
3582                match snapshot.conv[il].as_mut() {
3583                    Some(dst) => {
3584                        owner.copy_into(dst, 0, &recur.conv_state, recur.conv_state.len())?
3585                    }
3586                    None => snapshot.conv[il] = Some(owner.clone_dtod(&recur.conv_state)?),
3587                }
3588                match snapshot.ssm[il].as_mut() {
3589                    Some(dst) => {
3590                        owner.copy_into(dst, 0, &recur.ssm_state, recur.ssm_state.len())?
3591                    }
3592                    None => snapshot.ssm[il] = Some(owner.clone_dtod(&recur.ssm_state)?),
3593                }
3594            }
3595            None if snapshot.conv[il].is_some() || snapshot.ssm[il].is_some() => {
3596                return Err(
3597                    format!("optipipe stage-owned snapshot layer {il} changed shape").into(),
3598                );
3599            }
3600            None => {}
3601        }
3602    }
3603    snapshot.pos = cache.pos;
3604    Ok(())
3605}
3606
3607/// Increment-1 persistent fork state. Exactly two snapshot/seed slots alternate; a live ticket
3608/// names its generation and keeps teardown fail-closed. Only stage 0 is allowed to mutate before
3609/// resolve, so the reconcile tables and conditional restores are stage-local.
3610struct OptiForkState {
3611    mode: OptiForkGateMode,
3612    controller: Option<OptiControllerPolicy>,
3613    generations: OptiForkGenerationTracker,
3614    active_snapshot_slot: usize,
3615    alternate_snapshot: crate::cache::CacheSnapshot,
3616    seeds: [OptiForkSeedGeneration; 2],
3617    rt: &'static crate::pp::PpNRt,
3618    fence: [usize; 3],
3619    split: usize,
3620    len_ptrs: CudaSlice<u64>,
3621    saved_lens: CudaSlice<i32>,
3622    forced_acc: CudaSlice<u32>,
3623    valid: CudaSlice<u32>,
3624    stage0_stream: std::sync::Arc<cudarc::driver::CudaStream>,
3625    logical_payload_bytes: [usize; 2],
3626}
3627
3628struct OptiForkTicket {
3629    generation: OptiForkGeneration,
3630    boundary: Option<VerifyBoundaryTicket>,
3631    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3632    settled: bool,
3633}
3634
3635struct OptiControllerTicket {
3636    generation: OptiForkGeneration,
3637    boundary: Option<VerifyBoundaryTicket>,
3638    ckpt: Option<VerifyCkpt>,
3639    verify_tokens: [u32; 2],
3640    draft_prob: f32,
3641    eager_seed: Option<CudaSlice<f32>>,
3642    q_proxy: f32,
3643    scratch_len: usize,
3644    issued_at: std::time::Instant,
3645    drain: std::sync::Arc<cudarc::driver::CudaStream>,
3646    settled: bool,
3647}
3648
3649struct OptiControllerPrepared {
3650    verify_tokens: [u32; 2],
3651    draft_prob: f32,
3652    eager_seed: Option<CudaSlice<f32>>,
3653    q_proxy: f32,
3654    scratch_len: usize,
3655}
3656
3657impl OptiControllerTicket {
3658    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3659        self.boundary
3660            .take()
3661            .expect("controller boundary ticket already consumed")
3662    }
3663
3664    fn take_ckpt(&mut self) -> VerifyCkpt {
3665        self.ckpt
3666            .take()
3667            .expect("controller verify checkpoint already consumed")
3668    }
3669
3670    fn take_eager_seed(&mut self) -> Option<CudaSlice<f32>> {
3671        self.eager_seed.take()
3672    }
3673
3674    fn settle(&mut self) {
3675        self.settled = true;
3676    }
3677}
3678
3679impl Drop for OptiControllerTicket {
3680    fn drop(&mut self) {
3681        if !self.settled {
3682            let _ = self.drain.synchronize();
3683            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3684        }
3685    }
3686}
3687
3688impl OptiForkTicket {
3689    fn take_boundary(&mut self) -> VerifyBoundaryTicket {
3690        self.boundary
3691            .take()
3692            .expect("fork ticket boundary already consumed")
3693    }
3694
3695    fn settle(&mut self) {
3696        self.settled = true;
3697    }
3698}
3699
3700impl Drop for OptiForkTicket {
3701    fn drop(&mut self) {
3702        if !self.settled {
3703            let _ = self.drain.synchronize();
3704            OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3705        }
3706    }
3707}
3708
3709impl OptiForkState {
3710    #[allow(clippy::too_many_arguments)]
3711    fn new(
3712        e: &Engine,
3713        cache: &Cache,
3714        mode: OptiForkGateMode,
3715        alternate_snapshot: crate::cache::CacheSnapshot,
3716        h_seed: &CudaSlice<f32>,
3717        fill_prev: &CudaSlice<f32>,
3718        rt: &'static crate::pp::PpNRt,
3719        split: usize,
3720        n_layer: usize,
3721    ) -> Result<Self, Box<dyn std::error::Error>> {
3722        let fence = [0, split, n_layer];
3723        let mut logical_payload_bytes = [0usize; 2];
3724        for stage in 0..2 {
3725            for il in fence[stage]..fence[stage + 1] {
3726                logical_payload_bytes[stage] += alternate_snapshot.conv[il]
3727                    .as_ref()
3728                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3729                logical_payload_bytes[stage] += alternate_snapshot.ssm[il]
3730                    .as_ref()
3731                    .map_or(0, |v| v.len() * std::mem::size_of::<f32>());
3732            }
3733        }
3734        let seeds = [
3735            OptiForkSeedGeneration {
3736                h_seed: e.clone_dtod(h_seed)?,
3737                fill_prev: e.clone_dtod(fill_prev)?,
3738                scratch_len: 0,
3739            },
3740            OptiForkSeedGeneration {
3741                h_seed: e.clone_dtod(h_seed)?,
3742                fill_prev: e.clone_dtod(fill_prev)?,
3743                scratch_len: 0,
3744            },
3745        ];
3746        let (len_ptrs, saved_lens, forced_acc, valid, stage0_stream) = {
3747            let _stage = rt.enter(0);
3748            let e0 = rt.engine(0, e);
3749            (
3750                crate::round_stream::kv_len_ptr_table_range(e0, cache, 0..split, None)?,
3751                e0.htod_i32(&vec![0; split])?,
3752                e0.alloc_u32_zeroed(2)?,
3753                e0.alloc_u32_zeroed(1)?,
3754                e0.stream(),
3755            )
3756        };
3757        logical_payload_bytes[0] += seeds
3758            .iter()
3759            .map(|seed| (seed.h_seed.len() + seed.fill_prev.len()) * std::mem::size_of::<f32>())
3760            .sum::<usize>();
3761        logical_payload_bytes[0] += len_ptrs.len() * std::mem::size_of::<u64>()
3762            + saved_lens.len() * std::mem::size_of::<i32>()
3763            + forced_acc.len() * std::mem::size_of::<u32>()
3764            + valid.len() * std::mem::size_of::<u32>();
3765        Ok(Self {
3766            mode,
3767            controller: (mode == OptiForkGateMode::Controller)
3768                .then(OptiControllerPolicy::configured),
3769            generations: OptiForkGenerationTracker::default(),
3770            active_snapshot_slot: 0,
3771            alternate_snapshot,
3772            seeds,
3773            rt,
3774            fence,
3775            split,
3776            len_ptrs,
3777            saved_lens,
3778            forced_acc,
3779            valid,
3780            stage0_stream,
3781            logical_payload_bytes,
3782        })
3783    }
3784
3785    fn reserve(
3786        &mut self,
3787        current_snapshot: &mut crate::cache::CacheSnapshot,
3788    ) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3789        let generation = self.generations.reserve()?;
3790        if generation.slot != self.active_snapshot_slot {
3791            std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3792            self.active_snapshot_slot = generation.slot;
3793        }
3794        Ok(generation)
3795    }
3796
3797    fn capture_seed(
3798        &mut self,
3799        e: &Engine,
3800        generation: OptiForkGeneration,
3801        h_seed: &CudaSlice<f32>,
3802        fill_prev: &CudaSlice<f32>,
3803        scratch_len: usize,
3804    ) -> Result<(), Box<dyn std::error::Error>> {
3805        let seed = &mut self.seeds[generation.slot];
3806        e.copy_into(&mut seed.h_seed, 0, h_seed, h_seed.len())?;
3807        e.copy_into(&mut seed.fill_prev, 0, fill_prev, fill_prev.len())?;
3808        seed.scratch_len = scratch_len;
3809        Ok(())
3810    }
3811
3812    fn ticket(
3813        &self,
3814        generation: OptiForkGeneration,
3815        boundary: VerifyBoundaryTicket,
3816    ) -> OptiForkTicket {
3817        OptiForkTicket {
3818            generation,
3819            boundary: Some(boundary),
3820            drain: self.stage0_stream.clone(),
3821            settled: false,
3822        }
3823    }
3824
3825    #[allow(clippy::too_many_arguments)]
3826    fn controller_ticket(
3827        &self,
3828        generation: OptiForkGeneration,
3829        boundary: VerifyBoundaryTicket,
3830        ckpt: VerifyCkpt,
3831        verify_tokens: [u32; 2],
3832        draft_prob: f32,
3833        eager_seed: Option<CudaSlice<f32>>,
3834        q_proxy: f32,
3835        scratch_len: usize,
3836    ) -> OptiControllerTicket {
3837        OptiControllerTicket {
3838            generation,
3839            boundary: Some(boundary),
3840            ckpt: Some(ckpt),
3841            verify_tokens,
3842            draft_prob,
3843            eager_seed,
3844            q_proxy,
3845            scratch_len,
3846            issued_at: std::time::Instant::now(),
3847            drain: self.stage0_stream.clone(),
3848            settled: false,
3849        }
3850    }
3851
3852    fn reserve_successor(&mut self) -> Result<OptiForkGeneration, Box<dyn std::error::Error>> {
3853        self.generations.reserve()
3854    }
3855
3856    fn successor_snapshot_mut(&mut self) -> &mut crate::cache::CacheSnapshot {
3857        &mut self.alternate_snapshot
3858    }
3859
3860    fn promote_successor_snapshot(
3861        &mut self,
3862        current_snapshot: &mut crate::cache::CacheSnapshot,
3863        generation: OptiForkGeneration,
3864    ) {
3865        std::mem::swap(current_snapshot, &mut self.alternate_snapshot);
3866        self.active_snapshot_slot = generation.slot;
3867    }
3868
3869    fn queue_actual_reconcile(
3870        &mut self,
3871        e: &Engine,
3872        snapshot: &crate::cache::CacheSnapshot,
3873        acc: &CudaSlice<u32>,
3874        optimistic_pending: u32,
3875        base: usize,
3876    ) -> Result<(), Box<dyn std::error::Error>> {
3877        let saved: Vec<i32> = (0..self.split)
3878            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3879            .collect();
3880        // Serving keeps the caller/accept walk on the head (stage-1) device. Record the accept
3881        // decision point there and append a wait to stage 0 after its optimistic successor/TX;
3882        // the validity/reconcile kernels must never peer-read acc before it is written. The
3883        // increment-1 harness uses primary stage 0, where stream order already provides this.
3884        if self.rt.engine(0, e).ctx().ordinal() != e.ctx().ordinal() {
3885            self.rt.fence_stages_behind(&e.stream())?;
3886        }
3887        let _stage = self.rt.enter(0);
3888        let e0 = self.rt.engine(0, e);
3889        e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3890        e0.spec_fork_valid(acc, optimistic_pending, &mut self.valid)?;
3891        e0.spec_fork_reconcile_kv(
3892            &self.len_ptrs,
3893            &self.saved_lens,
3894            acc,
3895            &self.valid,
3896            base,
3897            self.split,
3898        )
3899    }
3900
3901    fn finish_actual_reconcile(
3902        &mut self,
3903        e: &Engine,
3904        cache: &mut Cache,
3905        snapshot: &crate::cache::CacheSnapshot,
3906        n_acc: usize,
3907        base: usize,
3908        hit: bool,
3909    ) -> Result<(), Box<dyn std::error::Error>> {
3910        if hit {
3911            return Ok(());
3912        }
3913        let len_delta = base + n_acc;
3914        for il in 0..self.split {
3915            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3916                kv.len = saved + len_delta;
3917            }
3918        }
3919        {
3920            let _stage = self.rt.enter(1);
3921            let e1 = self.rt.engine(1, e);
3922            for il in self.split..self.fence[2] {
3923                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3924                    kv.len = saved + len_delta;
3925                    e1.set_i32_one(&mut kv.len_d, kv.len as i32)?;
3926                }
3927            }
3928        }
3929        self.rt.publish_to(0, &e.stream())?;
3930        Ok(())
3931    }
3932
3933    fn cancel_controller_ticket(
3934        &mut self,
3935        e: &Engine,
3936        cache: &mut Cache,
3937        scratch: &mut MtpScratch,
3938        snapshot: &crate::cache::CacheSnapshot,
3939        ticket: &mut OptiControllerTicket,
3940    ) -> Result<(), Box<dyn std::error::Error>> {
3941        {
3942            let _stage = self.rt.enter(0);
3943            let e0 = self.rt.engine(0, e);
3944            for il in 0..self.split {
3945                if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
3946                    kv.len = saved;
3947                    e0.set_i32_one(&mut kv.len_d, saved as i32)?;
3948                }
3949            }
3950        }
3951        scratch.set_len(e, snapshot.pos)?;
3952        ticket.settle();
3953        self.generations.retire(ticket.generation)?;
3954        OPTI_FORK_ABORT_DRAINS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3955        OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
3956        eprintln!(
3957            "[opti-controller] tail-drain generation={} slot={}",
3958            ticket.generation.id, ticket.generation.slot,
3959        );
3960        Ok(())
3961    }
3962
3963    #[allow(clippy::too_many_arguments)]
3964    fn reconcile(
3965        &mut self,
3966        e: &Engine,
3967        cache: &mut Cache,
3968        scratch: &mut MtpScratch,
3969        snapshot: &crate::cache::CacheSnapshot,
3970        h_seed: &mut CudaSlice<f32>,
3971        fill_prev: &mut CudaSlice<f32>,
3972        generation: OptiForkGeneration,
3973        action: OptiForkAction,
3974        optimistic_pending: u32,
3975    ) -> Result<(), Box<dyn std::error::Error>> {
3976        debug_assert!(action != OptiForkAction::Abort);
3977        let miss_started = std::time::Instant::now();
3978        let keep = action == OptiForkAction::Hit;
3979        let saved: Vec<i32> = (0..self.split)
3980            .map(|il| snapshot.kv_len[il].map(|v| v as i32).unwrap_or(0))
3981            .collect();
3982        let seed = &self.seeds[generation.slot];
3983        {
3984            let _stage = self.rt.enter(0);
3985            let e0 = self.rt.engine(0, e);
3986            e0.htod_i32_into(&mut self.saved_lens, &saved)?;
3987            let forced = if keep {
3988                [1u32, optimistic_pending]
3989            } else {
3990                [0u32, optimistic_pending]
3991            };
3992            e0.htod_u32_into(&mut self.forced_acc, &forced)?;
3993            e0.spec_fork_valid(&self.forced_acc, optimistic_pending, &mut self.valid)?;
3994            e0.spec_fork_reconcile_kv(
3995                &self.len_ptrs,
3996                &self.saved_lens,
3997                &self.forced_acc,
3998                &self.valid,
3999                0,
4000                self.split,
4001            )?;
4002            for il in 0..self.split {
4003                if let Some(recur) = cache.recur[il].as_mut() {
4004                    let conv = snapshot.conv[il]
4005                        .as_ref()
4006                        .ok_or("optipipe stage0 snapshot missing conv state")?;
4007                    let ssm = snapshot.ssm[il]
4008                        .as_ref()
4009                        .ok_or("optipipe stage0 snapshot missing ssm state")?;
4010                    e0.spec_fork_restore_f32(conv, &mut recur.conv_state, &self.valid)?;
4011                    e0.spec_fork_restore_f32(ssm, &mut recur.ssm_state, &self.valid)?;
4012                }
4013            }
4014            e0.spec_fork_restore_f32(&seed.h_seed, h_seed, &self.valid)?;
4015            e0.spec_fork_restore_f32(&seed.fill_prev, fill_prev, &self.valid)?;
4016        }
4017
4018        if keep {
4019            OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4020            return Ok(());
4021        }
4022
4023        for il in 0..self.split {
4024            if let (Some(kv), Some(saved)) = (cache.kv[il].as_mut(), snapshot.kv_len[il]) {
4025                kv.len = saved;
4026            }
4027        }
4028        scratch.set_len(e, seed.scratch_len)?;
4029        // Targeted E_restart: publish only stage 0's reconcile to the caller, then bound the
4030        // forced diagnostic so the retained number is the actual miss cost, not enqueue time.
4031        let caller = e.stream();
4032        self.rt.publish_to(0, &caller)?;
4033        caller.synchronize()?;
4034        let miss_ms = miss_started.elapsed().as_secs_f64() * 1e3;
4035        eprintln!(
4036            "[opti-fork-reconcile] generation={} slot={} miss_ms={miss_ms:.3}",
4037            generation.id, generation.slot,
4038        );
4039        OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4040        Ok(())
4041    }
4042
4043    fn retire(&mut self, generation: OptiForkGeneration) -> Result<(), Box<dyn std::error::Error>> {
4044        self.generations.retire(generation)
4045    }
4046}
4047
4048fn rewind_tp_kv_verified_prefix(
4049    tp_kv: &mut [Option<crate::tp::ResidentTpKvCache>],
4050    saved_lens: &[Option<usize>],
4051    accepted: usize,
4052) -> Result<(), Box<dyn std::error::Error>> {
4053    if tp_kv.len() != saved_lens.len() {
4054        return Err("spec TP KV snapshot shape mismatch".into());
4055    }
4056    for (layer, (cache, saved)) in tp_kv.iter_mut().zip(saved_lens).enumerate() {
4057        match (cache.as_mut(), *saved) {
4058            (Some(cache), Some(saved)) => {
4059                let target = saved
4060                    .checked_add(accepted)
4061                    .ok_or("spec TP KV committed length overflow")?;
4062                cache.rewind_to(target)?;
4063            }
4064            (None, None) => {}
4065            _ => {
4066                return Err(
4067                    format!("spec TP KV layer {layer} changed shape since its snapshot").into(),
4068                );
4069            }
4070        }
4071    }
4072    Ok(())
4073}
4074
4075/// MEMRA_SPEC_ROUND_PROF counters: whole-round wall, so the round can be weighed against the
4076/// draft-step ([spec-anatomy]) and verify-walk ([tcol-prof]) splits we already print.
4077static ROUND_PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4078static ROUND_MS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4079static ROUND_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4080
4081impl HybridModel {
4082    fn mtp_head_count(&self) -> usize {
4083        usize::from(self.mtp.is_some()) + self.mtp_extra.len()
4084    }
4085
4086    fn mtp_head_at(&self, index: usize) -> &MtpHead {
4087        if index == 0 {
4088            self.mtp.as_ref().expect("MTP head 0 is unavailable")
4089        } else {
4090            &self.mtp_extra[index - 1]
4091        }
4092    }
4093
4094    fn new_mtp_scratch(
4095        &self,
4096        e: &Engine,
4097        cap: usize,
4098    ) -> Result<MtpScratch, Box<dyn std::error::Error>> {
4099        let mut scratch = MtpScratch::new(
4100            e,
4101            &self.cfg,
4102            &self.plan,
4103            cap,
4104            self.mtp.as_ref().and_then(|head| head.geom.as_ref()),
4105        )?;
4106        for head in &self.mtp_extra {
4107            scratch.push_plane(e, &self.cfg, &self.plan, head.geom.as_ref())?;
4108        }
4109        Ok(scratch)
4110    }
4111
4112    fn opti_graph_draft_step(
4113        &self,
4114        e: &Engine,
4115        mtp: &MtpHead,
4116        dctx: &mut DraftGraphCtx,
4117        scratch: &mut MtpScratch,
4118        d_vocab: usize,
4119    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4120        // dcw door: one replay appends one device-counter row; pre-arm ring headroom
4121        // host-side before launching (no-op on flat planes).
4122        if step35_draft_dcw_on() {
4123            scratch.ensure_dcw_headroom(e, 2)?;
4124        }
4125        dctx.graph
4126            .as_ref()
4127            .ok_or("optipipe controller requires the greedy draft graph")?
4128            .launch()?;
4129        scratch.kv.len += 1;
4130        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
4131        if (idx as usize) >= d_vocab {
4132            return Err(
4133                format!("optipipe draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}").into(),
4134            );
4135        }
4136        let probability = e.dtoh(&dctx.g_p)?[0];
4137        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4138            return Err(format!("optipipe draft probability is invalid: {probability}").into());
4139        }
4140        let token = match &mtp.d2t {
4141            Some(map) => map[idx as usize],
4142            None => idx,
4143        };
4144        if token != idx {
4145            e.set_u32_one(&mut dctx.g_tok, token)?;
4146        }
4147        Ok((token, probability))
4148    }
4149
4150    #[allow(clippy::too_many_arguments)]
4151    fn opti_controller_draft_step(
4152        &self,
4153        e: &Engine,
4154        mtp: &MtpHead,
4155        dctx: &mut DraftGraphCtx,
4156        scratch: &mut MtpScratch,
4157        d_vocab: usize,
4158        eager_state: &mut Option<(u32, CudaSlice<f32>)>,
4159        eager_pos: usize,
4160        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4161    ) -> Result<(u32, f32), Box<dyn std::error::Error>> {
4162        if dctx.graph.is_some() {
4163            return self.opti_graph_draft_step(e, mtp, dctx, scratch, d_vocab);
4164        }
4165        let (input_token, input_seed) = eager_state
4166            .take()
4167            .ok_or("optipipe eager continuation seed is unavailable")?;
4168        let (logits, next_seed) = self.mtp_head_forward_dev(
4169            e,
4170            mtp,
4171            input_token,
4172            &input_seed,
4173            scratch,
4174            eager_pos,
4175            embd_dev,
4176            None,
4177        )?;
4178        let token_d = e.argmax_token_device(&logits, d_vocab)?;
4179        let idx = e.dtoh_u32_one(&token_d)?;
4180        if (idx as usize) >= d_vocab {
4181            return Err(format!(
4182                "optipipe eager draft argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab}"
4183            )
4184            .into());
4185        }
4186        let probability_d = e.prob_of_token_device(&logits, &token_d, d_vocab)?;
4187        let probability = e.dtoh(&probability_d)?[0];
4188        if !probability.is_finite() || !(0.0..=1.0).contains(&probability) {
4189            return Err(
4190                format!("optipipe eager draft probability is invalid: {probability}").into(),
4191            );
4192        }
4193        let token = match &mtp.d2t {
4194            Some(map) => map[idx as usize],
4195            None => idx,
4196        };
4197        *eager_state = Some((token, next_seed));
4198        Ok((token, probability))
4199    }
4200
4201    /// NextN head forward for ONE draft token (§A ops 1-13, T=1).
4202    /// Inputs: `e_tok` = the token to predict FROM (last committed / previous draft); `h_seed` =
4203    /// the trunk's pre-output_norm hidden of that token (§A op 2 input). `mtp_pos` = absolute
4204    /// position of the token being predicted from. Returns (draft_logits[n_vocab] host, h_nextn dev).
4205    /// `h_nextn` (§A op 10) becomes `h_seed` for the next autoregressive draft step.
4206    /// Device-resident: returns draft logits ON DEVICE (no [n_vocab] dtoh). The greedy draft
4207    /// loop only needs argmax — paired with `argmax_token_device` this cuts the ~600KB logits
4208    /// transfer + host argmax per draft token from the K-token draft chain.
4209    #[allow(clippy::too_many_arguments)]
4210    fn mtp_head_forward_dev(
4211        &self,
4212        e: &Engine,
4213        mtp: &MtpHead,
4214        e_tok: u32,
4215        h_seed: &CudaSlice<f32>,
4216        scratch: &mut MtpScratch,
4217        mtp_pos: usize,
4218        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4219        mask: Option<(&CudaSlice<u32>, usize)>,
4220    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4221        self.mtp_head_forward_dev_at(e, mtp, e_tok, h_seed, scratch, 0, mtp_pos, embd_dev, mask)
4222    }
4223
4224    #[allow(clippy::too_many_arguments)]
4225    fn mtp_head_forward_dev_at(
4226        &self,
4227        e: &Engine,
4228        mtp: &MtpHead,
4229        e_tok: u32,
4230        h_seed: &CudaSlice<f32>,
4231        scratch: &mut MtpScratch,
4232        scratch_index: usize,
4233        mtp_pos: usize,
4234        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4235        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed set, words).
4236        // Applied to the head logits BEFORE they are returned, so every consumer (argmax,
4237        // gumbel draw, p-min prob) sees the grammar-legal row. None = unmasked (pre-lane).
4238        mask: Option<(&CudaSlice<u32>, usize)>,
4239    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4240        // MEMRA_SPEC_ANATOMY=1 — eager-step phase timers (diagnostic only). Phase boundaries
4241        // sync the stream, so absolute time inflates; the BREAKDOWN is the signal. Cumulative
4242        // summary on stderr every 128 steps: glue (embed..attn_norm), attn, ffn, head.
4243        use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
4244        static ANAT_NS: [AtomicU64; 5] = [
4245            AtomicU64::new(0),
4246            AtomicU64::new(0),
4247            AtomicU64::new(0),
4248            AtomicU64::new(0),
4249            AtomicU64::new(0),
4250        ];
4251        static ANAT_STEPS: AtomicU64 = AtomicU64::new(0);
4252        let anat = {
4253            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4254            *ON.get_or_init(|| std::env::var("MEMRA_SPEC_ANATOMY").as_deref() == Ok("1"))
4255        };
4256        if anat {
4257            e.stream().synchronize()?; // drain prior queue so phase 0 starts clean
4258        }
4259        let t_all = std::time::Instant::now();
4260        let mut t_ph = std::time::Instant::now();
4261        let mut anat_mark = |i: usize,
4262                             e: &Engine,
4263                             t: &mut std::time::Instant|
4264         -> Result<(), Box<dyn std::error::Error>> {
4265            if anat {
4266                e.stream().synchronize()?;
4267                ANAT_NS[i].fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
4268                *t = std::time::Instant::now();
4269            }
4270            Ok(())
4271        };
4272        let cfg = &self.cfg;
4273        let n_embd = cfg.n_embd as usize;
4274        // Distilled-student geometry: the block runs at the INNER width `di` (eh_proj out /
4275        // attn / ffn); the n_embd interface (embed, norms in, carrier out, head in) is unchanged.
4276        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
4277        let eps = cfg.rms_eps;
4278        let pos_d = e.htod_i32(&[mtp_pos as i32])?;
4279
4280        // op A: a resident table transfers one 4B token id. The exact host-row capacity path
4281        // expands this one row on CPU and transfers n_embd f32 values instead.
4282        let e_emb = match embd_dev {
4283            Some((g, qt, rb)) => e.embed_gather_device_t(g, &[e_tok], n_embd, qt, rb)?,
4284            None => e.htod(&self.embd.gather(n_embd, &[e_tok]))?,
4285        };
4286
4287        // op 1/2: e_norm = RMSNorm(e, enorm); h_norm = RMSNorm(h_seed, hnorm)
4288        let mut e_norm = e.zeros(n_embd)?;
4289        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
4290        let mut h_norm = e.zeros(n_embd)?;
4291        e.rms_norm(h_seed, mtp.hnorm.float_data(), &mut h_norm, n_embd, 1, eps)?;
4292
4293        // op 3: concat = [e_norm ; h_norm] -> [2*n_embd], e_norm in [0,n_embd), h_norm in [n_embd,2n_embd)
4294        let mut concat = e.zeros(2 * n_embd)?;
4295        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
4296        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
4297
4298        // op 4: inpSA = eh_proj @ concat  (eh_proj [2*n_embd, n_embd]) -> [n_embd]
4299        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
4300
4301        // op 5: a_norm = RMSNorm(inpSA, attn_norm)
4302        let mut a_norm = e.zeros(di)?;
4303        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
4304        anat_mark(0, e, &mut t_ph)?;
4305
4306        // op 6: attention on the scratch KV. SAME dc launcher as the graph path (bucket_max =
4307        // scratch.cap, length from the device len_d) so eager drafts match graph drafts
4308        // bit-for-bit at any t_kv (the parity gate). Host len mirrored here (the dc append
4309        // advances only the device counter).
4310        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
4311            // step35 MTP block, dcw door armed: the SAME windowed device-counter launcher as
4312            // the captured chain (draft parity by construction). Per-step ring headroom runs
4313            // HERE (eager is host-len work, a rebase is legal); host len mirrored like the
4314            // plain dc arm below.
4315            (Mixer::Full(fa), Some(g))
4316                if self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) =>
4317            {
4318                {
4319                    let (kv, _) = scratch.plane_mut(scratch_index);
4320                    let retain = match kv.ring.as_ref() {
4321                        Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4322                        None => 0,
4323                    };
4324                    e.prepare_kv_append(kv, retain, 1)?;
4325                }
4326                let out =
4327                    self.mtp_step35_attn_dcw(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?;
4328                scratch.plane_mut(scratch_index).0.len += 1;
4329                out
4330            }
4331            // step35 MTP block, door off (MEMRA_STEP35_DRAFT_DCW=0 rollback) or class-
4332            // ineligible: PER-LAYER geometry + a separate head-wise gate + an SWA window,
4333            // none of which the plain dc launcher can express (see `mtp_step35_attn`).
4334            // Host-len arm. Advances BOTH the
4335            // host len and the device counter itself (unlike the dc arm, whose host-side
4336            // mirror the caller does).
4337            (Mixer::Full(fa), Some(g)) => {
4338                self.mtp_step35_attn(e, fa, g, &a_norm, &pos_d, scratch, scratch_index)?
4339            }
4340            (Mixer::Full(fa), None) => {
4341                let out = self.mtp_full_attn_dc(
4342                    e,
4343                    fa,
4344                    &a_norm,
4345                    &pos_d,
4346                    scratch,
4347                    scratch_index,
4348                    mtp.geom.as_ref(),
4349                )?;
4350                scratch.plane_mut(scratch_index).0.len += 1;
4351                out
4352            }
4353            (Mixer::Linear(_), _) => {
4354                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
4355            }
4356            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
4357        };
4358        anat_mark(1, e, &mut t_ph)?;
4359
4360        // op 7: x1 = inpSA + attn_out
4361        let mut x1 = e.zeros(di)?;
4362        e.add(&inp_sa, &attn_out, &mut x1, di)?;
4363
4364        // op 8: z = RMSNorm(x1, post_attn_norm)  (pre-FFN norm)
4365        let mut z = e.zeros(di)?;
4366        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
4367
4368        // op 9: FFN (Dense or MoE) — same as the trunk decode FFN
4369        let ffn_out = match &mtp.ffn {
4370            crate::hybrid::Ffn::Dense {
4371                ffn_gate,
4372                ffn_up,
4373                ffn_down,
4374            } => {
4375                let n_ff = ffn_gate.out_features();
4376                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
4377                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
4378                    (
4379                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
4380                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
4381                    )
4382                } else {
4383                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
4384                };
4385                let mut act = e.zeros(n_ff)?;
4386                // step35: a DENSE FFN reads the per-layer SHEXP clamp (upstream's one `build_ffn`
4387                // serves the dense MLP and the shared expert off `swiglu_clamp_shexp` —
4388                // llama-graph.cpp:1751), resolved for the MTP block's OWN index. Every other arch
4389                // passes None, which is `ffn_act`'s dispatch verbatim.
4390                Self::ffn_act_lim(
4391                    e,
4392                    &self.cfg,
4393                    &gate,
4394                    &up,
4395                    1.0,
4396                    1.0,
4397                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
4398                    &mut act,
4399                    n_ff,
4400                )?;
4401                e.matmul(ffn_down, &act, 1)?
4402            }
4403            // MTP head is a distinct block — key its experts under a separate layer index (u16::MAX)
4404            // so they never alias trunk layer 0's cache keys.
4405            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, 1, u16::MAX)?,
4406        };
4407        anat_mark(2, e, &mut t_ph)?;
4408
4409        // op 10: h_nextn = x1 + ffn_out (at di)
4410        let mut h_inner = e.zeros(di)?;
4411        e.add(&x1, &ffn_out, &mut h_inner, di)?;
4412
4413        // op 10.5 (student): up-project the inner hidden back to n_embd — training semantics:
4414        // the chain carrier AND the head input are out_up(h_inner) (pre-final-norm).
4415        let h_nextn = match mtp.geom.as_ref() {
4416            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
4417            None => h_inner,
4418        };
4419
4420        // op 11: final = RMSNorm(h_nextn, shared_head_norm OR output_norm)
4421        let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
4422        let mut final_h = e.zeros(n_embd)?;
4423        e.rms_norm(
4424            &h_nextn,
4425            final_norm.float_data(),
4426            &mut final_h,
4427            n_embd,
4428            1,
4429            eps,
4430        )?;
4431
4432        // op 12: draft_logits = (shared_head_head OR output) @ final — stays ON DEVICE.
4433        let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
4434        let mut logits = e.matmul(head, &final_h, 1)?;
4435        // op 12b (lane/draft-mask): grammar mask over the DRAFT vocab, applied here so the
4436        // caller's argmax / gumbel draw / p-min prob all read the grammar-legal row.
4437        if let Some((mask_d, mw)) = mask {
4438            let d_vocab = head.out_features();
4439            e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
4440        }
4441        anat_mark(3, e, &mut t_ph)?;
4442        if anat {
4443            ANAT_NS[4].fetch_add(t_all.elapsed().as_nanos() as u64, Relaxed);
4444            let n = ANAT_STEPS.fetch_add(1, Relaxed) + 1;
4445            if n % 128 == 0 {
4446                let us = |i: usize| ANAT_NS[i].load(Relaxed) / n / 1000;
4447                eprintln!(
4448                    "[spec-anatomy] steps={n} avg us/step: glue={} attn={} ffn={} head={} total={}",
4449                    us(0),
4450                    us(1),
4451                    us(2),
4452                    us(3),
4453                    us(4)
4454                );
4455            }
4456        }
4457        // Chain recurrence hand-over: pre-norm h_nextn (default) or post-norm final_h
4458        // (MEMRA_SPEC_HPOST — llama.cpp #24025's t_h_nextn is taken AFTER the head norm).
4459        Ok((logits, if spec_hpost() { final_h } else { h_nextn }))
4460    }
4461
4462    #[allow(clippy::too_many_arguments)]
4463    fn mtp_chain_forward_dev(
4464        &self,
4465        e: &Engine,
4466        tokens: &[u32],
4467        seeds: &[CudaSlice<f32>],
4468        scratch: &mut MtpScratch,
4469        committed_scratch_len: usize,
4470        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
4471        mask: Option<(&CudaSlice<u32>, usize)>,
4472    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
4473        if tokens.is_empty() || tokens.len() != seeds.len() {
4474            return Err("multi-head MTP prefix tokens/seeds are malformed".into());
4475        }
4476        let index = mtp_chain_head_index(tokens.len() - 1, self.mtp_head_count());
4477        let head = self.mtp_head_at(index);
4478        scratch.set_plane_len(e, index, committed_scratch_len)?;
4479
4480        let mut last = None;
4481        for row in 0..tokens.len() {
4482            let is_last = row + 1 == tokens.len();
4483            last = Some(self.mtp_head_forward_dev_at(
4484                e,
4485                head,
4486                tokens[row],
4487                &seeds[row],
4488                scratch,
4489                index,
4490                committed_scratch_len + row + 1,
4491                embd_dev,
4492                if is_last { mask } else { None },
4493            )?);
4494        }
4495        Ok(last.expect("non-empty MTP prefix produced no row"))
4496    }
4497
4498    /// step35 MTP-block attention, T=1, on the scratch KV — the EAGER-ONLY twin of
4499    /// `mtp_full_attn_dc`. Three things force a separate arm rather than a geometry parameter on
4500    /// the dc path, and all three are properties of this arch's MTP block:
4501    ///
4502    /// 1. **The SWA window.** Block 45 is an SWA-type block (`sliding_window_pattern[45]=true`,
4503    ///    window 512). Windowed decode in memra is a token-aligned VIEW OFFSET into the quantized
4504    ///    cache (the gemma4 R6 / `step35_decode_attn` pattern: keys carry absolute rope and the
4505    ///    mask is purely positional, so one query at `len-1` attending the last `win` rows IS the
4506    ///    windowed result). `fa_decode_dc` takes the key count from a DEVICE counter and always
4507    ///    starts at row 0 — it cannot express a nonzero offset. The windowed dc arm is
4508    ///    `mtp_step35_attn_dcw` (`fa_decode_dcw`, doored via MEMRA_STEP35_DRAFT_DCW —
4509    ///    default ON since lane/step37-draft-graph-serving-20260830); this host-len arm is
4510    ///    the =0 rollback and the class-ineligibility fallback.
4511    /// 2. **Per-layer head count.** 96 q heads over 8 KV (GQA 12) at this block, vs the trunk's 64
4512    ///    on its full-attn layers. The trunk cfg's `n_head` scalar is the MAX over layers, and the
4513    ///    trunk ARTIFACT's per-layer arrays stop at index 44 — so the count must come from the
4514    ///    resolved `Step35MtpGeom`, never from `cfg`.
4515    /// 3. **The separate head-wise gate.** `blk.45.attn_gate.weight [n_embd, 96]` produces one
4516    ///    sigmoid scalar per head (broadcast over head_dim) — `attn_head_gate`, not the qwen35
4517    ///    fused-into-wq `q_gate_split` form the dc arm handles.
4518    ///
4519    /// DOOR STATE: with MEMRA_STEP35_DRAFT_DCW=0 (or a sub-eligible kernel class),
4520    /// `mtp_head_forward_cap` refuses step35 heads explicitly (rather than silently capturing
4521    /// a window-less, wrong-past-`win` graph) and this eager chain IS the served path. With
4522    /// the door armed (the default), BOTH draft modes run the `mtp_step35_attn_dcw` twin
4523    /// instead of this arm.
4524    ///
4525    /// Unlike the dc arm this advances BOTH the host `kv.len` and the device counter, so the
4526    /// caller must not mirror.
4527    fn mtp_step35_attn(
4528        &self,
4529        e: &Engine,
4530        fa: &FullAttnLayer,
4531        g: &crate::hybrid::Step35MtpGeom,
4532        h: &CudaSlice<f32>,
4533        pos_d: &CudaSlice<i32>,
4534        scratch: &mut MtpScratch,
4535        scratch_index: usize,
4536    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4537        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4538        // MTP-GEOM RECEIPT, once per process, on the SERVED draft path. Slot-0 acceptance is
4539        // 0.725 here against 0.994 for vLLM MTP3 on the same checkpoint family and card class, and
4540        // the first three explanations for that gap were all wrong: head assignment (step-modulo
4541        // is index 0 at K=1, correct), MEMRA_SPEC_HPOST (identical 84/116 both arms), and this
4542        // block's geometry. Geometry was the one that could have failed SILENTLY — a wrong window
4543        // makes the draft attend the whole context instead of Step-3.7's 512, stays fluent, and
4544        // shows up only as acceptance — so it gets a standing receipt rather than another reading
4545        // of the source. Prints the resolved Step35MtpGeom the served path actually runs on;
4546        // `full_attention_geometry_at`'s missing-row fallback (window: None) does NOT reach here.
4547        {
4548            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4549            ONCE.get_or_init(|| {
4550                eprintln!(
4551                    "[mtp-geom] arm=eager block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4552                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4553                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4554                );
4555            });
4556        }
4557        let eps = self.cfg.rms_eps;
4558        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4559        let n_embd = self.cfg.n_embd as usize;
4560        let gw = fa
4561            .attn_gate
4562            .as_ref()
4563            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4564
4565        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4566            && e.uses_q8_1_fast(&fa.wk)
4567            && e.uses_q8_1_fast(&fa.wv)
4568            && e.uses_q8_1_fast(gw)
4569        {
4570            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4571            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4572                Some(t3) => t3,
4573                None => (
4574                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4575                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4576                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4577                ),
4578            };
4579            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4580        } else {
4581            (
4582                e.matmul(&fa.wq, h, 1)?,
4583                e.matmul(&fa.wk, h, 1)?,
4584                e.matmul(&fa.wv, h, 1)?,
4585                e.matmul(gw, h, 1)?,
4586            )
4587        };
4588
4589        let mut q = e.uninit(nh * hd)?;
4590        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4591        let mut k = e.uninit(nkv * hd)?;
4592        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4593        // `rope_freqs.weight` (llama3 factors) applies to the FULL-attn layers ONLY; SWA passes
4594        // null (llama-hparams / step35.cpp). Block 45 is SWA, so `ff` is None there — but read
4595        // the resolved flag, not the constant, so an all-full sibling stays correct.
4596        let ff = if g.swa {
4597            None
4598        } else {
4599            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4600        };
4601        #[cfg(debug_assertions)]
4602        if let Some(ff) = ff {
4603            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_step35_attn.rope_freqs");
4604        }
4605        e.rope_neox2(
4606            &mut q,
4607            &mut k,
4608            pos_d,
4609            hd,
4610            g.n_rot,
4611            nh,
4612            nkv,
4613            1,
4614            g.rope_base,
4615            1.0,
4616            ff,
4617        )?;
4618
4619        // Append at the HOST slot, then re-stamp the device counter: the eager chain has the
4620        // length on the host anyway, and the windowed view below needs it there to compute the
4621        // offset. The device counter is kept in lockstep so `mtp_kv_fill`'s `set_i32_one` and any
4622        // dc-family consumer of this scratch still agree.
4623        let (kv, scratch_cap) = scratch.plane_mut(scratch_index);
4624        assert!(
4625            kv.len < scratch_cap,
4626            "step35 MTP scratch overflow ({} >= {})",
4627            kv.len,
4628            scratch_cap
4629        );
4630        let next_len = kv.len + 1;
4631        let (off, t_kv) = if g.swa && next_len > g.window {
4632            (next_len - g.window, g.window)
4633        } else {
4634            (0, next_len)
4635        };
4636        // `off`/`t_kv` stay the ATTENTION view; the retain is a separate, lower bound so the
4637        // rewind that follows this append is still resident. THIS is the only site that rebases
4638        // this plane (MEMRA_KV_REBASE_TRACE, one run: 1 rebase, all from here), so it is the site
4639        // that decides `base` for everyone.
4640        let retain_from = match kv.ring.as_ref() {
4641            Some(ring) => memra_kv::swa_retain_from(kv.len, ring.window(), ring.base()),
4642            None => off & !31usize,
4643        };
4644        let write_row = e.prepare_kv_append(kv, retain_from, 1)?;
4645        e.append_kv_quantized(
4646            &k,
4647            &v0,
4648            &mut kv.k,
4649            &mut kv.v,
4650            write_row,
4651            kv.kv_dim_k,
4652            kv.kv_dim_v,
4653            kv.k_tok_bytes,
4654            kv.v_tok_bytes,
4655            false,
4656        )?;
4657        kv.len = next_len;
4658        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
4659        // SWA view offset (see note 1). The draft chain is short (k+2 rows), but the scratch is
4660        // PERSISTENT across rounds — `mtp_kv_fill` leaves one row per committed token behind, so
4661        // `kv.len` tracks absolute position and crosses 512 in any real generation. The window is
4662        // therefore live, not theoretical.
4663        let physical = kv.physical_rows(off, off + t_kv)?;
4664        let k_view = e.view_u8_range(
4665            &kv.k,
4666            physical.start * kv.k_tok_bytes,
4667            physical.end * kv.k_tok_bytes,
4668        );
4669        let v_view = e.view_u8_range(
4670            &kv.v,
4671            physical.start * kv.v_tok_bytes,
4672            physical.end * kv.v_tok_bytes,
4673        );
4674        let mut attn = e.uninit(nh * hd)?;
4675        e.fa_decode_kvmod(
4676            &q,
4677            &k_view,
4678            &v_view,
4679            &mut attn,
4680            hd,
4681            nh,
4682            nkv,
4683            t_kv,
4684            scale,
4685            kv.k_tok_bytes,
4686            kv.v_tok_bytes,
4687            false,
4688        )?;
4689
4690        let mut ag = e.uninit(nh * hd)?;
4691        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4692        Ok(e.matmul(&fa.wo, &ag, 1)?)
4693    }
4694
4695    /// The dcw draft arm's kernel-class precondition, mirrored from `fa_decode_dcw`'s own
4696    /// refusal plus the v3 walk's format contract (`fa_v3_active`), so the DEV dispatch can
4697    /// never pick an arm the launcher would refuse mid-chain (the eager chain has no graceful
4698    /// fallback point) and the CAP site refuses with the named reason instead.
4699    ///
4700    /// `cap` = the SESSION's scratch-plane row capacity: the launcher's vec gate reads
4701    /// `bucket_max = min(window, cap)`, so a SMALL session (tiny prompt + tiny max_tokens,
4702    /// e.g. a max_tokens=8 probe: cap ~62 < the 96 vec floor) is OUTSIDE the dcw domain even
4703    /// though the WINDOW clears the floor. Mirroring the window alone shipped exactly that
4704    /// hole when the door default flipped ON (2026-08-30, vision-cell receipt: sampled
4705    /// capture WARN + `[engine-error] fa_decode_dcw supports the default v3-vec class only`
4706    /// hard-failing the burst — the eager dcw arm has no graceful fallback point). Sub-floor
4707    /// sessions now take the host-len kvmod arm, byte-for-byte the door-off serving.
4708    fn step35_dcw_eligible(&self, g: &crate::hybrid::Step35MtpGeom, cap: usize) -> bool {
4709        let hd = self.cfg.head_dim_k as usize;
4710        step35_draft_dcw_on()
4711            && g.swa
4712            && g.window.min(cap) >= crate::fa_vec_min_tkv()
4713            && std::env::var("MEMRA_NO_FA_VEC").is_err()
4714            && crate::fa_v3_active(hd)
4715            && hd <= 256
4716            && hd % 32 == 0
4717    }
4718
4719    /// step35 MTP-block attention, T=1, on the scratch KV: the WINDOWED DEVICE-COUNTER twin
4720    /// of `mtp_step35_attn`, serving BOTH draft paths when `step35_draft_dcw_on`. Write slot,
4721    /// key bound and SWA view offset all derive from device state (`len_d`, `base_d` written
4722    /// only at host-side rebases, and the block's `window`), so ONE captured graph serves the
4723    /// whole chain and replays see KV growth through the counter: the `mtp_full_attn_dc`
4724    /// contract plus the view offset the plain `_dc` kernel could not express (the old
4725    /// capture-refusal root cause). The three step35 properties stay per-geom exactly as in
4726    /// the eager twin: nh/nkv from `Step35MtpGeom`, the separate head-wise gate
4727    /// (`attn_head_gate`), per-layer rope width/base with SWA passing null freqs.
4728    ///
4729    /// bucket_max = min(cap, window): the windowed view never exceeds `window` rows, so the
4730    /// capture-time grid stays valid for every replayed len, and the kernel derives ns_eff
4731    /// from the LIVE T_kv at the fixed split_keys (one-partition law). Both arms call THIS
4732    /// launcher at THIS bucket, so eager and captured drafts are bit-identical by
4733    /// construction; vs the retired-by-flag `mtp_step35_attn` the only numeric-class deltas
4734    /// are the sub-vec-floor region (t_kv < 96: kvmod ran scalar, dcw stays vec) and any
4735    /// live-len split-ladder rung below the bucket's, both draft-side only (the verify
4736    /// arbitrates emitted bytes; acceptance is gated by the battery).
4737    ///
4738    /// Host len is NOT advanced here (graph contract); callers mirror. The EAGER caller runs
4739    /// `prepare_kv_append` per step (ring headroom, rebase legal there); the CAPTURED path
4740    /// pre-arms headroom at capture time and round start (`MtpScratch::ensure_dcw_headroom`)
4741    /// because a rebase is host work no captured chain may contain.
4742    fn mtp_step35_attn_dcw(
4743        &self,
4744        e: &Engine,
4745        fa: &FullAttnLayer,
4746        g: &crate::hybrid::Step35MtpGeom,
4747        h: &CudaSlice<f32>,
4748        pos_d: &CudaSlice<i32>,
4749        scratch: &mut MtpScratch,
4750        scratch_index: usize,
4751    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4752        let (nh, nkv, hd) = (g.n_head, g.n_head_kv, self.cfg.head_dim_k as usize);
4753        // MTP-GEOM RECEIPT (dcw twin of the `mtp_step35_attn` receipt): once per process,
4754        // naming the arm, so a serving log proves WHICH draft attention program ran (the
4755        // engagement receipt for the flag door, both directions).
4756        {
4757            static ONCE: std::sync::OnceLock<()> = std::sync::OnceLock::new();
4758            ONCE.get_or_init(|| {
4759                eprintln!(
4760                    "[mtp-geom] arm=dcw block={} swa={} window={} n_head={nh} n_head_kv={nkv} \
4761                     head_dim_k={hd} n_rot={} rope_base={} clamp_shexp={:?}",
4762                    g.il, g.swa, g.window, g.n_rot, g.rope_base, g.clamp_shexp,
4763                );
4764            });
4765        }
4766        let eps = self.cfg.rms_eps;
4767        let scale = 1.0 / (hd as f32).sqrt(); // step35.cpp:255 kq_scale
4768        let n_embd = self.cfg.n_embd as usize;
4769        let gw = fa
4770            .attn_gate
4771            .as_ref()
4772            .ok_or("step35 MTP block is missing attn_gate.weight (head-wise attention gate)")?;
4773
4774        let (q0, k0, v0, gt) = if e.uses_q8_1_fast(&fa.wq)
4775            && e.uses_q8_1_fast(&fa.wk)
4776            && e.uses_q8_1_fast(&fa.wv)
4777            && e.uses_q8_1_fast(gw)
4778        {
4779            let (hq, hdq) = e.quantize_q8_1(h, 1, n_embd)?;
4780            let (a, b, c) = match e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hdq)? {
4781                Some(t3) => t3,
4782                None => (
4783                    e.matmul_pre(&fa.wq, &hq, &hdq, h, 1)?,
4784                    e.matmul_pre(&fa.wk, &hq, &hdq, h, 1)?,
4785                    e.matmul_pre(&fa.wv, &hq, &hdq, h, 1)?,
4786                ),
4787            };
4788            (a, b, c, e.matmul_pre(gw, &hq, &hdq, h, 1)?)
4789        } else {
4790            (
4791                e.matmul(&fa.wq, h, 1)?,
4792                e.matmul(&fa.wk, h, 1)?,
4793                e.matmul(&fa.wv, h, 1)?,
4794                e.matmul(gw, h, 1)?,
4795            )
4796        };
4797
4798        let mut q = e.zeros(nh * hd)?;
4799        e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, nh, eps)?;
4800        let mut k = e.zeros(nkv * hd)?;
4801        e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, nkv, eps)?;
4802        // rope_freqs (llama3 factors) apply to the FULL-attn layers ONLY; SWA passes null
4803        // (the eager twin's rule, resolved from the flag, not the constant).
4804        let ff = if g.swa {
4805            None
4806        } else {
4807            self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4808        };
4809        #[cfg(debug_assertions)]
4810        if let Some(ff) = ff {
4811            crate::debug_assert_tensor_stream_device(
4812                ff,
4813                &e.stream(),
4814                "mtp_step35_attn_dcw.rope_freqs",
4815            );
4816        }
4817        e.rope_neox2(
4818            &mut q,
4819            &mut k,
4820            pos_d,
4821            hd,
4822            g.n_rot,
4823            nh,
4824            nkv,
4825            1,
4826            g.rope_base,
4827            1.0,
4828            ff,
4829        )?;
4830
4831        let (kv, cap) = scratch.plane_mut(scratch_index);
4832        // Append at the DEVICE slot's PHYSICAL row (len_d - base_d), then advance the counter
4833        // in-graph. Physical room is the callers' headroom contract (see the fn doc).
4834        e.append_kv_quantized_dcw(
4835            &k,
4836            &v0,
4837            &mut kv.k,
4838            &mut kv.v,
4839            &kv.len_d,
4840            kv.base_d.as_ref(),
4841            kv.kv_dim_k,
4842            kv.kv_dim_v,
4843            kv.k_tok_bytes,
4844            kv.v_tok_bytes,
4845        )?;
4846        e.inc_seqlen(&mut kv.len_d)?;
4847        // Full-buffer views (any in-round physical row stays in range under the headroom
4848        // contract); the kernel bounds and offsets the key range from (len_d, base_d, window).
4849        let k_view = e.view_u8(&kv.k, kv.k.len());
4850        let v_view = e.view_u8(&kv.v, kv.v.len());
4851        let bucket = g.window.min(cap);
4852        let mut attn = e.zeros(nh * hd)?;
4853        e.fa_decode_dcw(
4854            &q,
4855            &k_view,
4856            &v_view,
4857            &mut attn,
4858            hd,
4859            nh,
4860            nkv,
4861            &kv.len_d,
4862            kv.base_d.as_ref(),
4863            if g.swa { g.window } else { 0 },
4864            bucket,
4865            scale,
4866            kv.k_tok_bytes,
4867            kv.v_tok_bytes,
4868            None,
4869        )?;
4870
4871        let mut ag = e.zeros(nh * hd)?;
4872        e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, 1)?;
4873        Ok(e.matmul(&fa.wo, &ag, 1)?)
4874    }
4875
4876    /// MTP-block full attention, T=1, on the scratch KV (BOTH draft paths — eager and graph):
4877    /// the scratch write slot and the attention bound come from `scratch.kv.len_d` (device i32[1])
4878    /// so the launch args are FIXED across draft steps — ONE captured graph serves the whole
4879    /// chain, and replays keep seeing KV growth through the device counter (no recapture).
4880    /// Geometry contract: n_splits is sized from `scratch.cap` (the persistent capacity); splits
4881    /// whose key range lies beyond the device t_kv exit empty and the shared combine skips them
4882    /// (fa_decode_dc bit-correct-for-any-t_kv<=bucket_max contract). The eager path uses the SAME
4883    /// launcher with the SAME bucket_max -> identical dispatch -> bit-identical draft tokens (the
4884    /// graph-vs-eager parity gate). Host len is NOT advanced here (graph contract); callers mirror.
4885    fn mtp_full_attn_dc(
4886        &self,
4887        e: &Engine,
4888        fa: &FullAttnLayer,
4889        h: &CudaSlice<f32>,
4890        pos_d: &CudaSlice<i32>,
4891        scratch: &mut MtpScratch,
4892        scratch_index: usize,
4893        geom: Option<&crate::hybrid::DraftGeom>,
4894    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4895        let cfg = &self.cfg;
4896        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
4897        let geometry = cfg.full_attention_geometry_at(mtp_il);
4898        let n_head = geom.map(|g| g.n_head).unwrap_or(geometry.n_head as usize);
4899        let n_head_kv = geom
4900            .map(|g| g.n_head_kv)
4901            .unwrap_or(geometry.n_head_kv as usize);
4902        let head_dim = geometry.head_dim_k as usize;
4903        let eps = cfg.rms_eps;
4904        let scale = geometry.attention_scale();
4905        let n_embd = geom.map(|g| g.d_inner).unwrap_or(cfg.n_embd as usize);
4906        let bucket_max = scratch.plane(scratch_index).1;
4907
4908        let (qf, mut k, v) =
4909            if e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv) {
4910                let (hq, hd) = e.quantize_q8_1(h, 1, n_embd)?;
4911                (
4912                    e.matmul_pre(&fa.wq, &hq, &hd, h, 1)?,
4913                    e.matmul_pre(&fa.wk, &hq, &hd, h, 1)?,
4914                    e.matmul_pre(&fa.wv, &hq, &hd, h, 1)?,
4915                )
4916            } else {
4917                (
4918                    e.matmul(&fa.wq, h, 1)?,
4919                    e.matmul(&fa.wk, h, 1)?,
4920                    e.matmul(&fa.wv, h, 1)?,
4921                )
4922            };
4923        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
4924        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
4925        let (mut q, gate) = if gated {
4926            let mut q = e.zeros(n_head * head_dim)?;
4927            let mut gate = e.zeros(n_head * head_dim)?;
4928            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, 1)?;
4929            (q, Some(gate))
4930        } else {
4931            (qf, None)
4932        };
4933
4934        let mut qn = e.zeros(n_head * head_dim)?;
4935        e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, n_head, eps)?;
4936        q = qn;
4937        let mut kn = e.zeros(n_head_kv * head_dim)?;
4938        e.rms_norm(
4939            &k,
4940            fa.k_norm.float_data(),
4941            &mut kn,
4942            head_dim,
4943            n_head_kv,
4944            eps,
4945        )?;
4946        k = kn;
4947        let rope_dims = geometry.n_rot as usize;
4948        e.rope_neox(
4949            &mut q,
4950            pos_d,
4951            head_dim,
4952            rope_dims,
4953            n_head,
4954            1,
4955            geometry.rope_base,
4956            1.0,
4957        )?;
4958        e.rope_neox(
4959            &mut k,
4960            pos_d,
4961            head_dim,
4962            rope_dims,
4963            n_head_kv,
4964            1,
4965            geometry.rope_base,
4966            1.0,
4967        )?;
4968
4969        let kv = scratch.plane_mut(scratch_index).0;
4970        // append at the DEVICE slot (kv.len_d == old len), then advance the counter in-graph.
4971        e.append_kv_quantized_dc(
4972            &k,
4973            &v,
4974            &mut kv.k,
4975            &mut kv.v,
4976            &kv.len_d,
4977            kv.kv_dim_k,
4978            kv.kv_dim_v,
4979            kv.k_tok_bytes,
4980            kv.v_tok_bytes,
4981            false,
4982        )?;
4983        e.inc_seqlen(&mut kv.len_d)?;
4984        // full-buffer views (any in-round t_kv stays in range on replay); the kernel bounds the
4985        // key range from the device counter.
4986        let k_view = e.view_u8(&kv.k, kv.k.len());
4987        let v_view = e.view_u8(&kv.v, kv.v.len());
4988        let (ktb, vtb) = (kv.k_tok_bytes, kv.v_tok_bytes);
4989        let mut attn = e.zeros(n_head * head_dim)?;
4990        e.fa_decode_dc(
4991            &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, &kv.len_d, bucket_max,
4992            scale, ktb, vtb, false,
4993        )?;
4994
4995        let attn_g = match &gate {
4996            Some(gate) => {
4997                let mut gsig = e.zeros(n_head * head_dim)?;
4998                e.sigmoid(gate, &mut gsig, n_head * head_dim)?;
4999                let mut ag = e.zeros(n_head * head_dim)?;
5000                e.mul(&attn, &gsig, &mut ag, n_head * head_dim)?;
5001                ag
5002            }
5003            None => attn,
5004        };
5005        Ok(e.matmul(&fa.wo, &attn_g, 1)?)
5006    }
5007
5008    /// PERSISTENT-DRAFT-KV fill (the reference engine's "mtp_update" analogue): compute the MTP
5009    /// block's K/V for `tokens` (committed tokens at positions pos0..pos0+T) from their EXACT
5010    /// trunk hiddens `h` ([T, n_embd] token-major, pre-output_norm) and append at slots pos0.. of
5011    /// the scratch KV. K/V-ONLY — ops A/1-5 plus the K-side of op 6 (wk/wv + k_norm + rope +
5012    /// quantized append); no wq/attention/FFN/lm_head, so per-token cost ~= eh_proj + wk/wv (a
5013    /// small fraction of one trunk layer), T-batched. Rope follows the chain convention
5014    /// rope(token@p) = p+1. Runs at round boundaries OUTSIDE the captured graph in BOTH draft
5015    /// modes -> draft parity by construction. Caller must have scratch.kv.len == pos0.
5016    #[allow(clippy::too_many_arguments)]
5017    fn mtp_kv_fill_at(
5018        &self,
5019        e: &Engine,
5020        mtp: &MtpHead,
5021        tokens: &[u32],
5022        h: &CudaSlice<f32>,
5023        pos0: usize,
5024        scratch: &mut MtpScratch,
5025        scratch_index: usize,
5026        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5027    ) -> Result<(), Box<dyn std::error::Error>> {
5028        let cfg = &self.cfg;
5029        let n_embd = cfg.n_embd as usize;
5030        let eps = cfg.rms_eps;
5031        let t = tokens.len();
5032        let (scratch_kv, scratch_cap) = scratch.plane(scratch_index);
5033        assert_eq!(scratch_kv.len, pos0, "mtp_kv_fill: append slot mismatch");
5034        assert!(pos0 + t <= scratch_cap, "mtp_kv_fill: scratch overflow");
5035        let Mixer::Full(fa) = &mtp.mixer else {
5036            panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5037        };
5038        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i + 1) as i32).collect();
5039        let pos_d = e.htod_i32(&pos_vec)?;
5040
5041        // ops A/1/2: embed + the two input norms, T-wide.
5042        let e_emb = match embd_dev {
5043            Some((g, qt, rb)) => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5044            None => e.htod(&self.embd.gather(n_embd, tokens))?,
5045        };
5046        let mut e_norm = e.zeros(t * n_embd)?;
5047        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, t, eps)?;
5048        let mut h_norm = e.zeros(t * n_embd)?;
5049        e.rms_norm(h, mtp.hnorm.float_data(), &mut h_norm, n_embd, t, eps)?;
5050
5051        // op 3: per-row [e_norm ; h_norm] concat, token-major [T, 2*n_embd].
5052        let mut concat = e.zeros(t * 2 * n_embd)?;
5053        for i in 0..t {
5054            e.copy_view_into(
5055                &mut concat,
5056                i * 2 * n_embd,
5057                &e_norm.slice(i * n_embd..(i + 1) * n_embd),
5058                n_embd,
5059            )?;
5060            e.copy_view_into(
5061                &mut concat,
5062                i * 2 * n_embd + n_embd,
5063                &h_norm.slice(i * n_embd..(i + 1) * n_embd),
5064                n_embd,
5065            )?;
5066        }
5067
5068        // ops 4/5: eh_proj + attn_norm, T-wide (at the student inner width when geom is set).
5069        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5070        let inp_sa = e.matmul(&mtp.eh_proj, &concat, t)?;
5071        let mut a_norm = e.zeros(t * di)?;
5072        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, t, eps)?;
5073
5074        // op 6 (K/V half): wk/wv + k_norm + rope + per-row quantized append. No wq/attention —
5075        // the fill only has to leave correct K/V rows behind for later chains to attend over.
5076        let n_head_kv = mtp
5077            .geom
5078            .as_ref()
5079            .map(|g| g.n_head_kv)
5080            .or(mtp.step35.as_ref().map(|s| s.n_head_kv))
5081            .unwrap_or_else(|| {
5082                let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5083                cfg.full_attention_geometry_at(mtp_il).n_head_kv as usize
5084            });
5085        let mtp_il = cfg.n_layer.saturating_sub(cfg.nextn_predict_layers);
5086        let geometry = cfg.full_attention_geometry_at(mtp_il);
5087        let head_dim = geometry.head_dim_k as usize;
5088        let mut k = e.matmul(&fa.wk, &a_norm, t)?;
5089        let v = e.matmul(&fa.wv, &a_norm, t)?;
5090        let mut kn = e.zeros(t * n_head_kv * head_dim)?;
5091        e.rms_norm(
5092            &k,
5093            fa.k_norm.float_data(),
5094            &mut kn,
5095            head_dim,
5096            n_head_kv * t,
5097            eps,
5098        )?;
5099        k = kn;
5100        // step35: rotary width AND base are per-layer, and the MTP block's values come from the
5101        // resolved `Step35MtpGeom` — NOT from `cfg.rope_dim_count`/`cfg.rope_freq_base`, which
5102        // carry the arch defaults (128 / 5e6, i.e. the FULL-attn layers' base). Getting this wrong
5103        // writes K rows the attention arm then re-derives at a different theta: correct-looking
5104        // output with dead acceptance, invisible to the exactness gates.
5105        let (rope_dims, rope_base, ff) = match mtp.step35.as_ref() {
5106            Some(s) => (
5107                s.n_rot,
5108                s.rope_base,
5109                if s.swa {
5110                    None
5111                } else {
5112                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
5113                },
5114            ),
5115            None => (geometry.n_rot as usize, geometry.rope_base, None),
5116        };
5117        #[cfg(debug_assertions)]
5118        if let Some(ff) = ff {
5119            crate::debug_assert_tensor_stream_device(ff, &e.stream(), "mtp_kv_fill.rope_freqs");
5120        }
5121        match ff {
5122            Some(f) => e.rope_neox_ff(
5123                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0, f,
5124            )?,
5125            None => e.rope_neox(
5126                &mut k, &pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
5127            )?,
5128        }
5129
5130        let kv = scratch.plane_mut(scratch_index).0;
5131        // Match the trunk prime contract: a chunk may need the aligned window immediately before
5132        // its first row, so preserve that prefix when the physical tail rebases at wrap.
5133        let retain_from = kv
5134            .ring
5135            .as_ref()
5136            .map(|ring| memra_kv::swa_retain_from(pos0, ring.window(), ring.base()))
5137            .unwrap_or(0);
5138        let write_row = e.prepare_kv_append(kv, retain_from, t)?;
5139        for i in 0..t {
5140            let k_row = k.slice(i * kv.kv_dim_k..(i + 1) * kv.kv_dim_k);
5141            let v_row = v.slice(i * kv.kv_dim_v..(i + 1) * kv.kv_dim_v);
5142            e.append_kv_quantized_view(
5143                &k_row,
5144                &v_row,
5145                &mut kv.k,
5146                &mut kv.v,
5147                write_row + i,
5148                kv.kv_dim_k,
5149                kv.kv_dim_v,
5150                kv.k_tok_bytes,
5151                kv.v_tok_bytes,
5152                false,
5153            )?;
5154        }
5155        kv.len = pos0 + t;
5156        e.set_i32_one(&mut kv.len_d, kv.len as i32)?;
5157        Ok(())
5158    }
5159
5160    #[allow(clippy::too_many_arguments)]
5161    fn mtp_kv_fill_all(
5162        &self,
5163        e: &Engine,
5164        tokens: &[u32],
5165        h: &CudaSlice<f32>,
5166        pos0: usize,
5167        scratch: &mut MtpScratch,
5168        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5169    ) -> Result<(), Box<dyn std::error::Error>> {
5170        debug_assert_eq!(self.mtp_head_count(), scratch.plane_count());
5171        for index in 0..self.mtp_head_count() {
5172            self.mtp_kv_fill_at(
5173                e,
5174                self.mtp_head_at(index),
5175                tokens,
5176                h,
5177                pos0,
5178                scratch,
5179                index,
5180                embd_dev,
5181            )?;
5182        }
5183        Ok(())
5184    }
5185
5186    /// CAPTURE body for the GRAPH DRAFT (stage 2 of graph-grade spec): ONE MTP head forward with
5187    /// every varying input device-resident —
5188    ///   - token id from the persistent `tok_d` (the previous replay's in-graph argmax wrote it,
5189    ///     so the chain feeds itself; the host reads the same 4 bytes for the draft list),
5190    ///   - h_seed from the persistent `h_seed_d` (h_nextn is copied BACK into it at the end),
5191    ///   - rope pos from the persistent `pos_d` counter (inc'd in-graph),
5192    ///   - scratch KV slot/bound from `scratch.kv.len_d` (see mtp_full_attn_dc).
5193    /// The p-min confidence lands in the persistent `p_d` iff `with_prob` (env is fixed per run).
5194    /// Same kernels, same dispatch as the eager mtp_head_forward_dev chain -> same draft tokens
5195    /// (exactness never depends on drafts — the verify arbitrates — but acceptance parity does).
5196    /// `with_head=false` captures the HEAD-LESS twin for the pseudo-seed replay (2026-07-03):
5197    /// the pseudo pass only needs h_nextn (op 10) + the scratch append — the lm_head read
5198    /// (~1.06ms q6_K on the 9B), argmax and prob are dead weight there. h_nextn's inputs are
5199    /// untouched, so the seed value is identical; round-start resets overwrite tok_d/p_d anyway.
5200    /// `sampled_cap` = Some((ctr_d, perturb_d, q_out_d, seed, temp)) captures the SAMPLED twin
5201    /// (step 3 of the sampled-spec arc): head logits are retained in the persistent `q_out_d`
5202    /// (host D2Ds them to the round's q slot after each replay), the DEVICE event counter is
5203    /// bumped in-graph, and the argmax reads GUMBEL-PERTURBED logits — one categorical draw per
5204    /// replay, bit-identical to the eager arm's gumbel_perturb at the same (seed, sctr, temp).
5205    /// seed/temp are capture-time constants (fixed per generate call, like p_min).
5206    #[allow(clippy::too_many_arguments)]
5207    fn mtp_head_forward_cap(
5208        &self,
5209        e: &Engine,
5210        mtp: &MtpHead,
5211        tok_d: &mut CudaSlice<u32>,
5212        pos_d: &mut CudaSlice<i32>,
5213        h_seed_d: &mut CudaSlice<f32>,
5214        p_d: &mut CudaSlice<f32>,
5215        scratch: &mut MtpScratch,
5216        // Which scratch plane this head appends to / attends over: 0 for the single-head
5217        // chain (every pre-lane caller), the head's own plane index for the multi-head
5218        // chain graphs (each head owns one plane — `mtp_chain_forward_dev`'s contract).
5219        scratch_index: usize,
5220        with_prob: bool,
5221        with_head: bool,
5222        embd_gpu: &CudaSlice<u8>,
5223        embd_qt: i32,
5224        embd_rb: usize,
5225        d_vocab: usize,
5226        sampled_cap: Option<SampledCapArgs<'_>>,
5227        stream_pack: Option<(&mut CudaSlice<u32>, usize, Option<&CudaSlice<u32>>)>,
5228        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask): (packed draft-vocab allowed-set buffer,
5229        // word count). Captured as ONE mask_logits_f32 node between the head matmul and the
5230        // in-graph argmax; the buffer address is baked, its CONTENTS are re-uploaded by the
5231        // host before every replay (the decode.rs graph-mask pattern). All-ones contents = a
5232        // no-op ban, so a position the grammar cannot constrain costs one pass over the row.
5233        mask_cap: Option<(&CudaSlice<u32>, usize)>,
5234    ) -> Result<(), Box<dyn std::error::Error>> {
5235        let cfg = &self.cfg;
5236        let n_embd = cfg.n_embd as usize;
5237        // step35: capturable through the WINDOWED device-counter arm (`mtp_step35_attn_dcw`)
5238        // once the dcw door is armed and the v3-vec class is live. Without the door this stays
5239        // the deliberate, named refusal: the plain `_dc` attention's key bound always starts at
5240        // row 0, cannot express this block's SWA view offset, and a captured chain would
5241        // silently attend OUTSIDE the window once the persistent scratch passes 512 rows.
5242        // Returning Err (not a panic) is what the capture sites already handle by degrading to
5243        // the eager chain (`mtp_head_forward_dev` -> `mtp_step35_attn`).
5244        // ROUND-STREAM stays refused EITHER WAY: the stream VERIFY has no step35 twin (see the
5245        // step35_verify refusal), so a stream capture that succeeded here would only move the
5246        // failure from capture time (graceful stream-off) to serve time (a failed round).
5247        if let Some(g) = mtp.step35.as_ref() {
5248            if stream_pack.is_some() {
5249                return Err(
5250                    "step35 has no ROUND-STREAM draft arm (the stream verify has no step35 \
5251                     twin); stream off"
5252                        .into(),
5253                );
5254            }
5255            if !self.step35_dcw_eligible(g, scratch.plane(scratch_index).1) {
5256                return Err(format!(
5257                    "step35 has no captured draft chain (fa_decode_dc cannot express the MTP \
5258                        block's SWA view offset; the windowed dcw capture needs \
5259                        MEMRA_STEP35_DRAFT_DCW armed [default ON, =0 disarms] and the v3-vec \
5260                        class live at bucket=min(window {}, scratch cap {})) - the eager draft \
5261                        chain serves this shape",
5262                    g.window,
5263                    scratch.plane(scratch_index).1,
5264                )
5265                .into());
5266            }
5267        }
5268        // student inner width (see mtp_head_forward_dev) — interface dims stay n_embd.
5269        let di = mtp.geom.as_ref().map(|g| g.d_inner).unwrap_or(n_embd);
5270        let eps = cfg.rms_eps;
5271        let e_emb = e.embed_gather_device(embd_gpu, tok_d, n_embd, embd_qt, embd_rb)?;
5272        let mut e_norm = e.zeros(n_embd)?;
5273        e.rms_norm(&e_emb, mtp.enorm.float_data(), &mut e_norm, n_embd, 1, eps)?;
5274        let mut h_norm = e.zeros(n_embd)?;
5275        e.rms_norm(
5276            &*h_seed_d,
5277            mtp.hnorm.float_data(),
5278            &mut h_norm,
5279            n_embd,
5280            1,
5281            eps,
5282        )?;
5283        let mut concat = e.zeros(2 * n_embd)?;
5284        e.copy_into(&mut concat, 0, &e_norm, n_embd)?;
5285        e.copy_into(&mut concat, n_embd, &h_norm, n_embd)?;
5286        let inp_sa = e.matmul(&mtp.eh_proj, &concat, 1)?;
5287        let mut a_norm = e.zeros(di)?;
5288        e.rms_norm(&inp_sa, mtp.attn_norm.float_data(), &mut a_norm, di, 1, eps)?;
5289        let attn_out = match (&mtp.mixer, mtp.step35.as_ref()) {
5290            // step35 (eligibility already enforced by the refusal above): the windowed dcw
5291            // arm, the SAME launcher the eager dev arm runs when the door is armed. No host
5292            // work here (this is the capture body); headroom is the callers' pre-arm.
5293            (Mixer::Full(fa), Some(g)) => {
5294                self.mtp_step35_attn_dcw(e, fa, g, &a_norm, pos_d, scratch, scratch_index)?
5295            }
5296            (Mixer::Full(fa), None) => self.mtp_full_attn_dc(
5297                e,
5298                fa,
5299                &a_norm,
5300                pos_d,
5301                scratch,
5302                scratch_index,
5303                mtp.geom.as_ref(),
5304            )?,
5305            (Mixer::Linear(_), _) => {
5306                panic!("MTP block is full-attn in qwen35; linear MTP not supported")
5307            }
5308            (Mixer::Mla(_), _) => crate::hybrid::mla_forward_unimplemented(),
5309        };
5310        let mut x1 = e.zeros(di)?;
5311        e.add(&inp_sa, &attn_out, &mut x1, di)?;
5312        let mut z = e.zeros(di)?;
5313        e.rms_norm(&x1, mtp.post_attn_norm.float_data(), &mut z, di, 1, eps)?;
5314        let ffn_out = match &mtp.ffn {
5315            crate::hybrid::Ffn::Dense {
5316                ffn_gate,
5317                ffn_up,
5318                ffn_down,
5319            } => {
5320                let n_ff = ffn_gate.out_features();
5321                let (gate, up) = if e.uses_q8_1_fast(ffn_gate) && e.uses_q8_1_fast(ffn_up) {
5322                    let (zq, zd) = e.quantize_q8_1(&z, 1, di)?;
5323                    (
5324                        e.matmul_pre(ffn_gate, &zq, &zd, &z, 1)?,
5325                        e.matmul_pre(ffn_up, &zq, &zd, &z, 1)?,
5326                    )
5327                } else {
5328                    (e.matmul(ffn_gate, &z, 1)?, e.matmul(ffn_up, &z, 1)?)
5329                };
5330                let mut act = e.zeros(n_ff)?;
5331                // step35: the dense FFN reads the per-layer SHEXP clamp, resolved for the MTP
5332                // block's own index (the mtp_head_forward_dev rule; None for every other arch,
5333                // which is `ffn_act`'s dispatch verbatim). The eager and captured chains must
5334                // run the ONE activation program.
5335                Self::ffn_act_lim(
5336                    e,
5337                    &self.cfg,
5338                    &gate,
5339                    &up,
5340                    1.0,
5341                    1.0,
5342                    mtp.step35.as_ref().and_then(|s| s.clamp_shexp),
5343                    &mut act,
5344                    n_ff,
5345                )?;
5346                e.matmul(ffn_down, &act, 1)?
5347            }
5348            // ROUND-STREAM: the 35B NextN block carries a MoE FFN. With RESIDENT experts the
5349            // dev path is pure device launches (device top-k + rows kernels, ZERO-DtoH by
5350            // design) — capture-legal. Non-resident (SLRU-lock) stays rejected: the capture
5351            // error arm degrades the caller to eager/stream-off.
5352            crate::hybrid::Ffn::Moe(m) if m.dev_exps.is_some() => {
5353                self.moe_ffn_il(e, m, &z, 1, u16::MAX)?
5354            }
5355            crate::hybrid::Ffn::Moe(_) => {
5356                return Err("graph draft requires a Dense (or resident-MoE) MTP FFN".into());
5357            }
5358        };
5359        let mut h_inner = e.zeros(di)?;
5360        e.add(&x1, &ffn_out, &mut h_inner, di)?;
5361        // student: up-project back to n_embd (carrier + head input; see mtp_head_forward_dev).
5362        let h_nextn = match mtp.geom.as_ref() {
5363            Some(g) => e.matmul(&g.out_up, &h_inner, 1)?,
5364            None => h_inner,
5365        };
5366        // MEMRA_SPEC_HPOST needs final_h even head-less (it IS the next seed under that convention).
5367        let final_h = if with_head || spec_hpost() {
5368            let final_norm = mtp.shared_head_norm.as_ref().unwrap_or(&self.output_norm);
5369            let mut fh = e.zeros(n_embd)?;
5370            e.rms_norm(&h_nextn, final_norm.float_data(), &mut fh, n_embd, 1, eps)?;
5371            Some(fh)
5372        } else {
5373            None
5374        };
5375        if with_head {
5376            let head = mtp.shared_head_head.as_ref().unwrap_or(&self.output);
5377            let mut logits = e.matmul(head, final_h.as_ref().unwrap(), 1)?;
5378            // DRAFT-SIDE GRAMMAR MASK: ban the grammar-illegal draft ids IN the captured chain,
5379            // before the argmax — proposals become legal by construction. Contents-only
5380            // per-replay upload keeps the capture valid.
5381            if let Some((mask_d, mw)) = mask_cap {
5382                e.mask_logits_col(&mut logits, mask_d, 0, d_vocab, mw)?;
5383            }
5384            if let Some(SampledCapArgs {
5385                ctr: ctr_d,
5386                perturb: perturb_d,
5387                q_out: q_out_d,
5388                seed,
5389                temp,
5390                filt,
5391            }) = sampled_cap
5392            {
5393                // SAMPLED chain: retain q (raw head logits -> persistent q_out_d; the matmul's
5394                // own buffer is pool-recycled after the capture body returns, so it can't be the
5395                // retention target), bump the device event counter, gumbel-perturb reading it,
5396                // and argmax the PERTURBED logits into tok_d — the in-graph categorical draw.
5397                e.copy_into(q_out_d, 0, &logits, d_vocab)?;
5398                e.sctr_inc(ctr_d)?;
5399                match filt {
5400                    // PURE-TEMP: gumbel over the raw softmax — byte-identical to the
5401                    // pre-lane capture body.
5402                    None => e.gumbel_perturb_ctr(&logits, perturb_d, d_vocab, seed, ctr_d, temp)?,
5403                    // FILTERED (lane/step37-draft-graph-serving-20260830): the SAME
5404                    // filter_stats program the eager arm and the accept path run (the
5405                    // wrapper's coop/plain choice is deployment-keyed, never per-call), then
5406                    // the device-stat/device-counter perturb twin — the draft draws from the
5407                    // exact filtered distribution the verify gathers `q` from. q was
5408                    // retained ABOVE, pre-perturb, so the accept path's post-replay stats
5409                    // recompute (same kernel, same bits) reconstructs these th/z exactly.
5410                    Some(f) => {
5411                        e.filter_stats(
5412                            &logits, d_vocab, f.rows0, f.th, f.z, f.mx, d_vocab, 1, temp, f.top_k,
5413                            f.top_p, f.min_p,
5414                        )?;
5415                        e.gumbel_perturb_filtered_ctr(
5416                            &logits, perturb_d, d_vocab, seed, ctr_d, temp, f.mx, f.th,
5417                        )?;
5418                    }
5419                }
5420                e.argmax_token_device_into(perturb_d, tok_d, d_vocab)?;
5421                // p-min prob = the head's RAW softmax confidence in the SAMPLED pick — same
5422                // semantics as the eager sampled arm's prob_of_token_device(dl_d, tok_d).
5423                if with_prob {
5424                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5425                }
5426            } else {
5427                // draft token -> persistent tok_d (next replay's embed reads it; host reads the 4 bytes).
5428                e.argmax_token_device_into(&logits, tok_d, d_vocab)?;
5429                // p-min under a draft mask reads the MASKED row: confidence relative to the
5430                // grammar-LEGAL alternatives (illegal ids leave the softmax denominator), which
5431                // is the right semantics for "does the drafter know what comes next here" and
5432                // the same row the pick came from. Draft-quality only — verify arbitrates.
5433                if with_prob {
5434                    e.prob_of_token_device_into(&logits, tok_d, p_d, d_vocab)?;
5435                }
5436            }
5437        }
5438        // ROUND-STREAM K-chain: pack (tok, p) into slot j, then remap tok through d2t so the
5439        // NEXT chained body's embed reads the TARGET id — zero host involvement per step.
5440        if let Some((out, slot, d2t)) = stream_pack {
5441            e.pack_tok_p(tok_d, p_d, out, slot)?;
5442            if let Some(map) = d2t {
5443                e.tok_map_u32(tok_d, map)?;
5444            }
5445        }
5446        // Next draft step's h_seed: pre-norm h_nextn (default) or post-norm final_h (HPOST).
5447        if spec_hpost() {
5448            e.copy_into(h_seed_d, 0, final_h.as_ref().unwrap(), n_embd)?;
5449        } else {
5450            e.copy_into(h_seed_d, 0, &h_nextn, n_embd)?;
5451        }
5452        // advance the draft rope position in-graph.
5453        e.inc_seqlen(pos_d)?;
5454        Ok(())
5455    }
5456
5457    /// Batched target verify forward over `tokens` at positions `pos0..pos0+T` (§D.3, T=K+1).
5458    /// Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV
5459    /// and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1).
5460    /// Advances `cache.pos` by T.
5461    pub fn decode_step_t(
5462        &self,
5463        e: &Engine,
5464        tokens: &[u32],
5465        pos0: usize,
5466        cache: &mut Cache,
5467    ) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
5468        if self.is_gemma4_e4b() {
5469            return Ok(self.gemma4_e4b_decode_step_t_h(e, tokens, pos0, cache)?.0);
5470        }
5471        if self.gemma_batch_program() {
5472            return self.gemma4_decode_step_t(e, tokens, pos0, cache);
5473        }
5474        Ok(self.decode_step_t_h(e, tokens, pos0, cache)?.0)
5475    }
5476
5477    /// Like `decode_step_t` but ALSO returns the LAST column's pre-output_norm hidden (h_seed for
5478    /// the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward
5479    /// (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads).
5480    /// At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.
5481    pub fn decode_step_t_h(
5482        &self,
5483        e: &Engine,
5484        tokens: &[u32],
5485        pos0: usize,
5486        cache: &mut Cache,
5487    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5488        self.decode_step_t_h_emb(e, tokens, pos0, cache, None)
5489    }
5490
5491    /// Like `decode_step_t_h` with an optional RESIDENT embed table (spec hot loop): device
5492    /// gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.
5493    pub fn decode_step_t_h_emb(
5494        &self,
5495        e: &Engine,
5496        tokens: &[u32],
5497        pos0: usize,
5498        cache: &mut Cache,
5499        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5500    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5501        let (logits_d, h_seed) = self.decode_step_t_h_emb_dev(e, tokens, pos0, cache, embd_dev)?;
5502        Ok((e.dtoh(&logits_d)?, h_seed))
5503    }
5504
5505    /// DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to
5506    /// `decode_step_t_h_emb` but returns the [T, n_vocab] logits ON DEVICE — the accept walk
5507    /// argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh'ing the full
5508    /// T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is
5509    /// UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.
5510    pub fn decode_step_t_h_emb_dev(
5511        &self,
5512        e: &Engine,
5513        tokens: &[u32],
5514        pos0: usize,
5515        cache: &mut Cache,
5516        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5517    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5518        let n_embd = self.cfg.n_embd as usize;
5519        let t = tokens.len();
5520        let (logits, x) = self.decode_step_t_core(e, tokens, pos0, cache, embd_dev, None)?;
5521        // h_seed for the next round = LAST column's pre-output_norm hidden ([n_embd]).
5522        let mut hs = vbuf(e, n_embd)?; // fully written by copy_view_into below
5523        e.copy_view_into(&mut hs, 0, &x.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
5524        Ok((logits, hs))
5525    }
5526
5527    /// CORE verify forward: the `decode_step_t_h_emb_dev` kernel chain, returning the FULL
5528    /// pre-output_norm hidden stack x ([T, n_embd], any column extractable) and optionally
5529    /// filling a `VerifyCkpt` (retained per-layer state-rebuild inputs) for the REPLAY-FREE
5530    /// partial accept. `ckpt: None` => byte-for-byte the old behavior (the ckpt writes are pure
5531    /// retains/copies — they never change what any kernel computes).
5532    fn decode_step_t_core(
5533        &self,
5534        e: &Engine,
5535        tokens: &[u32],
5536        pos0: usize,
5537        cache: &mut Cache,
5538        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5539        mut ckpt: Option<&mut VerifyCkpt>,
5540    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5541        self.decode_step_t_core_stream(
5542            e,
5543            tokens,
5544            pos0,
5545            cache,
5546            embd_dev,
5547            ckpt.take(),
5548            None,
5549            None,
5550            None,
5551            None,
5552        )
5553    }
5554
5555    /// [`Self::decode_step_t_core`] with the MTP route's verify-graph pool armed
5556    /// (`MEMRA_SPEC_VERIFY_GRAPH`). `graphs: None` reproduces `decode_step_t_core`
5557    /// argument-for-argument, so the eager walk stays the byte-identical fallback.
5558    fn decode_step_t_core_vg(
5559        &self,
5560        e: &Engine,
5561        tokens: &[u32],
5562        pos0: usize,
5563        cache: &mut Cache,
5564        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5565        mut ckpt: Option<&mut VerifyCkpt>,
5566        graphs: Option<&mut DsparkVerifyGraphs>,
5567    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5568        self.decode_step_t_core_stream(
5569            e,
5570            tokens,
5571            pos0,
5572            cache,
5573            embd_dev,
5574            ckpt.take(),
5575            None,
5576            None,
5577            None,
5578            graphs,
5579        )
5580    }
5581
5582    /// Increment-0 two-session PP seam: release the peer after this lane's stage-0 boundary TX.
5583    /// The two independent sessions keep their own cache/checkpoint state; only issue order moves.
5584    fn decode_step_t_core_pipelined(
5585        &self,
5586        e: &Engine,
5587        tokens: &[u32],
5588        pos0: usize,
5589        cache: &mut Cache,
5590        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5591        mut ckpt: Option<&mut VerifyCkpt>,
5592        pipe: &SpecPipeLane,
5593        round: usize,
5594    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5595        let fence = crate::pp::pp_cuts(self.layers.len())
5596            .ok_or("two-session speculative pipeline requires a PP stage cut")?;
5597        if crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
5598            return Err("two-session speculative pipeline requires the PP verify split".into());
5599        }
5600        let interval_fence = pipe.stage0_begin(round)?;
5601        let ticket = self.verify_stage0_issue(
5602            e,
5603            tokens,
5604            pos0,
5605            cache,
5606            embd_dev,
5607            ckpt.as_deref_mut(),
5608            None,
5609            &fence,
5610            Some(interval_fence),
5611            pipe.trace(round),
5612        )?;
5613        pipe.stage0_end(round);
5614        pipe.stage1_begin(round)?;
5615        let result = self.verify_stage1_finish(e, ticket, cache, ckpt, None, &fence, true)?;
5616        pipe.verify_end(round);
5617        Ok(result)
5618    }
5619
5620    /// ROUND-STREAM stage (c) 4: `stream` = (device verify tokens [t], device pos counter) —
5621    /// when Some, rope positions come from pos_iota over the counter, the embed gathers the
5622    /// device tokens, and full_attn_verify routes appends/FA through the _dc twins reading the
5623    /// SAME counter (every layer's kvl.len == cache.pos, one counter drives all three). The
5624    /// host `tokens`/`pos0` args still size buffers (t is FIXED K+1 in stream mode).
5625    /// `vtok_dev` (engine-bundle slice 2): device verify tokens for the EMBED only —
5626    /// unlike `stream` mode it changes nothing else (host pos iota, host-len KV appends).
5627    /// `tokens` then only sizes buffers (the dummy-slice pattern the round-stream arm uses).
5628    #[allow(clippy::too_many_arguments)]
5629    fn decode_step_t_core_stream(
5630        &self,
5631        e: &Engine,
5632        tokens: &[u32],
5633        pos0: usize,
5634        cache: &mut Cache,
5635        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5636        mut ckpt: Option<&mut VerifyCkpt>,
5637        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5638        pp_pipe: Option<bool>,
5639        vtok_dev: Option<&CudaSlice<u32>>,
5640        graphs: Option<&mut DsparkVerifyGraphs>,
5641    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5642        // PP DOOR (lane/pp2-spec 2026-08-06): the verify trunk now takes its OWN stage split,
5643        // exactly as the eager and batched steps do. This is the single funnel every verify
5644        // forward reaches (decode_step_t / _h / _h_emb / _h_emb_dev / _core all land here), so
5645        // wiring it here wires the whole spec surface — the draft/accept/commit machinery above
5646        // is untouched.
5647        //
5648        // History: pp2-hardening (2026-08-06) made this funnel FAIL CLOSED, because its trunk
5649        // walk was unsplit on one stream and a sharded cross-device placement peer-read every
5650        // remote layer's weights on every spec round (measured 13.9-28x on the batched twin).
5651        // The refusal below survives to cover the residue — MEMRA_SPEC_PP=0, MEMRA_PP_STREAMS=0,
5652        // or a placement whose PpNRt fails to build — so a config that would still walk the
5653        // whole trunk on one stream refuses instead of regressing 28x.
5654        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
5655            if !crate::pp::pp2_streams_off() && crate::pp::spec_pp_on() {
5656                if vtok_dev.is_some() {
5657                    return Err(
5658                        "device-token dspark verify (slice-2 deferred readback) has no PP \
5659                         stage-split arm; set MEMRA_DSPARK_DEFER_READBACK=0 or run the dspark \
5660                         route on one device"
5661                            .into(),
5662                    );
5663                }
5664                return self.decode_step_t_core_ppn(
5665                    e,
5666                    tokens,
5667                    pos0,
5668                    cache,
5669                    embd_dev,
5670                    ckpt.take(),
5671                    stream,
5672                    &fence,
5673                    pp_pipe,
5674                );
5675            }
5676        }
5677        crate::pp::refuse_unsplit_if_remote(
5678            "decode_step_t (spec verify)",
5679            "drop MEMRA_SPEC_PP=0 / MEMRA_PP_STREAMS=0 so the verify trunk takes its OWN stage \
5680             split (decode_step_t_core_ppn); or run spec on one device",
5681        )?;
5682        let cfg = &self.cfg;
5683        let n_embd = cfg.n_embd as usize;
5684        let eps = cfg.rms_eps;
5685        let t = tokens.len();
5686        let pos_d = match stream {
5687            Some((_, ctr)) => {
5688                let mut p = e.alloc_uninit::<i32>(t)?;
5689                e.pos_iota(ctr, &mut p, t)?;
5690                p
5691            }
5692            None => {
5693                let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5694                e.htod_i32(&pos_vec)?
5695            }
5696        };
5697
5698        // embed T tokens -> [T, n_embd] token-major (device gather on the spec hot loop)
5699        let x = match (stream, embd_dev) {
5700            (Some((vtok, _)), Some((g, qt, rb))) => {
5701                e.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
5702            }
5703            (None, Some((g, qt, rb))) => match vtok_dev {
5704                // slice 2: device verify tokens, same embed_gather_u32_t kernel —
5705                // bit-identical rows to the host-token arm (same per-dtype deq).
5706                Some(vt_d) => e.embed_gather_device_td(g, vt_d, t, n_embd, qt, rb)?,
5707                None => e.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
5708            },
5709            _ => {
5710                assert!(
5711                    vtok_dev.is_none(),
5712                    "device-token verify requires the resident embed table (embd_dev)"
5713                );
5714                e.htod(&self.embd.gather(n_embd, tokens))?
5715            }
5716        };
5717
5718        // TRUNK WALK: layers [0, n_layers) through the SAME range-scoped subgraph the PP-N
5719        // stage split calls per stage (`verify_layers`) — one code path, so the split cannot
5720        // drift from the unsplit dispatch mirroring. lane/pp2-spec 2026-08-06.
5721        let x = self.verify_layers(
5722            e,
5723            x,
5724            0,
5725            self.layers.len(),
5726            &pos_d,
5727            pos0,
5728            t,
5729            cache,
5730            ckpt.take(),
5731            stream,
5732            graphs,
5733        )?;
5734        if spec_nan_scan() {
5735            nan_scan_rows(e, &x, t, n_embd, &format!("verify trunk exit pos0={pos0}"))?;
5736        }
5737
5738        let mut hn = vbuf(e, t * n_embd)?;
5739        // Stage-A door: with the serving-class row-outer verify walk, the TAIL must be the
5740        // t=1 decode program per row too (rms_norm t=1 + the single-row bf16 head — the
5741        // split head's concat is receipted bit-identical to it). The batched cuBLASLt head
5742        // is a different ULP class and flips near-tie argmaxes off the greedy tape.
5743        let eager_tail = self.sliding_gated_moe_batch_program() && spec_verify_eager_on();
5744        if eager_tail {
5745            let n_vocab = self.cfg.n_vocab as usize;
5746            // MEMRA_SPEC_HEAD_ROWS=1 — THE VERIFY TAIL'S REDUNDANT HEAD READ.
5747            //
5748            // The loop below runs the head at m=1 once PER COLUMN, so the LM head's weights are
5749            // streamed t times per verify pass. On step37 that head is ~0.49 GiB per card after the
5750            // rank split, ~1.07 ms of pure re-read at t=2 and worse at every wider t — which is a
5751            // large part of why the fixed K ladder LOSES (K=1 81.2 > K=2 73.1 > K=3 62.7 tok/s).
5752            //
5753            // The loop's justification is the comment above: the batched cuBLASLt head is a
5754            // different ULP class and flips near-tie argmaxes off the greedy tape. That is true of
5755            // cuBLASLt and it does NOT apply here, because a FloatBf16 head at 1..=32 rows never
5756            // reaches cuBLASLt: `matmul` routes it to `matvec_bf16_rows_into` (lib.rs:12248), whose
5757            // own doc says `matvec_bf16_f32acc_x4_rows` "runs the t=1 decode head program PER ROW
5758            // (identical dot + reduce), so decode/verify tiers keep the t=1 numeric class". Under
5759            // the W8 doors both widths route to the q8 mirror instead, and the t-column mirror is
5760            // documented "bit-identical to t single-row calls". So the batched form is the SAME
5761            // arithmetic per row on both paths, with one weight read instead of t.
5762            //
5763            // rms_norm is row-wise, so norm(t) is per-row identical to t x norm(1) by construction.
5764            //
5765            // DEFAULT OFF for exactly one turn of the crank: "bit-identical by two documented
5766            // claims" is still an argument. The greedy byte tape decides, and the door flips only
5767            // once the tape is a receipt.
5768            if head_rows_on() {
5769                e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5770                let logits = e.matmul(&self.output, &hn, t)?;
5771                if stream.is_none() {
5772                    cache.pos += t;
5773                }
5774                return Ok((logits, if spec_hpost() { hn } else { x }));
5775            }
5776            let mut logits = vbuf(e, t * n_vocab)?;
5777            for r in 0..t {
5778                let mut row = e.uninit(n_embd)?;
5779                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
5780                let mut hr = e.uninit(n_embd)?;
5781                e.rms_norm(&row, self.output_norm.float_data(), &mut hr, n_embd, 1, eps)?;
5782                let lr = e.matmul(&self.output, &hr, 1)?;
5783                e.dtod_copy_into(&lr, &mut logits, r * n_vocab)?;
5784                e.dtod_copy_into(&hr, &mut hn, r * n_embd)?;
5785            }
5786            if stream.is_none() {
5787                cache.pos += t;
5788            }
5789            return Ok((logits, if spec_hpost() { hn } else { x }));
5790        }
5791        let serving_head =
5792            self.sliding_gated_moe_batch_program() || self.batched_serving_numeric_class();
5793        let logits = if serving_head {
5794            // Step35 and the qwen35 family (MoE 2026-08-14 AM, dense-hybrid same day PM — the
5795            // Q3.8 bring-up reproduced the identical near-tie class on dense: eager-class verify
5796            // vs batched-class live serving, ULP drift amplified through the GDN recurrence)
5797            // serve one batched numeric class at every live width, including B=1. Keep the
5798            // verify head in that same class; other generic families retain the decode-exact
5799            // head that their run-spec contract pins.
5800            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5801            e.matmul(&self.output, &hn, t)?
5802        } else {
5803            e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
5804            e.matmul_decode_exact(&self.output, &hn, t)?
5805        };
5806        // stream: the device pos counter owns position; host mirror reconciles at drain.
5807        if stream.is_none() {
5808            cache.pos += t;
5809        }
5810        // Hidden stack for seeds/refresh-fills: pre-norm x (default) or post-norm hn (HPOST).
5811        Ok((logits, if spec_hpost() { hn } else { x }))
5812    }
5813
5814    /// THE VERIFY TRUNK OVER PP-N (lane/pp2-spec 2026-08-06): `decode_step_t_core_stream`'s walk
5815    /// as N stage subgraphs, each on its own engine/stream (and, under `MEMRA_PP_DEVICES`, its own
5816    /// device), with a `[T, n_embd]` boundary transfer between them. T = K+1 (the verify batch),
5817    /// so this is the batched-boundary shape the pp2-batch lane's grow-only slots already handle
5818    /// (`tx(b, x, t*n_embd)`; the slot grows to the high-water T and the transport moves exactly
5819    /// the payload).
5820    ///
5821    /// Structure is `decode_step_batch_ppn`'s, which is `decode_step_h_ppn`'s. FOUR THINGS ARE
5822    /// PER-STAGE and each for a measured reason (see `decode_step_batch_ppn`'s header for the
5823    /// receipts):
5824    ///
5825    /// 1. THE ENGINE (`rt.engine(s, e)`) — `Engine` owns lazily-grown stable-pointer scratch
5826    ///    (`fa_part_pool`, `fa_vf16_scratch`, `argmax_partials`) that is single-stream-safe BY
5827    ///    DESIGN. Two stage streams through one Engine is the 2026-08-02 shared-scratch race
5828    ///    (35% flake, nondeterministic all-logits divergence). `PpNRt::build` gives every stage
5829    ///    s>0 its own Engine even on the primary device; honouring it here is what scopes the
5830    ///    pools. The verify path allocates MORE of that scratch than eager decode does (FA at
5831    ///    m=T, and the per-layer `GdnStash` retains), so this is load-bearing, not inherited.
5832    ///
5833    /// 2. `pos_d` — each stage uploads its OWN copy of the T rope positions on ITS stream, so the
5834    ///    buffer is allocated, consumed and freed on one stream. In `stream` mode that means each
5835    ///    stage runs its own `pos_iota` over the SHARED device counter (`pos_ctr`): the counter is
5836    ///    read-only during the forward (the round's `inc`/`copy_add` happen outside it), so every
5837    ///    stage derives the identical iota, and each stage's own output buffer is stream-local.
5838    ///
5839    /// 3. THE EMBED lives with stage 0 (`self.embd` / `embd_gpu` are host/primary-side; the
5840    ///    sharded loader leaves the table with stage 0 by construction).
5841    ///
5842    /// 4. THE HEAD (`output_norm` + `output`) runs on the LAST stage — the sharded loader uploaded
5843    ///    both through that stage's engine (`hybrid.rs`: `e_head = layer_engine(e, n_trunk,
5844    ///    n_trunk-1)`), so reading them anywhere else is a peer read of the biggest tensor in the
5845    ///    model, every round.
5846    ///
5847    /// WHAT STAYS ON THE PRIMARY, deliberately: the returned logits and hidden stack `x`. Both are
5848    /// last-stage-allocated device buffers, and every consumer (the device argmax walk, the accept
5849    /// kernels, `spec_seed_gather`, the ckpt rebuild in `commit_verified_prefix`) reads them
5850    /// through the primary context by UVA — the same read the batched serving epilogue's
5851    /// `last_logits_dev` park does. Those consumers are per-round O(T x n_vocab) and O(n_embd),
5852    /// not per-layer, so they are not the 28x class; splitting them is a separate lane.
5853    ///
5854    /// The MTP HEAD (draft side) is NOT split: it is one block, it lives wherever the loader put
5855    /// it (`load_mtp` uses the primary engine), and it is ~1-2 GB against the trunk's tens. Draft
5856    /// placement is measured, not assumed — see `research/pp2-spec-20260806`.
5857    ///
5858    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME bytes in
5859    /// the same order via the SAME `verify_layers` the unsplit body calls; the only change is
5860    /// where the residual is materialized, and the boundary is a straight f32 copy (dtod
5861    /// same-device / `cudaMemcpyPeerAsync` cross-device, no conversion). So the split MUST be
5862    /// BIT-IDENTICAL to the unsplit verify at the same T, in both placement orders. Gate:
5863    /// `decode-batch-gate --mode ppspec`. Acceptance counts are a DERIVED consequence — greedy
5864    /// accept argmaxes these logits, so bit-identical logits force identical accept walks; the
5865    /// `run-spec` K=1..8 arm checks that end-to-end rather than trusting the implication.
5866    #[allow(clippy::too_many_arguments)]
5867    fn decode_step_t_core_ppn(
5868        &self,
5869        e: &Engine,
5870        tokens: &[u32],
5871        pos0: usize,
5872        cache: &mut Cache,
5873        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5874        mut ckpt: Option<&mut VerifyCkpt>,
5875        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5876        fence: &[usize],
5877        pp_pipe: Option<bool>,
5878    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
5879        let ticket = self.verify_stage0_issue(
5880            e,
5881            tokens,
5882            pos0,
5883            cache,
5884            embd_dev,
5885            ckpt.as_deref_mut(),
5886            stream,
5887            fence,
5888            pp_pipe,
5889            None,
5890        )?;
5891        self.verify_stage1_finish(e, ticket, cache, ckpt, stream, fence, true)
5892    }
5893
5894    /// Enqueue embed, stage 0, and the first boundary TX, then return the actual boundary slot.
5895    /// The ordinary PP verify wrapper calls `verify_stage1_finish` immediately after this return.
5896    #[allow(clippy::too_many_arguments)]
5897    fn verify_stage0_issue(
5898        &self,
5899        e: &Engine,
5900        tokens: &[u32],
5901        pos0: usize,
5902        cache: &mut Cache,
5903        embd_dev: Option<(&CudaSlice<u8>, i32, usize)>,
5904        mut ckpt: Option<&mut VerifyCkpt>,
5905        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
5906        fence: &[usize],
5907        pp_pipe: Option<bool>,
5908        trace: Option<SpecPipeTraceCtx>,
5909    ) -> Result<VerifyBoundaryTicket, Box<dyn std::error::Error>> {
5910        assert!(
5911            !self.is_gemma4_e4b() && !self.gemma_batch_program(),
5912            "decode_step_t_core_ppn covers the hybrid non-gemma4 verify trunk only \
5913             (the gemma4 arms have their own decode_step_t twins)"
5914        );
5915        if crate::pp::pp_host_bounce_active() && (stream.is_some() || embd_dev.is_some()) {
5916            return Err(
5917                "decode_step_t_core_ppn: refused with MEMRA_PP_HOST_BOUNCE=1 — the trunk \
5918                 boundary itself is host-staged, but device-resident verify still peer-reads \
5919                 primary-device token/position/embedding buffers from stage 0. Run plain PP \
5920                 serving on this host class; spec requires local per-stage inputs first."
5921                    .into(),
5922            );
5923        }
5924        let rt = crate::pp::PpNRt::get(e)?;
5925        let n_st = fence.len() - 1;
5926        assert_eq!(
5927            rt.n_stages(),
5928            n_st,
5929            "PpNRt stage count {} != fence stages {n_st}",
5930            rt.n_stages()
5931        );
5932        let n_embd = self.cfg.n_embd as usize;
5933        let t = tokens.len();
5934        let payload = t * n_embd;
5935        if pp_pipe.is_some() {
5936            assert_eq!(n_st, 2, "spec pipeline requires exactly two PP stages");
5937        }
5938        // One-shot lane diagnostic: force natural PP-2 boundaries to completion so the server
5939        // log can price stage 0, the peer hop, the RX copy, and stage 1 + head separately without
5940        // nsys. The ordinary path keeps every enqueue asynchronous. N>2 is deliberately excluded:
5941        // the report below names exactly two stages and must never imply it measured middle ones.
5942        let pp_anatomy = n_st == 2 && std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
5943        let pp_started = std::time::Instant::now();
5944        let (mut reverse_ms, mut stage0_ms, mut tx_ms) = (0.0f64, 0.0f64, 0.0f64);
5945        // The CALLER's ambient stream, captured BEFORE any `rt.enter()` pushes a stage stream:
5946        // this body returns DEVICE-RESIDENT buffers (the device-argmax accept walk's contract),
5947        // so the exit needs the same publication the boundaries get — see `PpNRt::publish_to`.
5948        // Taken here, not at the end, because inside the last-stage scope `e.stream()` IS the
5949        // stage stream and the wait would self-order into a no-op.
5950        let caller_stream = e.stream();
5951        // #87 ROOT-CAUSE FENCE (lane/pp2spec-crash): the PREVIOUS round's stage-allocated
5952        // outputs (logits/hidden/ckpt stashes) freed stream-ordered on the STAGE streams while
5953        // the primary stream still holds queued reads of them — with event tracking elided,
5954        // nothing stops the pool from reusing those blocks for THIS round's stage allocations,
5955        // whose writes then race the queued reads (measured: 13/4096-NaN random-bits garbage in
5956        // the spec round seed; the full anatomy is on `PpNRt::fence_stages_behind`). Order every
5957        // stage stream behind the caller before enqueueing new stage work.
5958        let reverse_started = std::time::Instant::now();
5959        if pp_pipe != Some(false) {
5960            rt.fence_stages_behind(&caller_stream)?;
5961        }
5962        if pp_pipe == Some(true) {
5963            // Both session verifies must alternate boundary slots even when the ordinary
5964            // decode overlap experiment is off. Prewarm before A's stage 0 so B cannot grow
5965            // slot 1 by synchronizing the RX stream while A's stage 1 is in flight.
5966            rt.prepare_overlap_slots(0, payload)?;
5967        }
5968        if pp_anatomy {
5969            // Drain the reverse-publication dependency before timing stage 0 itself. At c=1 this
5970            // prices any primary-stream rollback/refresh tail inherited from the prior round.
5971            for s in 0..n_st {
5972                let _st = rt.enter(s);
5973                rt.engine(s, e).stream().synchronize()?;
5974            }
5975            reverse_ms = reverse_started.elapsed().as_secs_f64() * 1e3;
5976        }
5977
5978        // Per-stage rope positions: in host mode the same [T] iota each stage uploads itself; in
5979        // stream mode each stage's own `pos_iota` over the shared read-only device counter.
5980        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
5981            match stream {
5982                Some((_, ctr)) => {
5983                    let mut p = es.alloc_uninit::<i32>(t)?;
5984                    es.pos_iota(ctr, &mut p, t)?;
5985                    Ok(p)
5986                }
5987                None => {
5988                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
5989                    es.htod_i32(&pos_vec)
5990                }
5991            }
5992        };
5993
5994        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
5995        let slot = {
5996            let _st0 = rt.enter(0);
5997            let e0 = rt.engine(0, e);
5998            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "start", None)?;
5999            let stage0_started = std::time::Instant::now();
6000            let pos_d = stage_pos(e0)?;
6001            let x = match (stream, embd_dev) {
6002                (Some((vtok, _)), Some((g, qt, rb))) => {
6003                    e0.embed_gather_device_td(g, vtok, t, n_embd, qt, rb)?
6004                }
6005                (None, Some((g, qt, rb))) => e0.embed_gather_device_t(g, tokens, n_embd, qt, rb)?,
6006                _ => e0.htod(&self.embd.gather(n_embd, tokens))?,
6007            };
6008            let x = self.verify_layers(
6009                e0,
6010                x,
6011                fence[0],
6012                fence[1],
6013                &pos_d,
6014                pos0,
6015                t,
6016                cache,
6017                ckpt.as_deref_mut(),
6018                stream,
6019                None,
6020            )?;
6021            if pp_anatomy {
6022                e0.stream().synchronize()?;
6023                stage0_ms = stage0_started.elapsed().as_secs_f64() * 1e3;
6024            }
6025            let tx_started = std::time::Instant::now();
6026            let slot = if pp_pipe.is_some() {
6027                rt.tx_pipelined(0, &x, payload)?
6028            } else {
6029                rt.tx(0, &x, payload)?
6030            };
6031            enqueue_spec_pipe_trace_marker(&e0.stream(), trace.as_ref(), "S0", "end", Some(slot))?;
6032            if pp_anatomy {
6033                e0.stream().synchronize()?;
6034                tx_ms = tx_started.elapsed().as_secs_f64() * 1e3;
6035            }
6036            slot
6037            // x + pos_d drop here: freed stream-ordered on stage-0's stream after use.
6038        };
6039
6040        Ok(VerifyBoundaryTicket {
6041            rt,
6042            caller_stream,
6043            slot,
6044            pos0,
6045            t,
6046            payload,
6047            n_st,
6048            pipelined: pp_pipe.is_some(),
6049            pp_anatomy,
6050            pp_started,
6051            reverse_ms,
6052            stage0_ms,
6053            tx_ms,
6054            trace,
6055        })
6056    }
6057
6058    /// Consume a stage-0 boundary ticket and enqueue the remaining PP stages plus the head.
6059    /// On PP-2 this is exactly stage 1; PP-N keeps its pre-existing middle-stage walk here.
6060    #[allow(clippy::too_many_arguments)]
6061    fn verify_stage1_finish(
6062        &self,
6063        e: &Engine,
6064        ticket: VerifyBoundaryTicket,
6065        cache: &mut Cache,
6066        mut ckpt: Option<&mut VerifyCkpt>,
6067        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
6068        fence: &[usize],
6069        publish_to_caller: bool,
6070    ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
6071        let VerifyBoundaryTicket {
6072            rt,
6073            caller_stream,
6074            slot,
6075            pos0,
6076            t,
6077            payload,
6078            n_st,
6079            pipelined,
6080            pp_anatomy,
6081            pp_started,
6082            reverse_ms,
6083            stage0_ms,
6084            tx_ms,
6085            trace,
6086        } = ticket;
6087        let n_embd = self.cfg.n_embd as usize;
6088        let eps = self.cfg.rms_eps;
6089        let mut slot = slot;
6090        let (mut rx_ms, mut stage1_ms) = (0.0f64, 0.0f64);
6091        let stage_pos = |es: &Engine| -> Result<CudaSlice<i32>, Box<dyn std::error::Error>> {
6092            match stream {
6093                Some((_, ctr)) => {
6094                    let mut p = es.alloc_uninit::<i32>(t)?;
6095                    es.pos_iota(ctr, &mut p, t)?;
6096                    Ok(p)
6097                }
6098                None => {
6099                    let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
6100                    es.htod_i32(&pos_vec)
6101                }
6102            }
6103        };
6104
6105        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
6106        for s in 1..n_st - 1 {
6107            let _st = rt.enter(s);
6108            let es = rt.engine(s, e);
6109            let pos_d = stage_pos(es)?;
6110            let x = rt.rx(s - 1, slot, payload)?;
6111            let x = self.verify_layers(
6112                es,
6113                x,
6114                fence[s],
6115                fence[s + 1],
6116                &pos_d,
6117                pos0,
6118                t,
6119                cache,
6120                ckpt.as_deref_mut(),
6121                stream,
6122                None,
6123            )?;
6124            slot = if pipelined {
6125                rt.tx_pipelined(s, &x, payload)?
6126            } else {
6127                rt.tx(s, &x, payload)?
6128            };
6129        }
6130
6131        // ---- LAST STAGE: RX + final range + output_norm + lm head ----
6132        let _stl = rt.enter(n_st - 1);
6133        let el = rt.engine(n_st - 1, e);
6134        let pos_d = stage_pos(el)?;
6135        let rx_started = std::time::Instant::now();
6136        let x = rt.rx(n_st - 2, slot, payload)?;
6137        if pp_anatomy {
6138            el.stream().synchronize()?;
6139            rx_ms = rx_started.elapsed().as_secs_f64() * 1e3;
6140        }
6141        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "start", Some(slot))?;
6142        let stage1_started = std::time::Instant::now();
6143        let x = self.verify_layers(
6144            el,
6145            x,
6146            fence[n_st - 1],
6147            fence[n_st],
6148            &pos_d,
6149            pos0,
6150            t,
6151            cache,
6152            ckpt.as_deref_mut(),
6153            stream,
6154            None,
6155        )?;
6156
6157        let mut hn = vbuf(el, payload)?;
6158        let logits = if self.sliding_gated_moe_batch_program() {
6159            // The PP Step35 serving path uses rms_norm + matmul for B=1 as well as B>1.
6160            // Verify must not switch numeric class merely because the same session speculates.
6161            el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6162            el.matmul(&self.output, &hn, t)?
6163        } else {
6164            el.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
6165            el.matmul_decode_exact(&self.output, &hn, t)?
6166        };
6167        enqueue_spec_pipe_trace_marker(&el.stream(), trace.as_ref(), "S1", "end", Some(slot))?;
6168        if pp_anatomy {
6169            el.stream().synchronize()?;
6170            stage1_ms = stage1_started.elapsed().as_secs_f64() * 1e3;
6171        }
6172        // EXIT PUBLICATION: both returned buffers are still being produced on the last stage's
6173        // stream. Order the caller's stream behind that work before the buffers escape this
6174        // scope (the 2026-08-06 same-device ppspec find: without it the caller's primary-stream
6175        // consumer read unwritten logits — nondeterministic, one-device-only, and it poisoned
6176        // the following arm's KV in the same process).
6177        if publish_to_caller {
6178            rt.publish_to(n_st - 1, &caller_stream)?;
6179        }
6180        if pp_anatomy {
6181            if publish_to_caller {
6182                caller_stream.synchronize()?;
6183            }
6184            eprintln!(
6185                "[spec-pp-anatomy] t={t} reverse={reverse_ms:.3}ms stage0={stage0_ms:.3}ms \
6186                 tx={tx_ms:.3}ms rx={rx_ms:.3}ms stage1-head={stage1_ms:.3}ms total={:.3}ms",
6187                pp_started.elapsed().as_secs_f64() * 1e3,
6188            );
6189        }
6190        // stream: the device pos counter owns position; host mirror reconciles at drain.
6191        if stream.is_none() {
6192            cache.pos += t;
6193        }
6194        Ok((logits, if spec_hpost() { hn } else { x }))
6195    }
6196
6197    /// Step3.5/Step3.7 verify trunk in the serving batched numeric class.
6198    ///
6199    /// `step35_decode_batch_layers` is now authoritative at every live serving width, including
6200    /// B=1 (lane/cx-b1fix). The older verify walk deliberately mirrored the eager T=1 class:
6201    /// it replayed `step35_decode_attn` per row and used the eager/decode-exact FFN dispatch.
6202    /// Those classes are individually stable, but a near-tie prompt can choose different greedy
6203    /// bytes when a request moves from batched plain serving into speculative verify. Run the
6204    /// same authoritative B=1 stage subgraph for each verify row here. Rows still advance
6205    /// layer-by-layer, so every layer sees the preceding verify rows in its attention cache while
6206    /// every norm/projection/FFN uses exactly the live serving dispatch.
6207    #[allow(clippy::too_many_arguments)]
6208    /// PRIME-BY-T-ROWS (MEMRA_PRIME_TROWS=1): prefill the prompt through the same-session
6209    /// t-row walk in 32-row chunks — every row runs the t=1 decode program bit-for-bit
6210    /// (the TOKENWISE-prime ORACLE class), so this door is exact against the exactness
6211    /// reference while replacing the host-canonical per-token prime. Requires the walk
6212    /// doors (MEMRA_SPEC_VERIFY_EAGER/TCOL); returns the prime contract trio.
6213    #[allow(clippy::type_complexity)]
6214    pub(crate) fn step35_prime_trows(
6215        &self,
6216        e: &Engine,
6217        tokens: &[u32],
6218        cache: &mut Cache,
6219    ) -> Result<Option<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn std::error::Error>>
6220    {
6221        let dbg = std::env::var("MEMRA_SPEC_FA2_DEBUG").as_deref() == Ok("1");
6222        if !prime_trows_on() {
6223            return Ok(None);
6224        }
6225        if !self.uses_sliding_gated_moe_program()
6226            || cache.pos != 0
6227            || cache.dflash_taps.is_some()
6228            || !spec_verify_eager_on()
6229            || !spec_verify_tcol_on()
6230        {
6231            if dbg {
6232                eprintln!(
6233                    "[prime-trows] refuse: program={} pos={} taps={} eager={:?} tcol={:?}",
6234                    self.uses_sliding_gated_moe_program(),
6235                    cache.pos,
6236                    cache.dflash_taps.is_some(),
6237                    std::env::var("MEMRA_SPEC_VERIFY_EAGER").ok(),
6238                    std::env::var("MEMRA_SPEC_VERIFY_TCOL").ok()
6239                );
6240            }
6241            return Ok(None);
6242        }
6243        let n_embd = self.cfg.n_embd as usize;
6244        let n_layers = self.layers.len();
6245        let t_total = tokens.len();
6246        let Some(embd_gpu) = self.embd_gpu_try(e) else {
6247            if dbg {
6248                eprintln!("[prime-trows] refuse: no device embed table");
6249            }
6250            return Ok(None);
6251        };
6252        let embd_qtype = match self.embd.ggml_type {
6253            memra_gguf::GgmlType::BF16 => crate::QT_BF16,
6254            memra_gguf::GgmlType::Q8_0 => crate::QT_Q8_0,
6255            other => {
6256                if dbg {
6257                    eprintln!("[prime-trows] refuse: embed dtype {other:?}");
6258                }
6259                return Ok(None);
6260            }
6261        };
6262        let embd_row_bytes = self.embd.raw.len() / self.cfg.n_vocab as usize;
6263        // Chunk plan: 32-row chunks; a 1-token tail folds into the previous chunk
6264        // (the walk floor is t >= 2).
6265        let mut bounds = Vec::new();
6266        let mut start = 0usize;
6267        while start < t_total {
6268            let mut end = (start + 32).min(t_total);
6269            if t_total - end == 1 {
6270                end -= 1;
6271            }
6272            bounds.push((start, end));
6273            start = end;
6274        }
6275        if bounds.iter().any(|(a, b)| b - a < 2) {
6276            return Ok(None); // degenerate short prompt keeps the ordinary prime
6277        }
6278        let mut hiddens = e.uninit(t_total * n_embd)?;
6279        let mut last: Option<CudaSlice<f32>> = None;
6280        for &(a, b) in &bounds {
6281            let tc = b - a;
6282            let tok_d = e.stream().clone_htod(&tokens[a..b])?;
6283            let x =
6284                e.embed_gather_device_td(embd_gpu, &tok_d, tc, n_embd, embd_qtype, embd_row_bytes)?;
6285            let out = self.step35_verify_batch_layers(e, x, 0, n_layers, a, tc, cache)?;
6286            e.copy_into(&mut hiddens, a * n_embd, &out, tc * n_embd)?;
6287            if b == t_total {
6288                let mut h = e.uninit(n_embd)?;
6289                e.dtod_copy_view(&out.slice((tc - 1) * n_embd..tc * n_embd), &mut h)?;
6290                last = Some(h);
6291            }
6292        }
6293        let h_seed = last.expect("last chunk produced the seed row");
6294        let mut hn = e.uninit(n_embd)?;
6295        e.rms_norm_decode(
6296            &h_seed,
6297            self.output_norm.float_data(),
6298            &mut hn,
6299            n_embd,
6300            1,
6301            self.cfg.rms_eps,
6302        )?;
6303        let logits_d = e.matmul_decode_exact(&self.output, &hn, 1)?;
6304        let logits = e.dtoh(&logits_d)?;
6305        cache.pos = t_total;
6306        Ok(Some((logits, h_seed, hiddens)))
6307    }
6308
6309    fn step35_verify_batch_layers(
6310        &self,
6311        e: &Engine,
6312        mut x: CudaSlice<f32>,
6313        lo: usize,
6314        hi: usize,
6315        pos0: usize,
6316        t: usize,
6317        cache: &mut Cache,
6318    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6319        let n_embd = self.cfg.n_embd as usize;
6320        if !self.uses_sliding_gated_moe_program() {
6321            return Err(
6322                "serving-class verify requires sliding-gated-MoE canonical operations".into(),
6323            );
6324        }
6325        // SERVING-CLASS VERIFY (MEMRA_SPEC_VERIFY_EAGER=1, step37 MTP bring-up): each verify
6326        // column rides decode_layers_eager — the EXACT t=1 program live serving runs (all TP2
6327        // doors) — row-outer, so row r's appends land before row r+1 attends: bit-equal to
6328        // plain greedy by construction. Only the unsplit full-range walk qualifies; PP splits
6329        // and the tap path keep the batch-layer class.
6330        static VE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6331        let eager_verify =
6332            *VE.get_or_init(spec_verify_eager_on) && lo == 0 && hi == self.layers.len();
6333        if eager_verify {
6334            // T-COLUMN LAYER-OUTER WALK (MEMRA_SPEC_VERIFY_TCOL=1): per layer, one t-grid
6335            // attn norm + ONE weight-amortized QKV(+gate) over all T columns, then each
6336            // column runs the UNMODIFIED t=1 attention program via the col-select door and
6337            // the ordinary residual/FFN body. Values per column are bit-equal to the
6338            // row-outer walk: rms over the materialized residual == the fused add+norm
6339            // (kernel_check identity), the tcol kernel's per-column FP order == the t=1
6340            // kernel, and every downstream op IS the t=1 program.
6341            static TCOL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6342            let tcol = *TCOL.get_or_init(spec_verify_tcol_on);
6343            // T > 32 (prefill-class): run the SAME walk in 32-row chunks — each chunk's
6344            // rows are the t=1 program bit-for-bit and the rope pass advances the cache,
6345            // so a chunked call is value-identical to the row-outer loop it replaces.
6346            static TROWS_PREFILL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6347            // MEMRA_STEP_GEMM_PRIME outranks the walk: with the grouped GEMM prime armed, the
6348            // t-row walk defers so the batch path (GEMM trunk + grouped MoE) takes the prompt —
6349            // flag precedence between two existing doors, not a new flag. Without this, both
6350            // doors ON meant the walk still won and the GEMM prime needed PRIME_TROWS=0 by hand.
6351            let trows_prefill =
6352                *TROWS_PREFILL.get_or_init(|| prime_trows_on() && !crate::step_gemm_prime_on());
6353            // MEMRA_PRIME_TROWS_T=<w>: chunk width, default 8 = the REAL cap of this walk.
6354            // The workspace slabs go to 32 rows, but `matvec_bf16_qkvg_tcol_into` refuses
6355            // t > 8 (compile-time-T twins exist for 2/4/8 only; the runtime-t kernel spills
6356            // its accumulators to local memory), so a wider chunk fails the request with
6357            // "matvec_bf16_qkvg_tcol geometry" — which is exactly how the first server-path
6358            // TROWS arm died. Measured at 193 tokens: w=8 2.459 s, w=4 2.574 s.
6359            static TROWS_W: std::sync::OnceLock<Result<usize, String>> = std::sync::OnceLock::new();
6360            let trows_w = match TROWS_W.get_or_init(|| {
6361                let value = std::env::var("MEMRA_PRIME_TROWS_T").ok();
6362                parse_prime_trows_width(value.as_deref())
6363            }) {
6364                Ok(width) => *width,
6365                Err(err) => return Err(err.clone().into()),
6366            };
6367            if tcol && trows_prefill && t > trows_w {
6368                // One-time engagement receipt: without it a prefill gate cannot tell a
6369                // chunked walk from the row-outer fallback it is supposed to replace
6370                // (the first PRIME_TROWS gate passed vacuously on exactly that).
6371                static SEEN: std::sync::atomic::AtomicBool =
6372                    std::sync::atomic::AtomicBool::new(false);
6373                if !SEEN.swap(true, std::sync::atomic::Ordering::Relaxed) {
6374                    eprintln!(
6375                        "[prime-trows] ENGAGED t={t} width={trows_w} chunks={} layers={}..{}",
6376                        t.div_ceil(trows_w),
6377                        lo,
6378                        hi
6379                    );
6380                }
6381                let mut out = e.uninit(t * n_embd)?;
6382                let mut start = 0usize;
6383                while start < t {
6384                    let mut end = (start + trows_w).min(t);
6385                    if t - end == 1 {
6386                        end -= 1;
6387                    }
6388                    let tc = end - start;
6389                    let mut xc = e.uninit(tc * n_embd)?;
6390                    e.dtod_copy_view(&x.slice(start * n_embd..end * n_embd), &mut xc)?;
6391                    let oc =
6392                        self.step35_verify_batch_layers(e, xc, lo, hi, pos0 + start, tc, cache)?;
6393                    e.copy_into(&mut out, start * n_embd, &oc, tc * n_embd)?;
6394                    start = end;
6395                }
6396                return Ok(out);
6397            }
6398            if tcol && t >= 2 && t <= 32 {
6399                // MEMRA_TCOL_PROF=1: synchronized per-segment wall profile of the walk
6400                // (norm+QKV precompute / per-col attention / per-col residual+FFN). The
6401                // syncs serialize the stream, so the split is for TARGETING amortization
6402                // work only — never a perf claim.
6403                static PROF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6404                let prof =
6405                    *PROF.get_or_init(|| std::env::var("MEMRA_TCOL_PROF").as_deref() == Ok("1"));
6406                let mut prof_ms = [0f64; 3];
6407                let eps = self.cfg.rms_eps;
6408                let mut x_t = x;
6409                let mut h_t = e.uninit(t * n_embd)?;
6410                let mut h_row = e.uninit(n_embd)?; // real row: the non-dcw fallback reads it
6411                // Per-column pos buffers hoisted out of the layer loop (a per-col-per-layer
6412                // pageable htod was an in-stream engine turnaround x t x 45).
6413                let mut pos_rows = Vec::with_capacity(t);
6414                for r in 0..t {
6415                    pos_rows.push(e.htod_i32(&[(pos0 + r) as i32])?);
6416                }
6417                let mut ok = true;
6418                // MEMRA_TCOL_OPROJ=1: defer each column's o_proj — the finish seam
6419                // stashes `gated` instead of joining per column; one b4_tcol per rank +
6420                // one slab join produce every column's `mixed` after the attention pass.
6421                // Bit-exact per column (t=1 b4 program per column; elementwise join).
6422                // MEMRA_TCOL_FFN=1: today this only IMPLIES the o_proj defer above. Its
6423                // named feature, the two-column device-routed FFN sweep, rode the
6424                // slot-major v2 TP banks and was REMOVED with the MEMRA_NVFP4_BANK_V2 door
6425                // (2026-08-29, research/step37-bankv2-removal-20260829): the v2 layout
6426                // changed generated text in serving. The flag itself stays because it is
6427                // family-armed in the step37 serving defaults and killing it here would
6428                // silently drop the o_proj defer from the qualified serving shape.
6429                static FFN2: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6430                let ffn_batch = *FFN2.get_or_init(tcol_ffn_on);
6431                let oproj_batch = crate::tp::tcol_oproj_on() || ffn_batch;
6432                // MEMRA_SPEC_FA2=1 (T=2 only): eligible layers defer BOTH columns' fa —
6433                // the per-column pass norms/ropes/appends and stashes q+gate, then one
6434                // shared-KV fa_decode_dcw2 per rank + the o_proj join produce the
6435                // [2, o_out] mixed slab. The precheck runs before arming (stashing is
6436                // unrecoverable); ineligible/boundary layers run the ordinary program.
6437                let fa2 = crate::tp::spec_fa2_on() && t <= 32;
6438                let mut mixed_row = e.uninit(n_embd)?;
6439                let mut pos_staged = false;
6440                for il in lo..hi {
6441                    let layer = &self.layers[il];
6442                    // BEFORE this layer touches its planes: is the history it is about to
6443                    // attend already poisoned? Global (non-ring) layers only, which are the
6444                    // ones the level-2 bitmap implicates.
6445                    if kv_plane_scan_on() && self.step35_geom(il).window.is_none() {
6446                        if let Some(distributed) = cache.tp_kv[il].as_ref() {
6447                            scan_kv_plane(e, distributed, il, pos0)?;
6448                        }
6449                    }
6450                    let fa2_layer = fa2 && self.step35_fa_rows_precheck(cache, il, pos0, t)?;
6451                    let mut seg = std::time::Instant::now();
6452                    e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
6453                    if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
6454                        ok = false;
6455                        break;
6456                    }
6457                    // FULL t-row attention pass (rope/append + fa + combine + o_proj in
6458                    // 3 launches/rank): same-session rows, slot = len-base+r, one len
6459                    // advance by t. Host cache bookkeeping mirrors the per-column tail.
6460                    if fa2_layer {
6461                        if let Some(mixed_t) =
6462                            self.step35_verify_rope_fa_pass(e, il, cache, pos0, t, !pos_staged)?
6463                        {
6464                            pos_staged = true;
6465                            {
6466                                let tp_kv = cache.tp_kv[il]
6467                                    .as_mut()
6468                                    .expect("precheck verified the distributed cache");
6469                                let transaction = tp_kv.begin_transaction()?;
6470                                let crate::hybrid::Mixer::Full(fa) = &layer.mixer else {
6471                                    return Err("verify rope pass expects full attention".into());
6472                                };
6473                                let tp = fa
6474                                    .step_tp_qkv
6475                                    .as_ref()
6476                                    .ok_or("verify rope pass lost its TP state")?;
6477                                let empty: [CudaSlice<f32>; 0] = [];
6478                                tp.runtime.append_tp_kv_transaction_inner(
6479                                    tp_kv,
6480                                    transaction,
6481                                    &empty,
6482                                    &empty,
6483                                    t,
6484                                    true,
6485                                )?;
6486                                tp.runtime.commit_tp_kv_transaction_external(
6487                                    tp_kv,
6488                                    transaction,
6489                                    t,
6490                                )?;
6491                                if let Some(local) = cache.kv[il].as_mut() {
6492                                    local.len = pos0 + t;
6493                                    if !crate::tp::len_mirror_lazy_on() {
6494                                        e.set_i32_one(&mut local.len_d, local.len as i32)?;
6495                                    }
6496                                }
6497                            }
6498                            if prof {
6499                                e.stream().synchronize()?;
6500                                prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6501                                seg = std::time::Instant::now();
6502                            }
6503                            let o_out = mixed_t.len() / t;
6504                            let mut next = e.uninit(t * n_embd)?;
6505                            {
6506                                for r in 0..t {
6507                                    e.dtod_copy_view(
6508                                        &mixed_t.slice(r * o_out..(r + 1) * o_out),
6509                                        &mut mixed_row,
6510                                    )?;
6511                                    let mut x_row = e.uninit(n_embd)?;
6512                                    e.dtod_copy_view(
6513                                        &x_t.slice(r * n_embd..(r + 1) * n_embd),
6514                                        &mut x_row,
6515                                    )?;
6516                                    let (x1, ffn_out) = self.residual_norm_ffn(
6517                                        e, layer, &x_row, &mixed_row, n_embd, il, eps,
6518                                    )?;
6519                                    let mut x2 = e.uninit(n_embd)?;
6520                                    e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6521                                    e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
6522                                }
6523                            }
6524                            if prof {
6525                                e.stream().synchronize()?;
6526                                prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6527                            }
6528                            x_t = next;
6529                            if spec_nan_scan() {
6530                                // The scan MUST sit on this arm too. It used to live only on
6531                                // the non-fused tail, so a fused layer's poison was first
6532                                // reported by the next non-fused layer.
6533                                verify_arm_receipt(
6534                                    "fused",
6535                                    il,
6536                                    pos0,
6537                                    t,
6538                                    cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6539                                );
6540                                nan_scan_rows(
6541                                    e,
6542                                    &x_t,
6543                                    t,
6544                                    n_embd,
6545                                    &format!("tcol layer {il} pos0={pos0} arm=fused"),
6546                                )?;
6547                            }
6548                            continue;
6549                        }
6550                    }
6551                    if prof {
6552                        e.stream().synchronize()?;
6553                        prof_ms[0] += seg.elapsed().as_secs_f64() * 1e3;
6554                        seg = std::time::Instant::now();
6555                    }
6556                    let mut next = e.uninit(t * n_embd)?;
6557                    // Columns whose o_proj was deferred (their FFN runs after the join).
6558                    // A NON-deferred column's FFN must run INSIDE the column loop: the
6559                    // oproj-tail handoff is a single cell that the same column's
6560                    // residual_norm_ffn consumes before the next column's finish.
6561                    let mut deferred: Vec<usize> = Vec::new();
6562                    let mut fa2_deferred: Vec<usize> = Vec::new();
6563                    let mut ffn_col =
6564                        |r: usize,
6565                         mixed: &CudaSlice<f32>,
6566                         next: &mut CudaSlice<f32>|
6567                         -> Result<(), Box<dyn std::error::Error>> {
6568                            let mut x_row = e.uninit(n_embd)?;
6569                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
6570                            let (x1, ffn_out) =
6571                                self.residual_norm_ffn(e, layer, &x_row, mixed, n_embd, il, eps)?;
6572                            if spec_nan_scan_level() >= 2 {
6573                                nan_scan_rows(
6574                                    e,
6575                                    &ffn_out,
6576                                    1,
6577                                    n_embd,
6578                                    &format!("tcol layer {il} col {r} per-column FFN out"),
6579                                )?;
6580                            }
6581                            let mut x2 = e.uninit(n_embd)?;
6582                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
6583                            e.dtod_copy_into(&x2, next, r * n_embd)?;
6584                            Ok(())
6585                        };
6586                    for r in 0..t {
6587                        e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
6588                        let row_pos = &pos_rows[r];
6589                        crate::tp::set_verify_tcol(Some(r));
6590                        if fa2_layer {
6591                            crate::tp::set_spec_fa2_defer(Some(r));
6592                        } else if oproj_batch {
6593                            crate::tp::set_tcol_oproj_defer(Some(r));
6594                        }
6595                        let mixed = match &layer.mixer {
6596                            crate::hybrid::Mixer::Full(fa) => {
6597                                self.full_attn_decode(e, fa, &h_row, row_pos, pos0 + r, cache, il)
6598                            }
6599                            _ => Err("step35 verify expects full attention".into()),
6600                        };
6601                        crate::tp::set_verify_tcol(None);
6602                        crate::tp::set_spec_fa2_defer(None);
6603                        crate::tp::set_tcol_oproj_defer(None);
6604                        let mixed = mixed?;
6605                        if fa2_layer && crate::tp::take_spec_fa2_stashed() {
6606                            fa2_deferred.push(r);
6607                        } else if oproj_batch && crate::tp::take_tcol_oproj_stashed() {
6608                            deferred.push(r);
6609                        } else {
6610                            if spec_nan_scan_level() >= 2 {
6611                                let cols = mixed.len();
6612                                nan_scan_rows(
6613                                    e,
6614                                    &mixed,
6615                                    1,
6616                                    cols,
6617                                    &format!("tcol layer {il} col {r} per-column ATTN out"),
6618                                )?;
6619                            }
6620                            ffn_col(r, &mixed, &mut next)?;
6621                        }
6622                    }
6623                    if !fa2_deferred.is_empty() && fa2_deferred.len() != t {
6624                        // The precheck guarantees both columns stash or neither; a strict
6625                        // subset means a column's output was never produced anywhere.
6626                        return Err("spec fa2 stash engaged for a subset of columns".into());
6627                    }
6628                    if prof {
6629                        e.stream().synchronize()?;
6630                        prof_ms[1] += seg.elapsed().as_secs_f64() * 1e3;
6631                        seg = std::time::Instant::now();
6632                    }
6633                    if !fa2_deferred.is_empty() {
6634                        deferred = fa2_deferred;
6635                    }
6636                    if !deferred.is_empty() {
6637                        let mixed_t = if fa2_layer {
6638                            self.step35_verify_fa_rows_join(e, il, cache, pos0, t)?
6639                        } else {
6640                            self.step35_verify_oproj_tcol(e, il, t)?
6641                        };
6642                        let o_out = mixed_t.len() / t;
6643                        if spec_nan_scan_level() >= 2 {
6644                            nan_scan_rows(
6645                                e,
6646                                &mixed_t,
6647                                t,
6648                                o_out,
6649                                &format!("tcol layer {il} JOINED attn over deferred cols"),
6650                            )?;
6651                        }
6652                        // Batched t=2 residual+MoE: one t-grid add_rms_norm (per-row
6653                        // program == t=1; bit-identical to the oproj-tail join per the
6654                        // M2 verbatim-program contract) feeding the two-column routed
6655                        // sweep. Ineligible layers (dense FFN, non-nvfp4) fall through
6656                        // to the per-column body.
6657                        {
6658                            for &r in &deferred {
6659                                e.dtod_copy_view(
6660                                    &mixed_t.slice(r * o_out..(r + 1) * o_out),
6661                                    &mut mixed_row,
6662                                )?;
6663                                ffn_col(r, &mixed_row, &mut next)?;
6664                            }
6665                        }
6666                    }
6667                    if prof {
6668                        e.stream().synchronize()?;
6669                        prof_ms[2] += seg.elapsed().as_secs_f64() * 1e3;
6670                    }
6671                    drop(ffn_col);
6672                    x_t = next;
6673                    if spec_nan_scan() {
6674                        verify_arm_receipt(
6675                            if fa2_layer { "join" } else { "percol" },
6676                            il,
6677                            pos0,
6678                            t,
6679                            cache.tp_kv[il].as_ref().map(|d| d.staged_len()),
6680                        );
6681                        nan_scan_rows(
6682                            e,
6683                            &x_t,
6684                            t,
6685                            n_embd,
6686                            &format!(
6687                                "tcol layer {il} pos0={pos0} arm={}",
6688                                if fa2_layer { "join" } else { "percol" }
6689                            ),
6690                        )?;
6691                    }
6692                }
6693                if prof {
6694                    eprintln!(
6695                        "[tcol-prof] t={t} norm+qkv={:.3}ms attn={:.3}ms ffn={:.3}ms",
6696                        prof_ms[0], prof_ms[1], prof_ms[2]
6697                    );
6698                }
6699                if ok {
6700                    return Ok(x_t);
6701                }
6702                // fall through to the row-outer walk on ineligible layers
6703                x = x_t;
6704            }
6705            let mut next = e.uninit(t * n_embd)?;
6706            let scan = spec_nan_scan();
6707            for r in 0..t {
6708                let mut row = e.uninit(n_embd)?;
6709                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6710                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6711                let out = if scan {
6712                    // Diagnostic arm: the same range walked one layer at a time so the first
6713                    // poisoned layer names itself. `decode_layers_eager(lo, hi)` is range-scoped
6714                    // and executes its trailing residual add, so a per-layer chain is the same
6715                    // program with the cross-layer add+norm fusion unrolled.
6716                    nan_scan_rows(
6717                        e,
6718                        &row,
6719                        1,
6720                        n_embd,
6721                        &format!("embed row r={r} pos={}", pos0 + r),
6722                    )?;
6723                    let mut acc = row;
6724                    for il in lo..hi {
6725                        acc = self.decode_layers_eager(
6726                            e,
6727                            acc,
6728                            il,
6729                            il + 1,
6730                            &row_pos,
6731                            pos0 + r,
6732                            cache,
6733                        )?;
6734                        nan_scan_rows(
6735                            e,
6736                            &acc,
6737                            1,
6738                            n_embd,
6739                            &format!("row-outer layer {il} r={r} pos={}", pos0 + r),
6740                        )?;
6741                    }
6742                    acc
6743                } else {
6744                    self.decode_layers_eager(e, row, lo, hi, &row_pos, pos0 + r, cache)?
6745                };
6746                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6747            }
6748            // dflash taps are NOT produced on this arm (they need per-layer hiddens the
6749            // row-outer walk does not materialize); the door is a step37 MTP bring-up
6750            // surface where taps are unused.
6751            return Ok(next);
6752        }
6753        let mut ph_last = std::time::Instant::now();
6754        for il in lo..hi {
6755            let mut next = e.uninit(t * n_embd)?;
6756            for r in 0..t {
6757                let mut row = e.uninit(n_embd)?;
6758                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
6759                // The caller owns this verify's position. During controller overlap, cache.pos
6760                // still describes generation N while this stage-0 walk belongs to N+1.
6761                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
6762                let mut one = [&mut *cache];
6763                let out = self.step35_decode_batch_layers(
6764                    e,
6765                    row,
6766                    &mut one,
6767                    &[(pos0 + r) as i32],
6768                    &row_pos,
6769                    il,
6770                    il + 1,
6771                    &mut ph_last,
6772                )?;
6773                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
6774            }
6775            self.dflash_tap(e, cache, il, &next, t)?;
6776            x = next;
6777            if spec_nan_scan() {
6778                nan_scan_rows(e, &x, t, n_embd, &format!("batch-layer {il} pos0={pos0}"))?;
6779            }
6780        }
6781        Ok(x)
6782    }
6783
6784    /// DSpark drafter verify (lane/dspark-q38-recover): one t-row forward through the
6785    /// SERVING-CLASS verify funnel (`decode_step_t_core_stream` — the same numeric class
6786    /// MTP verify rides, GDN state advanced in place), returning per-row argmax tokens.
6787    /// Advances `cache.pos += t`; the caller owns snapshot/rollback (block acceptance is
6788    /// prefix-keep, not all-or-nothing).
6789    pub(crate) fn dspark_verify_t_am(
6790        &self,
6791        e: &Engine,
6792        tokens: &[u32],
6793        pos0: usize,
6794        cache: &mut Cache,
6795    ) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
6796        let (logits, _hn) = self.decode_step_t_core_stream(
6797            e, tokens, pos0, cache, None, None, None, None, None, None,
6798        )?;
6799        let t = tokens.len();
6800        let v = self.output.out_features();
6801        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6802        for r in 0..t {
6803            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6804        }
6805        Ok(e.dtoh_u32(&am_d)?)
6806    }
6807
6808    /// DSpark verify returning the RAW verify logits [t, n_vocab] (device-resident) instead
6809    /// of per-row argmaxes — the sampled-admission arm's input (rejection-sampling accept
6810    /// gathers filtered p from these columns; lane/dspark-sampled-admission-20260820). Same
6811    /// forward as `dspark_verify_t_am`; the greedy arm keeps its argmax wrapper untouched.
6812    pub(crate) fn dspark_verify_t_logits(
6813        &self,
6814        e: &Engine,
6815        tokens: &[u32],
6816        pos0: usize,
6817        cache: &mut Cache,
6818    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
6819        let (logits, _hn) = self.decode_step_t_core_stream(
6820            e, tokens, pos0, cache, None, None, None, None, None, None,
6821        )?;
6822        Ok(logits)
6823    }
6824
6825    /// DSpark verify with the MTP column-stash armed: identical forward to
6826    /// `dspark_verify_t_am`, but fills a `VerifyCkpt` so a partial accept can restore
6827    /// column state directly (`dspark_commit_prefix`) instead of snapshot-replay.
6828    /// The ckpt type is opaque outside spec.rs (newtype) — dflash.rs threads it through.
6829    pub(crate) fn dspark_verify_t_am_ckpt(
6830        &self,
6831        e: &Engine,
6832        tokens: &[u32],
6833        pos0: usize,
6834        cache: &mut Cache,
6835    ) -> Result<(Vec<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6836        let mut ck = VerifyCkpt::new(self.layers.len());
6837        let (logits, _hn) = self.decode_step_t_core_stream(
6838            e,
6839            tokens,
6840            pos0,
6841            cache,
6842            None,
6843            Some(&mut ck),
6844            None,
6845            None,
6846            None,
6847            None,
6848        )?;
6849        let t = tokens.len();
6850        let v = self.output.out_features();
6851        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6852        for r in 0..t {
6853            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6854        }
6855        Ok((e.dtoh_u32(&am_d)?, DsparkVerifyCkpt(ck)))
6856    }
6857
6858    /// Engine-bundle slice 2: `dspark_verify_t_am_ckpt` with DEVICE tokens and NO readback.
6859    /// The verify tokens are the round's `chain_d` (cand layout: [anchor, drafts...]); the
6860    /// embed gathers its first `t` entries on-device (`embed_gather_u32_t` — bit-identical
6861    /// rows to the host arm), so the host never blocks on the draft chain before dispatching
6862    /// verify. Returns the device per-row argmax buffer; the caller merges its readback with
6863    /// the chain's into ONE sync. Forward, ckpt fill and argmax walk are `_ckpt` verbatim.
6864    pub(crate) fn dspark_verify_t_am_ckpt_dev(
6865        &self,
6866        e: &Engine,
6867        vtok: &CudaSlice<u32>,
6868        t: usize,
6869        pos0: usize,
6870        cache: &mut Cache,
6871        embd_dev: (&CudaSlice<u8>, i32, usize),
6872        graphs: Option<&mut DsparkVerifyGraphs>,
6873    ) -> Result<(CudaSlice<u32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6874        debug_assert!(
6875            vtok.len() >= t,
6876            "verify window exceeds the device token buffer"
6877        );
6878        // The slab flag is a per-round statement: clear it here so a verify that never
6879        // reaches the graphs door (rowwise env, a non-tparallel arm) cannot leave a
6880        // stale `true` steering the commit at slabs the round never wrote.
6881        let mut graphs = graphs;
6882        if let Some(g) = graphs.as_deref_mut() {
6883            g.round_slab = false;
6884        }
6885        let mut ck = VerifyCkpt::new(self.layers.len());
6886        // Dummy host tokens size the funnel; the embed reads `vtok` (the round-stream
6887        // arm's established pattern — spec.rs stream-mode verify does the same).
6888        let dummy = vec![0u32; t];
6889        let (logits, _hn) = self.decode_step_t_core_stream(
6890            e,
6891            &dummy,
6892            pos0,
6893            cache,
6894            Some(embd_dev),
6895            Some(&mut ck),
6896            None,
6897            None,
6898            Some(vtok),
6899            graphs,
6900        )?;
6901        let v = self.output.out_features();
6902        let mut am_d = e.stream().alloc_zeros::<u32>(t)?;
6903        for r in 0..t {
6904            e.argmax_token_device_col(&logits, r, v, &mut am_d, r)?;
6905        }
6906        Ok((am_d, DsparkVerifyCkpt(ck)))
6907    }
6908
6909    /// Ckpt-armed twin of [`Self::dspark_verify_t_logits`] (sampled-admission arm).
6910    pub(crate) fn dspark_verify_t_logits_ckpt(
6911        &self,
6912        e: &Engine,
6913        tokens: &[u32],
6914        pos0: usize,
6915        cache: &mut Cache,
6916    ) -> Result<(CudaSlice<f32>, DsparkVerifyCkpt), Box<dyn std::error::Error>> {
6917        let mut ck = VerifyCkpt::new(self.layers.len());
6918        let (logits, _hn) = self.decode_step_t_core_stream(
6919            e,
6920            tokens,
6921            pos0,
6922            cache,
6923            None,
6924            Some(&mut ck),
6925            None,
6926            None,
6927            None,
6928            None,
6929        )?;
6930        Ok((logits, DsparkVerifyCkpt(ck)))
6931    }
6932
6933    /// Restore the round to `keep` accepted columns from the verify stash: KV lens and
6934    /// pos from the pre-verify snapshot + keep, GDN conv/ssm from the stashed column
6935    /// state — no replay forward. The exact `commit_verified_prefix` the MTP path ships.
6936    pub(crate) fn dspark_commit_prefix(
6937        &self,
6938        e: &Engine,
6939        cache: &mut Cache,
6940        snap: &crate::cache::CacheSnapshot,
6941        ckpt: &DsparkVerifyCkpt,
6942        keep: usize,
6943    ) -> Result<(), Box<dyn std::error::Error>> {
6944        self.commit_verified_prefix(e, cache, snap, &ckpt.0, keep, false, None)
6945    }
6946
6947    /// Slice-3 commit twin: restore to `keep` accepted columns when the round's linear
6948    /// column stash lives in the graphs ctx's persistent slabs (`DsparkVerifyGraphs`) —
6949    /// the cols arm's exact semantics (KV lens + pos from the snapshot, GDN conv/ssm
6950    /// from the stash of column keep-1), slab-addressed and batched into two copy
6951    /// launches. `MEMRA_STATE_COPY_BATCH=0` falls back to per-layer view copies.
6952    pub(crate) fn dspark_commit_prefix_slab(
6953        &self,
6954        e: &Engine,
6955        cache: &mut Cache,
6956        snap: &crate::cache::CacheSnapshot,
6957        ctx: &DsparkVerifyGraphs,
6958        keep: usize,
6959    ) -> Result<(), Box<dyn std::error::Error>> {
6960        use cudarc::driver::DevicePtr;
6961        debug_assert!(keep >= 1, "keep==0 rounds take the legacy rollback");
6962        let mut conv_src: Vec<u64> = Vec::new();
6963        let mut ssm_src: Vec<u64> = Vec::new();
6964        let mut conv_dst: Vec<u64> = Vec::new();
6965        let mut ssm_dst: Vec<u64> = Vec::new();
6966        for il in 0..self.layers.len() {
6967            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
6968                kvl.len = saved + keep;
6969                e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
6970            }
6971            if let Some(rl) = cache.recur[il].as_ref() {
6972                let (pc, ps, _cw, _sw) = ctx
6973                    .slab_row(e, il, keep - 1)
6974                    .ok_or("slab commit: linear layer missing from the graphs ctx")?;
6975                conv_src.push(pc);
6976                ssm_src.push(ps);
6977                let st = &e.gpu.stream();
6978                let (dc, _g0) = rl.conv_state.device_ptr(st);
6979                let (ds, _g1) = rl.ssm_state.device_ptr(st);
6980                conv_dst.push(dc as u64);
6981                ssm_dst.push(ds as u64);
6982            }
6983        }
6984        let n = conv_src.len();
6985        if n > 0 {
6986            if state_copy_batch_on() {
6987                let mut tt = vec![0u64; 2 * n];
6988                tt[..n].copy_from_slice(&conv_src);
6989                tt[n..].copy_from_slice(&conv_dst);
6990                let ct = e.htod_u64(&tt)?;
6991                tt[..n].copy_from_slice(&ssm_src);
6992                tt[n..].copy_from_slice(&ssm_dst);
6993                let st = e.htod_u64(&tt)?;
6994                e.copy_batch_uniform_f32(&ct, n, ctx.conv_words)?;
6995                e.copy_batch_uniform_f32(&st, n, ctx.ssm_words)?;
6996            } else {
6997                let (cw, sw) = (ctx.conv_words, ctx.ssm_words);
6998                let row = keep - 1;
6999                for il in 0..self.layers.len() {
7000                    let Some(rl) = cache.recur[il].as_mut() else {
7001                        continue;
7002                    };
7003                    let k = ctx.lin_pos[&il];
7004                    {
7005                        let sv = e.view(&ctx.stash_conv[k], (row + 1) * cw);
7006                        let win = sv.slice(row * cw..(row + 1) * cw);
7007                        e.copy_view_into(&mut rl.conv_state, 0, &win, cw)?;
7008                    }
7009                    {
7010                        let sv = e.view(&ctx.stash_ssm[k], (row + 1) * sw);
7011                        let win = sv.slice(row * sw..(row + 1) * sw);
7012                        e.copy_view_into(&mut rl.ssm_state, 0, &win, sw)?;
7013                    }
7014                }
7015            }
7016        }
7017        cache.pos = snap.pos + keep;
7018        Ok(())
7019    }
7020
7021    /// Qwen35-family verify trunk in the live serving numeric class.
7022    ///
7023    /// Serving intentionally keeps this architecture in the generic batched program even at
7024    /// B=1. The older verify walk used its own mirrored dispatch and can flip near-tie argmaxes.
7025    ///
7026    /// Two arms, one numeric class:
7027    /// - DENSE GDN (`DenseMlp`, t<=16): `qwen35_verify_tparallel` — the weight ops (norms,
7028    ///   projections, FFN) hoist to m=T through the exact-tier batched kernels whose per-row
7029    ///   program IS the m=1 program (`matmul_pre == fused2 per (tensor,row); _bN mmvq per-row
7030    ///   == m=1` — decode_batch.rs v2 note), while the state ops (conv ring, gdn scan, KV
7031    ///   append, fa decode) stay a per-row loop running the b_n=1 serving kernels with each
7032    ///   row's own t_kv-driven arm pick (the straddle law: every row executes the exact
7033    ///   program its isolated serving step would). One weight read per layer per round
7034    ///   instead of T — this is what makes MTP profitable in the exact class (the per-row
7035    ///   walk measured verify(K+1) ~= (K+1) plain steps: 69 -> 44 tok/s served, 2026-08-15).
7036    /// - MoE / t>16 / `MEMRA_SPEC_VERIFY_ROWWISE=1`: the per-row replay of the authoritative
7037    ///   serving layer body, preserving single-session autoregressive cache order (the
7038    ///   correctness reference; also the rollback seam for the t-parallel arm).
7039    ///
7040    /// Bit-identity of the t-parallel arm vs the rowwise arm is gated by spec-serve-gate
7041    /// (zero differing logits at T=1..4, K arms) + the 8-prompt ON/OFF canary before ship.
7042    #[allow(clippy::too_many_arguments)]
7043    fn qwen35_verify_batch_layers(
7044        &self,
7045        e: &Engine,
7046        x: CudaSlice<f32>,
7047        lo: usize,
7048        hi: usize,
7049        pos0: usize,
7050        t: usize,
7051        cache: &mut Cache,
7052        ckpt: Option<&mut VerifyCkpt>,
7053        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7054        graphs: Option<&mut DsparkVerifyGraphs>,
7055    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7056        // Qwen35Moe admitted 2026-08-20 (lane/draftcost-moe): the t-parallel arm already
7057        // carries the MoE FFN (`moe_ffn_il_zq8` at m=T) and the GDN per-row state loop; the
7058        // arch fence was a qualification gate, not a mechanism gap. Measured disease on the
7059        // 35B-A3B class: rowwise verify ~= 5.6 ms per drafted token (one full trunk step
7060        // each) — the same (K+1)-plain-steps wall the dense admission fixed on 2026-08-15.
7061        // Rollback seam unchanged: MEMRA_SPEC_VERIFY_ROWWISE=1.
7062        let rowwise = std::env::var("MEMRA_SPEC_VERIFY_ROWWISE").as_deref() == Ok("1")
7063            || !self.batched_serving_numeric_class()
7064            || t > 16;
7065        if rowwise {
7066            if stream.is_some() {
7067                // rowwise replays per row with host cache.pos — irreconcilable with a
7068                // device position counter. Burst callers must keep t <= 16 and the
7069                // ROWWISE env unset; refusing beats silently mispositioned rows.
7070                return Err("qwen35 rowwise verify has no ROUND-STREAM arm \
7071                            (t > 16 or MEMRA_SPEC_VERIFY_ROWWISE=1)"
7072                    .into());
7073            }
7074            self.qwen35_verify_rowwise(e, x, lo, hi, pos0, t, cache, ckpt)
7075        } else {
7076            self.qwen35_verify_tparallel(e, x, lo, hi, pos0, t, cache, ckpt, stream, graphs)
7077        }
7078    }
7079
7080    /// The per-row correctness reference: replay each verify row through the authoritative
7081    /// serving layer body (`decode_batch_layers` at b_n=1). T full weight reads per layer.
7082    #[allow(clippy::too_many_arguments)]
7083    fn qwen35_verify_rowwise(
7084        &self,
7085        e: &Engine,
7086        mut x: CudaSlice<f32>,
7087        lo: usize,
7088        hi: usize,
7089        pos0: usize,
7090        t: usize,
7091        cache: &mut Cache,
7092        mut ckpt: Option<&mut VerifyCkpt>,
7093    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7094        let n_embd = self.cfg.n_embd as usize;
7095        let saved_pos = cache.pos;
7096        let mut ph_last = std::time::Instant::now();
7097        for il in lo..hi {
7098            let mut next = e.uninit(t * n_embd)?;
7099            let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7100                if ckpt.is_some() && t >= 2 && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7101                    Some(Vec::with_capacity(t - 1))
7102                } else {
7103                    None
7104                };
7105            for r in 0..t {
7106                cache.pos = pos0 + r;
7107                let mut row = e.uninit(n_embd)?;
7108                e.dtod_copy_view(&x.slice(r * n_embd..(r + 1) * n_embd), &mut row)?;
7109                let row_pos = e.htod_i32(&[(pos0 + r) as i32])?;
7110                let mut one = [&mut *cache];
7111                let ctx = self.batch_layer_ctx(e, &one, il, il + 1)?;
7112                let out = match self.decode_batch_layers(
7113                    e,
7114                    row,
7115                    &mut one,
7116                    &ctx,
7117                    &row_pos,
7118                    &mut ph_last,
7119                ) {
7120                    Ok(out) => out,
7121                    Err(error) => {
7122                        cache.pos = saved_pos;
7123                        return Err(error);
7124                    }
7125                };
7126                e.dtod_copy_into(&out, &mut next, r * n_embd)?;
7127                if r + 1 < t {
7128                    if let Some(states) = col_states.as_mut() {
7129                        let recur = cache.recur[il]
7130                            .as_ref()
7131                            .ok_or("Qwen35-MoE linear verify layer has no recurrent state")?;
7132                        states.push((
7133                            e.clone_dtod(&recur.conv_state)?,
7134                            e.clone_dtod(&recur.ssm_state)?,
7135                        ));
7136                    }
7137                }
7138            }
7139            if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
7140                checkpoint.cols[il] = Some(states);
7141            }
7142            x = next;
7143        }
7144        cache.pos = saved_pos;
7145        Ok(x)
7146    }
7147
7148    /// T-PARALLEL VERIFY IN THE SERVING NUMERIC CLASS (lane/tparallel-verify, 2026-08-15).
7149    ///
7150    /// The weight ops run ONCE per layer at m=T; the state ops run per row through the same
7151    /// b_n=1 serving kernels the rowwise replay uses. Per-row bit-identity rests on the two
7152    /// pins the serving batch tier already carries:
7153    ///   * `matmul_pre` / `_bN` mmvq: per-row program == m=1 program (decode_batch.rs v2 note,
7154    ///     kernel-check pinned) — so a [T, n_embd] projection row equals the row projected
7155    ///     alone;
7156    ///   * row-indexed norms/elementwise (`rms_norm`, `quantize_q8_1`, `add_rms_norm`,
7157    ///     `gated_rmsnorm[_q8_1]`, `silu_mul`, `rope_neox` with per-row positions): the T-row
7158    ///     launch is the per-row program (same pin the generic verify's fused norms rely on).
7159    /// The sequential dependencies keep their exact serving order: the conv ring / gdn scan
7160    /// chain state row -> row through the `_b` kernels at b_n=1 (ping-pong via a 6-entry
7161    /// alternating pointer table, host handles swapped per row so VerifyCkpt clones the
7162    /// canonical state exactly as the rowwise arm does), and each row's KV append + fa decode
7163    /// picks its arm from ITS OWN t_kv (append: format-only; fa: `fa_seqs_eligible` + its own
7164    /// `fa_split_keys` rung at b_n=1) — the straddle law per row, so every row executes the
7165    /// program its isolated B=1 serving step would.
7166    ///
7167    /// Cost: 1 weight read per layer per round + T state micro-launches, vs the rowwise arm's
7168    /// T weight reads. Gated bit-identical vs the rowwise arm by spec-serve-gate + canary.
7169    #[allow(clippy::too_many_arguments)]
7170    fn qwen35_verify_tparallel(
7171        &self,
7172        e: &Engine,
7173        mut x: CudaSlice<f32>,
7174        lo: usize,
7175        hi: usize,
7176        pos0: usize,
7177        t: usize,
7178        cache: &mut Cache,
7179        mut ckpt: Option<&mut VerifyCkpt>,
7180        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
7181        mut graphs: Option<&mut DsparkVerifyGraphs>,
7182    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7183        let seqs_append =
7184            std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0") && !Engine::kv_fp8_on();
7185        let batch_fa_on = std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0");
7186
7187        // Merge guard (v0.98 train, re-affirmed on the v0.100 train over slice 4c): the
7188        // ROUND-STREAM arm (lane/draftcost-moe, device position counter) and the dspark
7189        // verify graphs (engine-bundle slice 3 / trunk slice 4c) have no common caller —
7190        // stream rides the qwen35moe burst, graphs ride the dspark route. If a future
7191        // caller arms both, refuse loudly instead of silently dropping the graphs ctx
7192        // (the stream linear arm takes linear_attn_verify_t, not the graphed segment or
7193        // full-verify bodies).
7194        if stream.is_some() && graphs.is_some() {
7195            return Err(
7196                "qwen35 tparallel verify: ROUND-STREAM and dspark verify graphs \
7197                        cannot arm together"
7198                    .into(),
7199            );
7200        }
7201        // Engine-bundle slice 3 + slice 4c: with a graphs ctx armed, pointer tables are
7202        // refreshed once per verify (the gdn ping-pong moves handles; a fresh generation
7203        // moves the kv caches). Then:
7204        //  - slice 4c: when the WHOLE round rides one seqs rung (every row batchable, one
7205        //    split-ladder step, rung covers the round), the ENTIRE walk replays as ONE
7206        //    full-verify graph per (vt, rung) — linear layers through the shared
7207        //    `qwen35_tparallel_linear_layer` body, full-attention layers through the
7208        //    shared `qwen35_tparallel_fa_layer` body in graph mode.
7209        //  - fallback (straddle rounds, below the vec floor, partial walks): runs of
7210        //    consecutive LINEAR layers replay the slice-3 per-(segment, vt) graphs and
7211        //    the full-attention layers run eager (batched rows when eligible).
7212        if let Some(g) = graphs.as_deref_mut() {
7213            g.refresh_tables(e, cache)?;
7214            g.round_slab = false;
7215            if let Some(rung) = g.full_rung(self, cache, lo, hi, t, seqs_append && batch_fa_on) {
7216                // Pool ceiling (dspark_vg_cap): an existing key always replays; a NEW
7217                // full capture past the ceiling falls through to the segment/eager arms.
7218                if g.full.contains_key(&(t, rung, hi)) || g.can_capture() {
7219                    let out = g.run_full(self, e, lo, hi, &x, t, pos0, rung, cache)?;
7220                    g.round_slab = true;
7221                    return Ok(out);
7222                }
7223            }
7224            // Round-atomic ceiling check for the segment door: if any linear run in this
7225            // walk would need a NEW capture past the ceiling, the whole round runs the
7226            // eager cols-ckpt walk (mixing slab- and cols-stashed layers in one round
7227            // would corrupt the commit).
7228            if !g.segments_ready(self, lo, hi, t) {
7229                graphs = None;
7230            }
7231        }
7232        // STREAM (2b, lane/draftcost-moe): positions come from the device round counter
7233        // (pos_iota / i32_copy_add) so a burst round needs no host position knowledge.
7234        let pos_d = match stream {
7235            Some((_, ctr)) => {
7236                let mut p = e.alloc_uninit::<i32>(t)?;
7237                e.pos_iota(ctr, &mut p, t)?;
7238                p
7239            }
7240            None => {
7241                let pos_host: Vec<i32> = (0..t).map(|r| (pos0 + r) as i32).collect();
7242                e.htod_i32(&pos_host)?
7243            }
7244        };
7245        // Per-row 1-element position buffers, built ONCE per verify (the append/fa wrappers
7246        // take owned pos slices; building these inside the layer x row loops cost 16xT H2Ds).
7247        // LAZY since slice 4: the batched fa/append arm never touches them — they are built
7248        // on the first per-row fallback layer only (stream-aware there; the stream FA arm
7249        // rides the dc rows kernels and never reaches the fallback).
7250        let mut pos_rows: Option<Vec<CudaSlice<i32>>> = None;
7251        let mut il = lo;
7252        while il < hi {
7253            if graphs.is_some() && matches!(self.layers[il].mixer, Mixer::Linear(_)) {
7254                let mut end = il;
7255                while end < hi && matches!(self.layers[end].mixer, Mixer::Linear(_)) {
7256                    end += 1;
7257                }
7258                let g = graphs.as_deref_mut().expect("checked above");
7259                x = g.run_segment(self, e, il, end, &x, t, cache)?;
7260                g.round_slab = true;
7261                il = end;
7262                continue;
7263            }
7264            let layer = &self.layers[il];
7265            if stream.is_none() && matches!(layer.mixer, Mixer::Linear(_)) {
7266                // Eager linear layer (no graphs ctx): the shared body, legacy cols-ckpt arm.
7267                // Under ROUND-STREAM the linear layers ride the fa-body match's stream arm
7268                // below (linear_attn_verify_t — the stream COMMIT needs its GdnStash).
7269                x = self.qwen35_tparallel_linear_layer(
7270                    e,
7271                    il,
7272                    &x,
7273                    t,
7274                    cache,
7275                    ckpt.as_deref_mut(),
7276                    None,
7277                    None,
7278                )?;
7279                il += 1;
7280                continue;
7281            }
7282            // Full-attention (or stream-Linear, or MLA-refusing) layer: the extracted
7283            // shared body — eager arm (fresh per-verify pos/table, exact t_kv sizing,
7284            // in-body len bump). The slice-4c captured full-verify graphs run the SAME
7285            // body in graph mode; under ROUND-STREAM the body's dc-rows / GDN stream arms
7286            // run (lane/draftcost-moe).
7287            x = self.qwen35_tparallel_fa_layer(
7288                e,
7289                il,
7290                &x,
7291                t,
7292                cache,
7293                FaLayerArgs {
7294                    pos_d: &pos_d,
7295                    pos_rows: &mut pos_rows,
7296                    pos0,
7297                    seqs_append,
7298                    batch_fa_on,
7299                    graph_cap: None,
7300                    stream,
7301                    ckpt: ckpt.as_deref_mut(),
7302                },
7303            )?;
7304            il += 1;
7305        }
7306        Ok(x)
7307    }
7308
7309    /// SHARED dense-FFN body for the qwen35 t-parallel layers (trunk-kernels slice B) —
7310    /// ONE copy for the fa and linear layer bodies (the verify_layers extraction lesson).
7311    /// Dual arm (MEMRA_TK_FFN_DUAL, default on): gate+up in ONE dual launch from the
7312    /// pre-quantized activation with macro-scales DEFERRED into the fused SwiGLU+q8_1
7313    /// epilogue, then ffn_down from the fused (aq, ad) — the q27 verify chain verbatim.
7314    /// Every door is the bit-identical proven one: `matmul_decode_exact_dual_pre` (per
7315    /// (tensor,token,row) == the two singles), `silu_mul_scaled_q8_1` (y*s inline == the
7316    /// scale_inplace store, value-exact; fused quantize == quantize_q8_1 bytes),
7317    /// `matmul_decode_exact_pre` (dispatch mirror of the singles' q8_1-fast tail).
7318    /// Dual-refused (t outside 2..=7, non-NVFP4, layout mismatch) or seam off -> the
7319    /// original singles chain, byte-for-byte.
7320    #[allow(clippy::too_many_arguments)]
7321    fn qwen35_tparallel_dense_ffn(
7322        &self,
7323        e: &Engine,
7324        ffn_gate: &crate::model::GpuTensor,
7325        ffn_up: &crate::model::GpuTensor,
7326        ffn_down: &crate::model::GpuTensor,
7327        zn: &CudaSlice<f32>,
7328        t: usize,
7329        n_embd: usize,
7330    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7331        let n_ff = ffn_gate.out_features();
7332        let (zq, zd) = e.quantize_q8_1(zn, t, n_embd)?;
7333        if Engine::tk_ffn_dual_on() {
7334            if let Some(((g, gs), (u, us))) =
7335                e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, &zq, &zd, t)?
7336            {
7337                if e.uses_q8_1_fast(ffn_down) {
7338                    let (aq, ad) = e.silu_mul_scaled_q8_1(&g, &u, gs, us, t * n_ff)?;
7339                    return e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t);
7340                }
7341                let mut act = e.uninit(t * n_ff)?;
7342                e.silu_mul_scaled(&g, &u, gs, us, &mut act, t * n_ff)?;
7343                let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7344                return e.matmul_pre(ffn_down, &aq, &ad, &act, t);
7345            }
7346        }
7347        // v1 singles chain (seam off or dual-refused) — the pre-slice-B body verbatim.
7348        let g = e.matmul_pre(ffn_gate, &zq, &zd, zn, t)?;
7349        let u = e.matmul_pre(ffn_up, &zq, &zd, zn, t)?;
7350        let mut act = e.uninit(t * n_ff)?;
7351        e.silu_mul(&g, &u, &mut act, t * n_ff)?;
7352        let (aq, ad) = e.quantize_q8_1(&act, t, n_ff)?;
7353        e.matmul_pre(ffn_down, &aq, &ad, &act, t)
7354    }
7355
7356    /// ONE t-parallel FULL-ATTENTION layer (attn_norm + fa mixer + post_attn_norm + FFN +
7357    /// tap) — extracted from the walk exactly like `qwen35_tparallel_linear_layer` so the
7358    /// eager walk and the slice-4c captured full-verify graphs execute the SAME body (a
7359    /// second copy is how dispatch mirrors drift — the verify_layers extraction lesson).
7360    ///
7361    /// `args.graph_cap = Some((table, off, rung_end))` is the captured-graph mode:
7362    /// - kv base-pointer pairs come from the ctx-owned persistent table at `off` (a fresh
7363    ///   generation's cache lands at new addresses that only the per-verify table refresh
7364    ///   knows — the slice-3 baked-address lesson);
7365    /// - the seqs twins size partials/grid at `rung_end` and pin `split_keys` to the
7366    ///   rung's ladder value: `n_splits_max` is pure stride, splits >= ns_eff write the
7367    ///   EMPTY partial the combine never reads, and every per-row T_kv derives in-kernel
7368    ///   from `pos_seq[z]` — so one captured launch replays bit-identically for every
7369    ///   round whose rows all sit inside the rung;
7370    /// - the host len bump moves to the replay caller (captured host code does not
7371    ///   re-run at replay).
7372    /// Graph mode REFUSES any round the batched arm cannot take: the per-row fallback
7373    /// host-branches on t_kv and must never be captured.
7374    #[allow(clippy::too_many_arguments)]
7375    fn qwen35_tparallel_fa_layer(
7376        &self,
7377        e: &Engine,
7378        il: usize,
7379        x: &CudaSlice<f32>,
7380        t: usize,
7381        cache: &mut Cache,
7382        args: FaLayerArgs<'_>,
7383    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7384        use cudarc::driver::DevicePtr;
7385        let cfg = &self.cfg;
7386        let n_embd = cfg.n_embd as usize;
7387        let eps = cfg.rms_eps;
7388        let head_dim_global = cfg.head_dim_k as usize;
7389        let layer = &self.layers[il];
7390        let FaLayerArgs {
7391            pos_d,
7392            pos_rows,
7393            pos0,
7394            seqs_append,
7395            batch_fa_on,
7396            graph_cap,
7397            stream,
7398            mut ckpt,
7399        } = args;
7400
7401        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7402        let anorm = layer.attn_norm.float_data();
7403        let mut xn = e.uninit(t * n_embd)?;
7404        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7405        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7406
7407        let mixed: CudaSlice<f32> = match &layer.mixer {
7408            Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
7409            // STREAM ARM (2b, lane/draftcost-moe): under a device position counter the
7410            // per-row serving-kernel chain cannot run (host state swaps keyed on host
7411            // row index are fine, but the stream COMMIT needs the GdnStash for its _dc
7412            // rebuild — the per-row chain only produces per-column clones). GDN rides
7413            // `linear_attn_verify_t`: batched q8_1-class projections, stash-producing,
7414            // and its one-scan recurrence is pinned bit-identical to T chained T=1
7415            // steps (its header + kernel-check). Position-independent, so no counter
7416            // plumbing is needed. Guards mirror the generic call site exactly.
7417            Mixer::Linear(la) if stream.is_some() => {
7418                if !(t >= 3 || (t == 2 && spec_m2()))
7419                    || !self.mixer_in_q8_1_fast(e, &layer.mixer)
7420                    || !e.uses_q8_1_fast(&la.ssm_out)
7421                {
7422                    return Err("qwen35 stream verify: GDN batched arm requires t>=3 \
7423                                (or MEMRA_SPEC_M2 at t=2) and q8_1-fast projections"
7424                        .into());
7425                }
7426                let want = ckpt.is_some();
7427                let (out, stash) =
7428                    self.linear_attn_verify_t(e, la, &xn, Some((&hq, &hd)), t, cache, il, want)?;
7429                if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
7430                    ck.gdn[il] = Some(st);
7431                }
7432                out
7433            }
7434            Mixer::Linear(_) => {
7435                unreachable!("linear layers ride qwen35_tparallel_linear_layer")
7436            }
7437            Mixer::Full(fa) => {
7438                let geometry = cfg.full_attention_geometry_at(il as u32);
7439                let n_head = geometry.n_head as usize;
7440                let n_head_kv = geometry.n_head_kv as usize;
7441                let head_dim = geometry.head_dim_k as usize;
7442                let rope_dims = geometry.n_rot as usize;
7443                let rope_base = geometry.rope_base;
7444                let scale = geometry.attention_scale();
7445                // Batched projections: one weight read serves all T rows.
7446                // GROUP-3 twin (trunk-kernels slice D): q/k/v in ONE launch — the group4
7447                // kernel with n3=0, bit-identical per (tensor, token, row) to the three
7448                // singles; refused or MEMRA_TK_FA_GROUP=0 -> singles byte-for-byte.
7449                let (qf, mut k, v) = match e.matmul_decode_exact_group3_pre(
7450                    [&fa.wq, &fa.wk, &fa.wv],
7451                    &hq,
7452                    &hd,
7453                    t,
7454                )? {
7455                    Some(mut g3) => {
7456                        let v = g3.pop().unwrap();
7457                        let k = g3.pop().unwrap();
7458                        let qf = g3.pop().unwrap();
7459                        (qf, k, v)
7460                    }
7461                    None => (
7462                        e.matmul_pre(&fa.wq, &hq, &hd, &xn, t)?,
7463                        e.matmul_pre(&fa.wk, &hq, &hd, &xn, t)?,
7464                        e.matmul_pre(&fa.wv, &hq, &hd, &xn, t)?,
7465                    ),
7466                };
7467                let gated =
7468                    geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
7469                let (mut q, gate) = if gated {
7470                    let mut qs = e.uninit(t * n_head * head_dim)?;
7471                    let mut gs = e.uninit(t * n_head * head_dim)?;
7472                    e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, t)?;
7473                    (qs, Some(gs))
7474                } else {
7475                    (qf, None)
7476                };
7477                let mut qn = e.uninit(t * n_head * head_dim)?;
7478                e.rms_norm(
7479                    &q,
7480                    fa.q_norm.float_data(),
7481                    &mut qn,
7482                    head_dim,
7483                    t * n_head,
7484                    eps,
7485                )?;
7486                q = qn;
7487                let mut kn = e.uninit(t * n_head_kv * head_dim)?;
7488                e.rms_norm(
7489                    &k,
7490                    fa.k_norm.float_data(),
7491                    &mut kn,
7492                    head_dim,
7493                    t * n_head_kv,
7494                    eps,
7495                )?;
7496                k = kn;
7497                e.rope_neox(
7498                    &mut q, pos_d, head_dim, rope_dims, n_head, t, rope_base, 1.0,
7499                )?;
7500                e.rope_neox(
7501                    &mut k, pos_d, head_dim, rope_dims, n_head_kv, t, rope_base, 1.0,
7502                )?;
7503
7504                // Per-row append + attend: row r sees rows 0..r in KV (causal within the
7505                // draft), each through the b_n=1 serving kernels at its own t_kv.
7506                let q_dim = n_head * head_dim;
7507                let kv_dim = n_head_kv * head_dim;
7508                let mut attn = e.uninit(t * q_dim)?;
7509                let (kdk, kdv, ktb, vtb, len0, kv_local) = {
7510                    let kvl = cache.kv[il].as_ref().unwrap();
7511                    // [2T] interleaved k,v base pointers: entry pair z serves row z of
7512                    // the batched twins; the per-row fallback reads pair 0 (same cache
7513                    // for every row of one layer). Graph mode reads the ctx table.
7514                    let local: Option<CudaSlice<u64>> = match graph_cap {
7515                        Some(_) => None,
7516                        None => {
7517                            let s = &e.gpu.stream();
7518                            let (pk, _g) = kvl.k.device_ptr(s);
7519                            let (pv, _g2) = kvl.v.device_ptr(s);
7520                            let mut tbl = Vec::with_capacity(2 * t);
7521                            for _ in 0..t {
7522                                tbl.push(pk as u64);
7523                                tbl.push(pv as u64);
7524                            }
7525                            Some(e.htod_u64(&tbl)?)
7526                        }
7527                    };
7528                    (
7529                        kvl.kv_dim_k,
7530                        kvl.kv_dim_v,
7531                        kvl.k_tok_bytes,
7532                        kvl.v_tok_bytes,
7533                        kvl.len,
7534                        local,
7535                    )
7536                };
7537                let (kv_tbl, kv_off): (&CudaSlice<u64>, usize) = match graph_cap {
7538                    Some((tb, off, _)) => (tb, off),
7539                    None => (kv_local.as_ref().expect("built above"), 0),
7540                };
7541                // Slice 4 (fa/append rows — see dspark_fa_rows_on): the whole per-row
7542                // section batches into the z-batched serving twins when every row of
7543                // this round takes the v4-seqs arm on ONE fa_split_keys rung. Both
7544                // guards are evaluated at the round's FIRST and LAST t_kv — the
7545                // eligibility window (vec floor .. v4 max) and each split-ladder rung
7546                // are intervals in t_kv, so ends-inside means all-inside (the straddle
7547                // law). Appending all T rows before any attend is read-equivalent to
7548                // the interleaved order: row r's walk reads keys 0..len0+r only, and
7549                // rows > r land at slots it never touches; every written cache row is
7550                // the per-token appender's exact warp program (kernel-check pinned).
7551                let t_kv_first = len0 + 1;
7552                let t_kv_last = len0 + t;
7553                let rows_batched = t >= 2
7554                    && seqs_append
7555                    && batch_fa_on
7556                    && dspark_fa_rows_on()
7557                    // the z-batched twins read stacked rows at the CACHE's kv dims;
7558                    // the projection stack is [T, n_head_kv*head_dim] — they must be
7559                    // the same stride or row z misaligns (true for this family; the
7560                    // guard keeps any asymmetric-kv model on the per-row loop).
7561                    && kdk == kv_dim
7562                    && kdv == kv_dim
7563                    && crate::fa_seqs_eligible(t_kv_first, head_dim_global)
7564                    && crate::fa_seqs_eligible(t_kv_last, head_dim_global)
7565                    && crate::fa_split_keys(t_kv_first, cfg.n_head_kv as usize)
7566                        == crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize);
7567                // Sizing: eager = exact round bound; graph mode = the rung end (stride +
7568                // grid only — bytes proven equal above). Capture-time invariants refuse
7569                // loudly rather than bake a divergent body.
7570                let (size_kv_max, sp) = match graph_cap {
7571                    Some((_, _, rung)) => {
7572                        if !rows_batched {
7573                            return Err(format!(
7574                                "fa graph capture: layer {il} round is not batchable \
7575                                 (t_kv {t_kv_first}..{t_kv_last}) — the per-row fallback \
7576                                 must never be captured"
7577                            )
7578                            .into());
7579                        }
7580                        let sp_r = crate::fa_split_keys(rung, cfg.n_head_kv as usize);
7581                        if t_kv_last > rung
7582                            || sp_r != crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize)
7583                        {
7584                            return Err(format!(
7585                                "fa graph capture: rung {rung} does not cover round \
7586                                 t_kv {t_kv_first}..{t_kv_last} on one split ladder step"
7587                            )
7588                            .into());
7589                        }
7590                        (rung, sp_r)
7591                    }
7592                    None => (
7593                        t_kv_last,
7594                        crate::fa_split_keys(t_kv_last, cfg.n_head_kv as usize),
7595                    ),
7596                };
7597                if let Some((_, ctr)) = stream {
7598                    // STREAM ARM (2b): one batched dc append + the multi-row dc attention
7599                    // — the generic stream arm's exact shape (rows kernels are pinned
7600                    // byte-identical to the per-row programs by kernel-check). Host len
7601                    // stays a stale lower bound; the burst drain reconciles it.
7602                    let kvl = cache.kv[il].as_mut().unwrap();
7603                    e.append_kv_quantized_rows_dc(
7604                        &k,
7605                        &v,
7606                        &mut kvl.k,
7607                        &mut kvl.v,
7608                        ctr,
7609                        t,
7610                        kdk,
7611                        kdv,
7612                        ktb,
7613                        vtb,
7614                        Engine::kv_fp8_on(),
7615                    )?;
7616                    let upper = (kvl.len + t + 64).min(cache.max_ctx);
7617                    let k_view = e.view_u8(&kvl.k, upper * ktb);
7618                    let v_view = e.view_u8(&kvl.v, upper * vtb);
7619                    e.fa_decode_rows_dc(
7620                        &q, &k_view, &v_view, &mut attn, head_dim, n_head, n_head_kv, ctr, upper,
7621                        t, scale, ktb, vtb, 0, false,
7622                    )?;
7623                } else if rows_batched {
7624                    e.append_kv_quantized_seqs(
7625                        &k,
7626                        &v,
7627                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7628                        pos_d,
7629                        t,
7630                        kdk,
7631                        kdv,
7632                        ktb,
7633                        vtb,
7634                    )?;
7635                    if graph_cap.is_none() {
7636                        cache.kv[il].as_mut().unwrap().len += t;
7637                    }
7638                    e.fa_decode_batch_seqs_v4(
7639                        &q,
7640                        &kv_tbl.slice(kv_off..kv_off + 2 * t),
7641                        pos_d,
7642                        &mut attn,
7643                        head_dim,
7644                        n_head,
7645                        n_head_kv,
7646                        t,
7647                        size_kv_max,
7648                        scale,
7649                        sp,
7650                        ktb,
7651                        vtb,
7652                    )?;
7653                } else {
7654                    if pos_rows.is_none() {
7655                        // Stream-aware for symmetry with pos_d (the stream FA arm rides
7656                        // the dc rows kernels above and never reaches this fallback).
7657                        *pos_rows = Some(match stream {
7658                            Some((_, ctr)) => (0..t)
7659                                .map(|r| {
7660                                    let mut b = e.alloc_uninit::<i32>(1)?;
7661                                    e.i32_copy_add(ctr, &mut b, r as i32)?;
7662                                    Ok(b)
7663                                })
7664                                .collect::<Result<_, Box<dyn std::error::Error>>>()?,
7665                            None => (0..t)
7666                                .map(|r| e.htod_i32(&[(pos0 + r) as i32]))
7667                                .collect::<Result<_, _>>()?,
7668                        });
7669                    }
7670                    let pos_rows = pos_rows.as_ref().unwrap();
7671                    for r in 0..t {
7672                        // Owned per-row scratch: the b_n=1 kernels take packed batch buffers
7673                        // whose row 0 is this row (arithmetic-free materialization copies,
7674                        // same as decode's per-seq fallback arm).
7675                        let mut k_row = e.uninit(kv_dim)?;
7676                        e.dtod_copy_view(&k.slice(r * kv_dim..(r + 1) * kv_dim), &mut k_row)?;
7677                        let mut v_row = e.uninit(kv_dim)?;
7678                        e.dtod_copy_view(&v.slice(r * kv_dim..(r + 1) * kv_dim), &mut v_row)?;
7679                        let pos_row = &pos_rows[r];
7680                        let kvl = cache.kv[il].as_mut().unwrap();
7681                        if seqs_append {
7682                            e.append_kv_quantized_seqs(
7683                                &k_row,
7684                                &v_row,
7685                                &kv_tbl.slice(kv_off..kv_off + 2),
7686                                pos_row,
7687                                1,
7688                                kdk,
7689                                kdv,
7690                                ktb,
7691                                vtb,
7692                            )?;
7693                            kvl.len += 1;
7694                        } else {
7695                            e.append_kv_quantized_view(
7696                                &k_row.slice(0..kv_dim),
7697                                &v_row.slice(0..kv_dim),
7698                                &mut kvl.k,
7699                                &mut kvl.v,
7700                                kvl.len,
7701                                kvl.kv_dim_k,
7702                                kvl.kv_dim_v,
7703                                kvl.k_tok_bytes,
7704                                kvl.v_tok_bytes,
7705                                Engine::kv_fp8_on(),
7706                            )?;
7707                            kvl.len += 1;
7708                        }
7709                        let t_kv = kvl.len;
7710                        let mut q_row = e.uninit(q_dim)?;
7711                        e.dtod_copy_view(&q.slice(r * q_dim..(r + 1) * q_dim), &mut q_row)?;
7712                        let mut a_row = e.uninit(q_dim)?;
7713                        if batch_fa_on && crate::fa_seqs_eligible(t_kv, head_dim_global) {
7714                            let sp0_r = crate::fa_split_keys(t_kv, cfg.n_head_kv as usize);
7715                            e.fa_decode_batch_seqs_v4(
7716                                &q_row,
7717                                &kv_tbl.slice(kv_off..kv_off + 2),
7718                                pos_row,
7719                                &mut a_row,
7720                                head_dim,
7721                                n_head,
7722                                n_head_kv,
7723                                1,
7724                                t_kv,
7725                                scale,
7726                                sp0_r,
7727                                ktb,
7728                                vtb,
7729                            )?;
7730                        } else {
7731                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
7732                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
7733                            let mut a_view = a_row.slice_mut(0..q_dim);
7734                            e.fa_decode_kvmod_view(
7735                                &q_row.slice(0..q_dim),
7736                                &k_view,
7737                                &v_view,
7738                                &mut a_view,
7739                                head_dim,
7740                                n_head,
7741                                n_head_kv,
7742                                t_kv,
7743                                scale,
7744                                kvl.k_tok_bytes,
7745                                kvl.v_tok_bytes,
7746                                Engine::kv_fp8_on(),
7747                            )?;
7748                        }
7749                        e.dtod_copy_into(&a_row, &mut attn, r * q_dim)?;
7750                    }
7751                }
7752
7753                // Output gate (element-wise) + o-proj at m=T.
7754                let attn_g = match &gate {
7755                    Some(g) => {
7756                        let n = t * q_dim;
7757                        let mut gsig = e.uninit(n)?;
7758                        e.sigmoid(g, &mut gsig, n)?;
7759                        let mut ag = e.uninit(n)?;
7760                        e.mul(&attn, &gsig, &mut ag, n)?;
7761                        ag
7762                    }
7763                    None => attn,
7764                };
7765                e.matmul(&fa.wo, &attn_g, t)?
7766            }
7767        };
7768
7769        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
7770        let pnorm = layer.post_attn_norm.float_data();
7771        let mut x1 = e.uninit(t * n_embd)?;
7772        let mut zn = e.uninit(t * n_embd)?;
7773        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
7774        let ffn_out = match &layer.ffn {
7775            crate::hybrid::Ffn::Dense {
7776                ffn_gate,
7777                ffn_up,
7778                ffn_down,
7779            } => {
7780                assert!(
7781                    self.cfg.m3.is_none(),
7782                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
7783                );
7784                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
7785            }
7786            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
7787        };
7788        let mut x2 = e.uninit(t * n_embd)?;
7789        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
7790        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
7791        self.dflash_tap(e, cache, il, &x2, t)?;
7792        Ok(x2)
7793    }
7794
7795    /// ONE t-parallel LINEAR layer (attn_norm + gdn mixer + post_attn_norm + FFN + tap) —
7796    /// the exact body the old in-loop Linear arm ran, extracted so the eager walk and the
7797    /// slice-3 captured segments execute the SAME code (a second copy is how dispatch
7798    /// mirrors drift — the verify_layers extraction lesson). Two deliberate changes, both
7799    /// bit-identical by construction:
7800    /// - the gdn ping-pong host swap moves from per-row to ONE end-of-body swap (t odd):
7801    ///   the device sequence is driven entirely by the 6-entry pointer table, which
7802    ///   already encodes both parities; the ckpt stash reads name row r's out buffer
7803    ///   directly (r even -> alt handle, odd -> canonical) — the same physical bytes the
7804    ///   legacy post-swap clone read.
7805    /// - `stash` (slice-3 ctx): persistent per-layer slabs written by copy_into instead of
7806    ///   per-row clone_dtod allocs — same bytes, capture-legal (no per-round host objects).
7807    /// `table_src` = (persistent pointer table, offset) when the ctx owns the tables;
7808    /// None builds the per-verify table exactly as before.
7809    #[allow(clippy::too_many_arguments)]
7810    fn qwen35_tparallel_linear_layer(
7811        &self,
7812        e: &Engine,
7813        il: usize,
7814        x: &CudaSlice<f32>,
7815        t: usize,
7816        cache: &mut Cache,
7817        mut ckpt: Option<&mut VerifyCkpt>,
7818        stash: Option<(&mut CudaSlice<f32>, &mut CudaSlice<f32>)>,
7819        table_src: Option<(&CudaSlice<u64>, usize)>,
7820    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
7821        use cudarc::driver::DevicePtr;
7822        let cfg = &self.cfg;
7823        let n_embd = cfg.n_embd as usize;
7824        let eps = cfg.rms_eps;
7825        let layer = &self.layers[il];
7826        let Mixer::Linear(la) = &layer.mixer else {
7827            return Err("qwen35_tparallel_linear_layer on a non-linear layer".into());
7828        };
7829        // ---- attn_norm + q8_1 quantize at m=T (row-indexed == per-row) ----
7830        let anorm = layer.attn_norm.float_data();
7831        let mut xn = e.uninit(t * n_embd)?;
7832        e.rms_norm(x, anorm, &mut xn, n_embd, t, eps)?;
7833        let (hq, hd) = e.quantize_q8_1(&xn, t, n_embd)?;
7834
7835        let geometry = la.geometry;
7836        let d_state = geometry.key_head_dim as usize;
7837        let num_k = geometry.key_heads as usize;
7838        let num_v = geometry.value_heads as usize;
7839        let d_conv = geometry.conv_kernel as usize;
7840        let key_dim = d_state * num_k;
7841        let value_dim = geometry.value_head_dim as usize * num_v;
7842        let conv_dim = key_dim * 2 + value_dim;
7843        let gdn_scale = 1.0 / (d_state as f32).sqrt();
7844
7845        // ---- batched projections: one weight read for all T rows ----
7846        // GROUP-4 twin (trunk-kernels slice C): the whole 4-tuple in ONE launch, bit-identical
7847        // per (tensor, token, row) to the four singles; refused (layout/tier) or
7848        // MEMRA_TK_GDN_GROUP=0 -> the singles chain byte-for-byte.
7849        let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_decode_exact_group4_pre(
7850            [&la.wqkv, &la.wqkv_gate, &la.ssm_beta, &la.ssm_alpha],
7851            &hq,
7852            &hd,
7853            t,
7854        )? {
7855            Some(mut g4) => {
7856                let alpha = g4.pop().unwrap();
7857                let beta_raw = g4.pop().unwrap();
7858                let z = g4.pop().unwrap();
7859                let qkv_mixed = g4.pop().unwrap();
7860                (qkv_mixed, z, beta_raw, alpha)
7861            }
7862            None => (
7863                e.matmul_pre(&la.wqkv, &hq, &hd, &xn, t)?,
7864                e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, t)?,
7865                e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, t)?,
7866                e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, t)?,
7867            ),
7868        };
7869        let beta_w = la.ssm_beta.out_features();
7870        let alpha_w = la.ssm_alpha.out_features();
7871        let qkv_w = la.wqkv.out_features();
7872
7873        // ---- per-row state chain through the b_n=1 serving kernels ----
7874        // 6-entry alternating pointer table expresses the ping-pong without a rebuild per
7875        // row: even rows scan s0 -> s1, odd rows s1 -> s0.
7876        let table_local: Option<CudaSlice<u64>> = match table_src {
7877            Some(_) => None,
7878            None => {
7879                let rl = cache.recur[il].as_ref().unwrap();
7880                let s = &e.gpu.stream();
7881                let (pc, _g0) = rl.conv_state.device_ptr(s);
7882                let (p0, _g1) = rl.ssm_state.device_ptr(s);
7883                let (p1, _g2) = rl.ssm_state_alt.device_ptr(s);
7884                Some(e.htod_u64(&[
7885                    pc as u64, p0 as u64, p1 as u64, pc as u64, p1 as u64, p0 as u64,
7886                ])?)
7887            }
7888        };
7889        let (table, toff): (&CudaSlice<u64>, usize) = match table_src {
7890            Some((tb, off)) => (tb, off),
7891            None => (table_local.as_ref().unwrap(), 0),
7892        };
7893        let mut o_all = e.uninit(t * value_dim)?;
7894        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
7895            if ckpt.is_some() && stash.is_none() && t >= 2 {
7896                Some(Vec::with_capacity(t - 1))
7897            } else {
7898                None
7899            };
7900        let mut stash = stash;
7901        // Per-row scratch reused across rows (uninit is cheap but not free at
7902        // 48 layers x T rows); row inputs/outputs pass as VIEWS into the packed
7903        // [T, ...] buffers — zero arithmetic-free copies in this loop.
7904        let mut conv_out = e.uninit(conv_dim)?;
7905        let mut q_l2 = e.uninit(value_dim)?;
7906        let mut k_l2 = e.uninit(value_dim)?;
7907        let mut v_gd = e.uninit(value_dim)?;
7908        let mut beta_b = e.uninit(num_v)?;
7909        let mut g_log = e.uninit(num_v)?;
7910        for r in 0..t {
7911            let base = toff + if r % 2 == 0 { 0 } else { 3 };
7912            let conv_view = table.slice(base..base + 1);
7913            let in_view = table.slice(base + 1..base + 2);
7914            let out_view = table.slice(base + 2..base + 3);
7915            e.ssm_conv1d_fused_decode_b_view(
7916                &qkv_mixed.slice(r * qkv_w..(r + 1) * qkv_w),
7917                &conv_view,
7918                la.ssm_conv1d.float_data(),
7919                &mut conv_out,
7920                conv_dim,
7921                d_conv,
7922                1,
7923            )?;
7924            e.gdn_prep_decode_b_view(
7925                &conv_out,
7926                &beta_raw.slice(r * beta_w..(r + 1) * beta_w),
7927                &alpha.slice(r * alpha_w..(r + 1) * alpha_w),
7928                la.ssm_dt.float_data(),
7929                la.ssm_a.float_data(),
7930                &mut q_l2,
7931                &mut k_l2,
7932                &mut v_gd,
7933                &mut beta_b,
7934                &mut g_log,
7935                d_state,
7936                num_v,
7937                num_k,
7938                key_dim,
7939                eps,
7940                conv_dim,
7941                1,
7942            )?;
7943            let mut o_row = o_all.slice_mut(r * value_dim..(r + 1) * value_dim);
7944            e.gdn_scan_s128_batched_view(
7945                &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_row, num_v, 1,
7946                gdn_scale,
7947            )?;
7948            if r + 1 < t {
7949                // Row r's out buffer: even rows write s1 (the alt handle — no swaps ran),
7950                // odd rows write s0 — the same physical state the legacy post-swap
7951                // canonical clone read.
7952                let rl = cache.recur[il]
7953                    .as_ref()
7954                    .ok_or("qwen35 linear verify layer has no recurrent state")?;
7955                let ssm_src = if r % 2 == 0 {
7956                    &rl.ssm_state_alt
7957                } else {
7958                    &rl.ssm_state
7959                };
7960                match stash.as_mut() {
7961                    Some((conv_slab, ssm_slab)) => {
7962                        // BOTH stash reads go through the pointer table at run time: the
7963                        // ssm handles ping-pong between rounds, and the ctx (with its
7964                        // captured graphs) outlives the Cache — a fresh generation's
7965                        // conv/ssm buffers land at new addresses that only the per-round
7966                        // table refresh knows. A baked direct copy would read freed
7967                        // memory (parity was the slice-3 smoke divergence; cache
7968                        // lifetime is the cross-generation twin).
7969                        e.copy_indirect_src_f32(
7970                            &conv_view,
7971                            conv_slab,
7972                            r * conv_dim * (d_conv - 1),
7973                            conv_dim * (d_conv - 1),
7974                        )?;
7975                        // The ssm handles PING-PONG between rounds: a captured direct
7976                        // copy would bake the capture-time physical buffer and read the
7977                        // wrong parity after any odd-vt round (the slice-3 smoke
7978                        // divergence). Read the src address from row r's OUT table
7979                        // entry at run time — the same entry the scan just wrote.
7980                        e.copy_indirect_src_f32(
7981                            &out_view,
7982                            ssm_slab,
7983                            r * d_state * d_state * num_v,
7984                            d_state * d_state * num_v,
7985                        )?;
7986                    }
7987                    None => {
7988                        if let Some(states) = col_states.as_mut() {
7989                            states.push((e.clone_dtod(&rl.conv_state)?, e.clone_dtod(ssm_src)?));
7990                        }
7991                    }
7992                }
7993            }
7994        }
7995        // ONE end-of-body parity swap (t odd) — the legacy loop swapped per row; the net
7996        // handle motion is identical and the device sequence never read the handles.
7997        if t % 2 == 1 {
7998            let rl = cache.recur[il].as_mut().unwrap();
7999            std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8000        }
8001        if let (Some(checkpoint), Some(states)) = (ckpt.as_deref_mut(), col_states) {
8002            checkpoint.cols[il] = Some(states);
8003        }
8004
8005        // ---- batched gated norm + out-projection at m=T ----
8006        let mixed = if e.uses_q8_1_fast(&la.ssm_out) {
8007            let (gq, gd) = e.gated_rmsnorm_q8_1(
8008                &o_all,
8009                la.ssm_norm.float_data(),
8010                &z,
8011                d_state,
8012                t * num_v,
8013                eps,
8014            )?;
8015            let g0 = e.zeros(0)?;
8016            e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, t)?
8017        } else {
8018            let mut gn = e.uninit(t * value_dim)?;
8019            e.gated_rmsnorm(
8020                &o_all,
8021                la.ssm_norm.float_data(),
8022                &z,
8023                &mut gn,
8024                d_state,
8025                t * num_v,
8026                eps,
8027            )?;
8028            e.matmul(&la.ssm_out, &gn, t)?
8029        };
8030
8031        // ---- residual add + post_attn_norm + FFN at m=T (serving dispatch verbatim) ----
8032        let pnorm = layer.post_attn_norm.float_data();
8033        let mut x1 = e.uninit(t * n_embd)?;
8034        let mut zn = e.uninit(t * n_embd)?;
8035        e.add_rms_norm(x, &mixed, pnorm, &mut x1, &mut zn, n_embd, t, eps)?;
8036        let ffn_out = match &layer.ffn {
8037            crate::hybrid::Ffn::Dense {
8038                ffn_gate,
8039                ffn_up,
8040                ffn_down,
8041            } => {
8042                assert!(
8043                    self.cfg.m3.is_none(),
8044                    "qwen35 t-parallel verify: M3 swigluoai FFN not yet batched"
8045                );
8046                self.qwen35_tparallel_dense_ffn(e, ffn_gate, ffn_up, ffn_down, &zn, t, n_embd)?
8047            }
8048            crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &zn, None, t, il as u16)?,
8049        };
8050        let mut x2 = e.uninit(t * n_embd)?;
8051        e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
8052        // dspark drafter tap (no-op when no sink armed): post-layer residual verify rows
8053        self.dflash_tap(e, cache, il, &x2, t)?;
8054        Ok(x2)
8055    }
8056
8057    /// PP-N STAGE SUBGRAPH of the verify trunk: layers `[lo, hi)` of `decode_step_t_core_stream`'s
8058    /// walk, verbatim. Enters with a MATERIALIZED `[T, n_embd]` residual (no pending fusion pair
8059    /// carried in from outside the range) and exits with the range's final residual materialized
8060    /// (the trailing add executed) — exactly the `decode_layers_eager(lo, hi)` contract, T rows
8061    /// instead of one.
8062    ///
8063    /// EXTRACTED (lane/pp2-spec 2026-08-06) rather than duplicated: `decode_step_t_core_stream` IS
8064    /// the single funnel every verify forward reaches, and its per-layer dispatch MIRRORING (norm
8065    /// fusion per layer, the t>=3/spec_m2 batched-linear window, the fused-q8 FFN chain, the
8066    /// decode-exact projections) is what makes verify bit-identical to eager decode. A second copy
8067    /// for the split arm is how those mirrors drift apart on the next lever. The unsplit body now
8068    /// calls this with `(0, n_layers)`, so the whole-trunk path and every stage range run the SAME
8069    /// code — there is no "split version" of the verify math.
8070    ///
8071    /// Bit-identity of a cut rests on the same kernel-check-pinned identity the eager arm's cut
8072    /// does — `add_rms_norm_q8_1 == add then rms_norm_q8_1` at nrows=T — because the ONLY thing a
8073    /// fence changes is that the cross-layer fusion carry breaks at `hi-1` and is re-materialized
8074    /// as an explicit `add`. `decode-batch-gate --mode ppspec` verifies end-to-end on real weights.
8075    #[allow(clippy::too_many_arguments)]
8076    fn verify_layers(
8077        &self,
8078        e: &Engine,
8079        mut x: CudaSlice<f32>,
8080        lo: usize,
8081        hi: usize,
8082        pos_d: &CudaSlice<i32>,
8083        pos0: usize,
8084        t: usize,
8085        cache: &mut Cache,
8086        mut ckpt: Option<&mut VerifyCkpt>,
8087        stream: Option<(&CudaSlice<u32>, &CudaSlice<i32>)>,
8088        graphs: Option<&mut DsparkVerifyGraphs>,
8089    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
8090        if self.sliding_gated_moe_batch_program() {
8091            if stream.is_some() {
8092                return Err(
8093                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
8094                            cannot express the SWA offset KV view)"
8095                        .into(),
8096                );
8097            }
8098            return self.step35_verify_batch_layers(e, x, lo, hi, pos0, t, cache);
8099        }
8100        if self.batched_serving_numeric_class() {
8101            return self.qwen35_verify_batch_layers(
8102                e,
8103                x,
8104                lo,
8105                hi,
8106                pos0,
8107                t,
8108                cache,
8109                ckpt.take(),
8110                stream,
8111                graphs,
8112            );
8113        }
8114        let n_embd = self.cfg.n_embd as usize;
8115        let eps = self.cfg.rms_eps;
8116        // CROSS-LAYER ADD+NORM FUSION (lane/vt-fixes fix 2, mirroring decode_step_h's
8117        // launch-arc form): layer il's post-FFN residual add (x2 = x1 + ffn_out) and layer
8118        // il+1's attn_norm(+quantize) are consecutive row-wise ops — ONE add_rms_norm_q8_1
8119        // launch at nrows=t does all three (bit-identity pinned by the T-row kernel-check
8120        // arms). Carry the un-added (x1, ffn_out) pair; the fused launch materializes x2 (the
8121        // residual the next layer needs) as its `res` output. Falls back to the separate add
8122        // when the next layer is off the fused-q8 path.
8123        let mut pending: Option<(CudaSlice<f32>, CudaSlice<f32>)> = None;
8124        for il in lo..hi {
8125            let layer = &self.layers[il];
8126            // DISPATCH-MIRRORED attn-input RMSNorm (FP-order lesson #8): eager decode fuses the
8127            // 1024-thread rms_norm_q8_1 ONLY when every mixer projection is q8_1-fast; layers with
8128            // Float projections (ssm_beta/ssm_alpha on layers 1/2/4 of the 9B NVFP4 GGUF) take the
8129            // UNFUSED 256-thread rms_norm. The verify norm must mirror that PER-LAYER choice —
8130            // blockDim changes the sum-of-squares reduce order, and the ULP shift amplifies through
8131            // the GDN recurrence into argmax flips (measured: 9B text prompt, 1 ULP at layer 2 ->
8132            // 2.3e-1 logit maxdiff at the head -> K=1..8 divergence at a 0.03-margin token).
8133            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8134            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8135            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2, 2026-08-03): when the norm is
8136            // dispatch-fused AND every consumer of `h` reads only its q8_1 form (Full mixer:
8137            // projections only; Linear mixer: the batched arm — the per-column fallback needs
8138            // f32 h), emit the attn-input norm DIRECTLY as q8_1 via `rms_norm_q8_1` at nrows=t
8139            // (row-indexed kernel — the T-row launch is the per-row m=1 program, kernel-check
8140            // pins bit-identity vs rms_norm_decode -> quantize_q8_1). Kills the standalone
8141            // quantize launch(es) + the f32 h HBM round-trip that decode never pays.
8142            // step35 (Full mixer) is the third case that needs f32 `h`: its verify arm is a
8143            // per-ROW replay of the eager decode mixer, whose `pre_q` contract is a single row —
8144            // a T-row q8_1 pair cannot be handed to it, and re-deriving per-row q8_1 from the f32
8145            // rows is exactly the dispatch being mirrored. Keep step35 on the unfused arm.
8146            let lin_q8_only = match &layer.mixer {
8147                Mixer::Linear(la) => {
8148                    (t >= 3 || (t == 2 && spec_m2())) && e.uses_q8_1_fast(&la.ssm_out)
8149                }
8150                Mixer::Full(_) if self.sliding_gated_moe_batch_program() => false,
8151                _ => true,
8152            };
8153            // NOTE decode.rs's take()-first lesson: take the pending pair BEFORE branching so
8154            // a non-fused layer still performs the residual add.
8155            let taken = pending.take();
8156            let (h, h_q8) = if norm_fused && lin_q8_only {
8157                let pair = match taken {
8158                    // fused add + attn_norm + q8_1: ONE launch resolves the carried residual
8159                    // AND emits this layer's mixer input pre-quantized. res -> x2 (= new x).
8160                    Some((x1p, f1p)) => {
8161                        let mut x2 = vbuf(e, t * n_embd)?; // fully written (res output)
8162                        let p = e.add_rms_norm_q8_1(
8163                            &x1p,
8164                            &f1p,
8165                            layer.attn_norm.float_data(),
8166                            &mut x2,
8167                            n_embd,
8168                            t,
8169                            eps,
8170                        )?;
8171                        x = x2;
8172                        p
8173                    }
8174                    None => e.rms_norm_q8_1(&x, layer.attn_norm.float_data(), n_embd, t, eps)?,
8175                };
8176                (e.zeros(0)?, Some(pair)) // h unused on this path (q8-only consumers)
8177            } else {
8178                if let Some((x1p, f1p)) = taken {
8179                    let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8180                    e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8181                    x = x2;
8182                }
8183                let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8184                if norm_fused {
8185                    e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8186                } else {
8187                    e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8188                }
8189                (h, None)
8190            };
8191            let h_q8_ref = h_q8.as_ref().map(|(q, d)| (q, d));
8192
8193            let mixed = match &layer.mixer {
8194                Mixer::Full(fa) => self.full_attn_verify(
8195                    e,
8196                    fa,
8197                    &h,
8198                    h_q8_ref,
8199                    pos_d,
8200                    t,
8201                    cache,
8202                    il,
8203                    stream.map(|(_, c)| c),
8204                )?,
8205                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8206                Mixer::Linear(la) => {
8207                    // BATCHED linear verify (2026-07-03, the MTP-profit lever): one T-token pass —
8208                    // batched projections (weight read ONCE, hits the m=2-4 weight-resident matvec),
8209                    // carried-state conv (ssm_conv1d_tm_state), GDN prep on the prefill kernels, and
8210                    // ONE gdn_scan whose internal sequential t-loop is the SAME recurrence as T
8211                    // chained T=1 steps (bit-identical). Falls back to the sequential per-column
8212                    // chain when T < d_conv-1 (conv ring update needs T >= pad) — or when ANY
8213                    // projection is off the q8_1 fast path: matmul_decode_exact would route a Float
8214                    // tensor to cuBLAS at m=t (different FP accumulation than eager's per-token
8215                    // GEMV), so mixed-dtype layers stay on the eager-identical per-column chain.
8216                    // MEMRA_SPEC_M2 (lane/spec-m2): the t==2 batch rides the same arm — the conv
8217                    // wrapper handles t<pad with a pure-copy ring rebuild; see spec_m2() header.
8218                    if (t >= 3 || (t == 2 && spec_m2()))
8219                        && mixer_fast
8220                        && e.uses_q8_1_fast(&la.ssm_out)
8221                    {
8222                        let want = ckpt.is_some();
8223                        let (out, stash) =
8224                            self.linear_attn_verify_t(e, la, &h, h_q8_ref, t, cache, il, want)?;
8225                        if let (Some(ck), Some(st)) = (ckpt.as_deref_mut(), stash) {
8226                            ck.gdn[il] = Some(st);
8227                        }
8228                        out
8229                    } else {
8230                        let mut out = vbuf(e, t * n_embd)?; // every col written by copy_into
8231                        let mut col_states: Option<Vec<(CudaSlice<f32>, CudaSlice<f32>)>> =
8232                            if ckpt.is_some() && t >= 2 {
8233                                Some(Vec::with_capacity(t - 1))
8234                            } else {
8235                                None
8236                            };
8237                        for col in 0..t {
8238                            let mut h_col = vbuf(e, n_embd)?; // fully written by copy_view_into
8239                            let src = h.slice(col * n_embd..(col + 1) * n_embd);
8240                            e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8241                            let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8242                            e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8243                            // REPLAY-FREE ckpt: clone the chain's ACTUAL state after this column
8244                            // (pure dtod — cannot change any computed value). Last column skipped:
8245                            // rebuild targets are j <= t-1 columns.
8246                            if let Some(cs) = col_states.as_mut() {
8247                                if col + 1 < t {
8248                                    let rl = cache.recur[il].as_ref().unwrap();
8249                                    cs.push((
8250                                        e.clone_dtod(&rl.conv_state)?,
8251                                        e.clone_dtod(&rl.ssm_state)?,
8252                                    ));
8253                                }
8254                            }
8255                        }
8256                        if let (Some(ck), Some(cs)) = (ckpt.as_deref_mut(), col_states) {
8257                            // ReplaySSM-assessment instrumentation (2026-07-30): the
8258                            // per-column clones are the only true state snapshots left in
8259                            // the verify (the batched path stashes INPUTS and replays).
8260                            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
8261                                static ONCE: std::sync::Once = std::sync::Once::new();
8262                                let bytes: usize =
8263                                    cs.iter().map(|(c, s)| (c.len() + s.len()) * 4).sum();
8264                                ONCE.call_once(|| eprintln!(
8265                                    "[verify-ckpt] per-column layer il={il}: {} clones, {:.2} MB/layer/round",
8266                                    cs.len(), bytes as f64 / 1e6));
8267                            }
8268                            ck.cols[il] = Some(cs);
8269                        }
8270                        out
8271                    }
8272                }
8273            };
8274
8275            // DISPATCH-MIRRORED post-attn norm: eager residual_norm_ffn fuses add+norm+quant
8276            // (1024-thread add_rms_norm_q8_1) only for Dense FFNs whose gate+up are q8_1-fast;
8277            // otherwise (and for MoE) it runs the 256-thread fused add_rms_norm. Mirror per layer.
8278            let ffn_fuse = match &layer.ffn {
8279                crate::hybrid::Ffn::Dense {
8280                    ffn_gate, ffn_up, ..
8281                } => {
8282                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8283                        && e.uses_q8_1_fast(ffn_gate)
8284                        && e.uses_q8_1_fast(ffn_up)
8285                }
8286                crate::hybrid::Ffn::Moe(_) => false,
8287            };
8288            // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): on the ffn_fuse path (Dense,
8289            // gate+up q8_1-fast, non-M3) the FFN input is emitted DIRECTLY as q8_1 by ONE
8290            // add_rms_norm_q8_1 launch at nrows=t (row-indexed kernel: the T-row launch is the
8291            // per-row m=1 program; kernel-check pins bit-identity vs the unfused
8292            // add_f32 -> rms_norm_decode -> quantize_q8_1 chain at T=2/4/5/8) — replacing the
8293            // add + rms_norm_decode launches AND the dual/singles' internal re-quantize.
8294            // M3's swigluoai must keep the f32 chain (the fused SwiGLU epilogue encodes plain
8295            // SiLU), mirroring residual_norm_ffn's m3 guard on the decode path.
8296            // step35: same guard per LAYER. A dense FFN's clamp is the SHEXP array (upstream's
8297            // one build_ffn serves dense + shared expert, llama-graph.cpp:1751), and verify MUST
8298            // mirror decode's dispatch or spec self-consistency fails.
8299            let dense_lim = self.cfg.clamp_shexp_at(il as u32);
8300            let fuse_q8 = ffn_fuse && self.cfg.m3.is_none() && dense_lim.is_none();
8301            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm*
8302            let mut z = e.zeros(0)?; // replaced below on the unfused arms
8303            let z_q8 = if fuse_q8 {
8304                Some(e.add_rms_norm_q8_1(
8305                    &x,
8306                    &mixed,
8307                    layer.post_attn_norm.float_data(),
8308                    &mut x1,
8309                    n_embd,
8310                    t,
8311                    eps,
8312                )?)
8313            } else {
8314                let mut zf = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8315                if ffn_fuse {
8316                    e.add(&x, &mixed, &mut x1, t * n_embd)?;
8317                    e.rms_norm_decode(
8318                        &x1,
8319                        layer.post_attn_norm.float_data(),
8320                        &mut zf,
8321                        n_embd,
8322                        t,
8323                        eps,
8324                    )?;
8325                } else {
8326                    e.add_rms_norm(
8327                        &x,
8328                        &mixed,
8329                        layer.post_attn_norm.float_data(),
8330                        &mut x1,
8331                        &mut zf,
8332                        n_embd,
8333                        t,
8334                        eps,
8335                    )?;
8336                }
8337                z = zf;
8338                None
8339            };
8340            // DECODE-EXACT FFN projections: force MMVQ for gate/up/down at any T to match the
8341            // T=1 decode FP accumulation order. At T>=5 the generic matmul/matmul_pre falls to dp4a
8342            // (128-thread, different FP sum order). At T=2-4 the batched MMVQ is already bit-identical.
8343            let ffn_out = match &layer.ffn {
8344                crate::hybrid::Ffn::Dense {
8345                    ffn_gate,
8346                    ffn_up,
8347                    ffn_down,
8348                } => {
8349                    let n_ff = ffn_gate.out_features();
8350                    if let Some((zq, zd)) = z_q8.as_ref() {
8351                        // FUSED CHAIN (fix 2): pre-quantized z feeds the projections; the SwiGLU
8352                        // epilogue emits act pre-quantized for ffn_down (silu_mul_scaled_q8_1,
8353                        // bit-identical to silu_mul + quantize — kernel-check-pinned) with the
8354                        // NVFP4 macro-scales folded (deferred-scale dual: y*s inline == the
8355                        // scale_inplace store, value-exact) — the exact m=1 decode epilogue
8356                        // structure at nrows=t.
8357                        let pair =
8358                            match e.matmul_decode_exact_dual_pre(ffn_gate, ffn_up, zq, zd, t)? {
8359                                Some(((g, gs), (u, us))) => Some((g, gs, u, us)),
8360                                None => None,
8361                            };
8362                        let (gate, gs, up, us) = match pair {
8363                            Some(x4) => x4,
8364                            None => (
8365                                e.matmul_decode_exact_pre(ffn_gate, zq, zd, t)?,
8366                                1.0, // scale already applied inside _pre
8367                                e.matmul_decode_exact_pre(ffn_up, zq, zd, t)?,
8368                                1.0,
8369                            ),
8370                        };
8371                        if e.uses_q8_1_fast(ffn_down) {
8372                            let (aq, ad) = e.silu_mul_scaled_q8_1(&gate, &up, gs, us, t * n_ff)?;
8373                            e.matmul_decode_exact_pre(ffn_down, &aq, &ad, t)?
8374                        } else {
8375                            let mut act = vbuf(e, t * n_ff)?;
8376                            e.silu_mul_scaled(&gate, &up, gs, us, &mut act, t * n_ff)?;
8377                            e.matmul_decode_exact(ffn_down, &act, t)?
8378                        }
8379                    } else {
8380                        // UNFUSED (pre-fix) chain — MoE-adjacent/M3/off-fast layers, unchanged.
8381                        // DUAL gate+up batched twin (lane/verify-economics, 2026-08-02): one launch
8382                        // for the pair at t=2..8 — bit-identical per (tensor,token,row) to the two
8383                        // singles (kernel-check pins bitwise; MEMRA_SPEC_DUAL_T=0 reverts). None
8384                        // (non-NVFP4 / t outside the tier / seam off) -> the two singles, unchanged.
8385                        let (gate, up) =
8386                            match e.matmul_decode_exact_dual(ffn_gate, ffn_up, &z, t)? {
8387                                Some(pair) => pair,
8388                                None => (
8389                                    e.matmul_decode_exact(ffn_gate, &z, t)?,
8390                                    e.matmul_decode_exact(ffn_up, &z, t)?,
8391                                ),
8392                            };
8393                        let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
8394                        Self::ffn_act_lim(
8395                            e,
8396                            &self.cfg,
8397                            &gate,
8398                            &up,
8399                            1.0,
8400                            1.0,
8401                            dense_lim,
8402                            &mut act,
8403                            t * n_ff,
8404                        )?;
8405                        e.matmul_decode_exact(ffn_down, &act, t)?
8406                    }
8407                }
8408                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
8409            };
8410            // CROSS-LAYER fusion: defer this layer's post-FFN residual add — the next layer's
8411            // fused-q8 attn norm folds it in (add_rms_norm_q8_1 == add; rms_norm; quantize,
8412            // kernel-check-pinned at nrows=T). Non-fused next layers add explicitly above.
8413            pending = Some((x1, ffn_out));
8414        }
8415        // RANGE's final add (no next norm INSIDE the range to fuse with; for the
8416        // whole-trunk call that is the last layer, whose next norm is output_norm — f32-out).
8417        if let Some((x1p, f1p)) = pending.take() {
8418            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
8419            e.add(&x1p, &f1p, &mut x2, t * n_embd)?;
8420            x = x2;
8421        }
8422        Ok(x)
8423    }
8424    /// BATCHED linear-attn verify (T=K+1): the whole layer in ~10 launches instead of T x the
8425    /// T=1 decode chain (T x ~12 launches + T weight reads of the four projections). The GDN
8426    /// recurrence itself is inherently sequential — gdn_scan_s128 runs its internal t-loop with
8427    /// the SAME per-token math as chained T=1 calls (bit-identical state evolution); everything
8428    /// around it (projections, conv, prep, gated norm, out-proj) batches. Advances conv ring +
8429    /// ssm state exactly like T sequential decode steps.
8430    /// `want_stash`: additionally RETAIN the gdn-scan inputs (pure buffer keep-alives, zero extra
8431    /// kernels) so a partial accept can rebuild the state after any column prefix (REPLAY-FREE).
8432    #[allow(clippy::too_many_arguments)]
8433    fn linear_attn_verify_t(
8434        &self,
8435        e: &Engine,
8436        la: &LinearAttnLayer,
8437        h: &CudaSlice<f32>,
8438        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
8439        t: usize,
8440        cache: &mut Cache,
8441        il: usize,
8442        want_stash: bool,
8443    ) -> Result<(CudaSlice<f32>, Option<GdnStash>), Box<dyn std::error::Error>> {
8444        let cfg = &self.cfg;
8445        let geometry = la.geometry;
8446        let d_state = geometry.key_head_dim as usize;
8447        let num_k = geometry.key_heads as usize;
8448        let num_v = geometry.value_heads as usize;
8449        let d_conv = geometry.conv_kernel as usize;
8450        let key_dim = d_state * num_k;
8451        let conv_dim = key_dim * 2 + geometry.value_head_dim as usize * num_v;
8452        let eps = cfg.rms_eps;
8453        let scale = 1.0 / (d_state as f32).sqrt();
8454
8455        // DECODE-EXACT projections: matmul_decode_exact forces the MMVQ (warp-per-row, 32-thread)
8456        // accumulation order for EVERY m, matching the T=1 decode path bit-for-bit. The generic
8457        // `matmul` at m>=5 falls to dp4a (128-thread, two-level reduce) which has a different FP
8458        // sum order — ULP differences propagate through gdn_scan and flip argmax on the 27B.
8459        // Q8 TRUNK-FUSION at T=1 (35B: wqkv+wqkv_gate both Q8_0): one fused2 launch, bit-identical
8460        // per (tensor,row) to the two m=1 MMVQ dispatches below — decode-exact contract holds.
8461        // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T, t=2-4): quantize h ONCE for every
8462        // fused-eligible same-input Q8_0 pair of this layer (35B wqkv+wqkv_gate; 9B
8463        // ssm_beta+ssm_alpha) — each fused2 batched launch then replaces two decode-exact
8464        // calls (each of which re-quantizes the same h + runs its own _b2/_b4 launch).
8465        // Bit-identical per (tensor,token,row) — see spec_fused_t().
8466        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm emitted
8467        // directly as q8_1 by the caller's fused rms_norm_q8_1 (bit-identical to the unfused
8468        // chain, kernel-check-pinned). When present it REPLACES the standalone quantize below
8469        // and feeds every projection; the caller guaranteed all four input projections are
8470        // q8_1-fast. When absent, the old shared-quantize (fused-t window) stands.
8471        let h_q8_t = if h_q8.is_none()
8472            && spec_fused_t()
8473            && (2..=4).contains(&t)
8474            && ((e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate))
8475                || (e.uses_q8_1_fast(&la.ssm_beta) && e.uses_q8_1_fast(&la.ssm_alpha)))
8476        {
8477            Some(e.quantize_q8_1(h, t, cfg.n_embd as usize)?)
8478        } else {
8479            None
8480        };
8481        // one view: the caller's fused-norm q8 or this fn's own shared quantize.
8482        let hq8_any: Option<(&CudaSlice<i8>, &CudaSlice<f32>)> =
8483            h_q8.or(h_q8_t.as_ref().map(|(q, d)| (q, d)));
8484        let (qkv_mixed, z) = {
8485            let mut fused = None;
8486            if t == 1 && e.uses_q8_1_fast(&la.wqkv) && e.uses_q8_1_fast(&la.wqkv_gate) {
8487                let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8488                fused = e.matmul_q8_fused2(&la.wqkv, &la.wqkv_gate, &hq, &hd)?;
8489            } else if let Some((hq, hd)) = hq8_any {
8490                if spec_fused_t() && (2..=4).contains(&t) {
8491                    fused = e.matmul_q8_fused2_t(&la.wqkv, &la.wqkv_gate, hq, hd, t)?;
8492                }
8493            }
8494            match (fused, hq8_any) {
8495                (Some(pair), _) => pair,
8496                (None, Some((hq, hd))) if h_q8.is_some() => (
8497                    e.matmul_decode_exact_pre(&la.wqkv, hq, hd, t)?,
8498                    e.matmul_decode_exact_pre(&la.wqkv_gate, hq, hd, t)?,
8499                ),
8500                (None, _) => (
8501                    e.matmul_decode_exact(&la.wqkv, h, t)?,
8502                    e.matmul_decode_exact(&la.wqkv_gate, h, t)?,
8503                ),
8504            }
8505        };
8506        // beta+alpha DUAL at T=1 (75% of p3 rounds run T=1 verify — p-min chain cuts): the dual
8507        // mr2 kernel is bit-identical per element to the m=1 MMVQ matmul_decode_exact dispatches
8508        // (same warp-per-row body, blockIdx.y picks the weight), so the decode-exact contract
8509        // holds; the run-spec battery is the arbiter. T>1 keeps the per-tensor decode-exact path.
8510        let (beta_raw, alpha) = if t == 1 {
8511            let (hq, hd) = e.quantize_q8_1(h, 1, cfg.n_embd as usize)?;
8512            match e.matmul_pre_dual_noscale(&la.ssm_beta, &la.ssm_alpha, &hq, &hd, 1)? {
8513                Some(((mut b, bs), (mut a, as_))) => {
8514                    if bs != 1.0 {
8515                        e.scale_inplace(&mut b, bs, la.ssm_beta.out_features())?;
8516                    }
8517                    if as_ != 1.0 {
8518                        e.scale_inplace(&mut a, as_, la.ssm_alpha.out_features())?;
8519                    }
8520                    (b, a)
8521                }
8522                // Q8_0 fused2 twin (9B stores beta/alpha as Q8_0): DISPATCH-MIRRORS the eager
8523                // decode's beta_alpha closure — the fused body is qmatvec_q8_0_mmvq verbatim,
8524                // bit-identical per row (kernel-check rel=0.00e0 gate), so decode==verify holds.
8525                None => match e.matmul_q8_fused2(&la.ssm_beta, &la.ssm_alpha, &hq, &hd)? {
8526                    Some((b, a)) => (b, a),
8527                    None => (
8528                        e.matmul_decode_exact(&la.ssm_beta, h, 1)?,
8529                        e.matmul_decode_exact(&la.ssm_alpha, h, 1)?,
8530                    ),
8531                },
8532            }
8533        } else {
8534            // fused-t twin (9B stores beta/alpha as Q8_0): same shared-quantize + one launch
8535            // contract as the wqkv pair above; 35B beta/alpha are Float -> None -> fallback.
8536            let mut nvfp4_fused = None;
8537            let mut q8_fused = None;
8538            if let Some((hq, hd)) = hq8_any {
8539                if t == 3 && std::env::var("MEMRA_NVFP4_AUX_DUAL").as_deref() != Ok("0") {
8540                    nvfp4_fused =
8541                        e.matmul_decode_exact_dual_pre(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8542                    if nvfp4_fused.is_some() && std::env::var("MEMRA_DEBUG").is_ok() {
8543                        static ONCE: std::sync::Once = std::sync::Once::new();
8544                        ONCE.call_once(|| {
8545                            eprintln!("[memra] NVFP4 beta+alpha batched aux dual ENGAGED (t={t})")
8546                        });
8547                    }
8548                }
8549                if nvfp4_fused.is_none() && spec_fused_t() && (2..=4).contains(&t) {
8550                    q8_fused = e.matmul_q8_fused2_t(&la.ssm_beta, &la.ssm_alpha, hq, hd, t)?;
8551                }
8552            }
8553            if let Some(((mut b, bs), (mut a, as_))) = nvfp4_fused {
8554                if bs != 1.0 {
8555                    e.scale_inplace(&mut b, bs, t * la.ssm_beta.out_features())?;
8556                }
8557                if as_ != 1.0 {
8558                    e.scale_inplace(&mut a, as_, t * la.ssm_alpha.out_features())?;
8559                }
8560                (b, a)
8561            } else if let Some(pair) = q8_fused {
8562                pair
8563            } else {
8564                match hq8_any {
8565                    Some((hq, hd)) if h_q8.is_some() => (
8566                        e.matmul_decode_exact_pre(&la.ssm_beta, hq, hd, t)?,
8567                        e.matmul_decode_exact_pre(&la.ssm_alpha, hq, hd, t)?,
8568                    ),
8569                    _ => (
8570                        e.matmul_decode_exact(&la.ssm_beta, h, t)?,
8571                        e.matmul_decode_exact(&la.ssm_alpha, h, t)?,
8572                    ),
8573                }
8574            }
8575        };
8576
8577        // conv with CARRIED state + ring roll (T >= pad rides the input-column update kernel;
8578        // T < pad — the MEMRA_SPEC_M2 t=2 arm — rolls via the pure-copy ring rebuild).
8579        let rl = cache.recur[il].as_mut().unwrap();
8580        let mut conv_out = e.uninit(conv_dim * t)?;
8581        e.ssm_conv1d_tm_state(
8582            &qkv_mixed,
8583            &mut rl.conv_state,
8584            la.ssm_conv1d.float_data(),
8585            &mut conv_out,
8586            conv_dim,
8587            t,
8588            d_conv,
8589        )?;
8590
8591        // GDN prep via the prefill kernels (repack + L2 + sigmoid + glog), T-wide.
8592        let mut q_g = e.uninit(d_state * num_v * t)?;
8593        let mut k_g = e.uninit(d_state * num_v * t)?;
8594        let mut v_g = e.uninit(d_state * num_v * t)?;
8595        e.qkv_to_gdn_repack(
8596            &conv_out, &mut q_g, &mut k_g, &mut v_g, d_state, num_v, num_k, key_dim, t,
8597        )?;
8598        let mut q_l2 = e.uninit(d_state * num_v * t)?;
8599        e.l2_norm_decode(&q_g, &mut q_l2, d_state, num_v * t, eps)?;
8600        let mut k_l2 = e.uninit(d_state * num_v * t)?;
8601        e.l2_norm_decode(&k_g, &mut k_l2, d_state, num_v * t, eps)?;
8602        let mut beta = e.uninit(t * num_v)?;
8603        e.sigmoid(&beta_raw, &mut beta, t * num_v)?;
8604        let mut g_log = e.uninit(t * num_v)?;
8605        e.gdn_glog(
8606            &alpha,
8607            la.ssm_dt.float_data(),
8608            la.ssm_a.float_data(),
8609            &mut g_log,
8610            num_v,
8611            t,
8612        )?;
8613
8614        // ONE gdn_scan over T tokens from the carried state (internal sequential loop ==
8615        // T chained T=1 steps). Ping-pong the resident buffers like eager decode.
8616        let mut o = e.uninit(d_state * num_v * t)?;
8617        {
8618            let crate::cache::RecurLayer {
8619                ssm_state,
8620                ssm_state_alt,
8621                ..
8622            } = rl;
8623            e.gdn_scan_s128(
8624                &q_l2,
8625                &k_l2,
8626                &v_g,
8627                &g_log,
8628                &beta,
8629                ssm_state,
8630                ssm_state_alt,
8631                &mut o,
8632                num_v,
8633                t,
8634                scale,
8635            )?;
8636        }
8637        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
8638
8639        // gated RMSNorm + out projection, T-wide. FUSED-QUANTIZE ARM (lane/vt-fixes fix 2,
8640        // mirroring the T=1 decode's launch-arc form): when ssm_out rides the q8_1 fast path,
8641        // emit q8_1 straight from the gated norm at nrows=num_v*t (row-indexed kernel, the
8642        // T-wide launch is the per-row program; kernel-check pins bit-identity vs
8643        // gated_rmsnorm -> quantize_q8_1 at T=1 and T=5) and feed the decode-exact dispatch
8644        // pre-quantized — one launch replaces norm + quantize. Fallback = the f32 chain.
8645        let out = if e.uses_q8_1_fast(&la.ssm_out) {
8646            let (gq, gd) =
8647                e.gated_rmsnorm_q8_1(&o, la.ssm_norm.float_data(), &z, d_state, num_v * t, eps)?;
8648            e.matmul_decode_exact_pre(&la.ssm_out, &gq, &gd, t)?
8649        } else {
8650            let mut gn = e.uninit(d_state * num_v * t)?;
8651            e.gated_rmsnorm(
8652                &o,
8653                la.ssm_norm.float_data(),
8654                &z,
8655                &mut gn,
8656                d_state,
8657                num_v * t,
8658                eps,
8659            )?;
8660            // DECODE-EXACT out-projection: same MMVQ path as the T=1 decode (ssm_out at m>=5
8661            // would fall to dp4a with a different FP reduction order — same class of bug as
8662            // the input projs).
8663            e.matmul_decode_exact(&la.ssm_out, &gn, t)?
8664        };
8665        let stash = if want_stash {
8666            Some(GdnStash {
8667                qkv_mixed,
8668                q_l2,
8669                k_l2,
8670                v_g,
8671                g_log,
8672                beta,
8673            })
8674        } else {
8675            None
8676        };
8677        Ok((out, stash))
8678    }
8679
8680    /// REPLAY-FREE partial-accept commit (2026-07-03): make the cache state == "committed through
8681    /// the first `j` verify columns" WITHOUT the legacy rollback + duplicate trunk replay.
8682    /// - Full-attn KV: truncate both the owning-stage shadow and every TP rank to snapshot + j.
8683    ///   The verify's appended rows for those columns are bit-identical to what an eager T=1
8684    ///   chain writes (the decode-exact contract the verify-probe gates), so keeping them ==
8685    ///   replaying them.
8686    /// - Linear layers, batched path: rebuild the conv ring by PURE COPIES (ring holds raw input
8687    ///   columns) and the ssm state by a prefix re-run of the SAME gdn_scan kernel (t=j) from the
8688    ///   snapshot state over the stash's identical inputs — the kernel's t-loop carries state in
8689    ///   registers and writes it once at the end, so iterations 0..j-1 are independent of T:
8690    ///   bit-identical to the verify's own state after j tokens == the eager chain state.
8691    /// - Linear layers, per-column path: restore the cloned actual state after column j-1.
8692    /// Caller guarantees 1 <= j <= t-1 (j==0 rounds take the legacy rollback; j==t is full accept).
8693    fn commit_verified_prefix(
8694        &self,
8695        e: &Engine,
8696        cache: &mut Cache,
8697        snap: &crate::cache::CacheSnapshot,
8698        ckpt: &VerifyCkpt,
8699        j: usize,
8700        kv_lens_done: bool,
8701        dev_j: Option<(&CudaSlice<u32>, usize, usize)>,
8702    ) -> Result<(), Box<dyn std::error::Error>> {
8703        // GDN geometry derives lazily inside recurrent-layer arms. Full-attention plans carry no
8704        // recurrent state and must never be forced through a synthetic SSM geometry.
8705        // Engine-bundle slice 1 (DSF-ROUNDCOST-20260820 §1.1): the per-column-arm restores
8706        // are 2 tiny D2D copies per linear layer (~96 dispatches/partial round on the q38
8707        // route). When every cols-arm layer shares uniform state sizes (single ssm cfg —
8708        // always true today), batch them into two `copy_batch_uniform_f32` launches. Bytes,
8709        // buffers and stream order are identical to the per-layer memcpy sequence; the
8710        // kernel-rebuild (gdn-stash) arm below is untouched. MEMRA_STATE_COPY_BATCH=0 reverts.
8711        let mut batched_cols = false;
8712        if state_copy_batch_on() && dev_j.is_none() {
8713            use cudarc::driver::DevicePtr;
8714            let s = &e.gpu.stream();
8715            let mut conv_pairs: Vec<(u64, u64)> = Vec::new();
8716            let mut ssm_pairs: Vec<(u64, u64)> = Vec::new();
8717            let (mut conv_words, mut ssm_words) = (0usize, 0usize);
8718            let mut uniform = true;
8719            for il in 0..self.layers.len() {
8720                let Some(rl) = cache.recur[il].as_ref() else {
8721                    continue;
8722                };
8723                if ckpt.gdn[il].is_some() {
8724                    continue; // kernel-rebuild arm restores below, per layer
8725                }
8726                let Some(cols) = &ckpt.cols[il] else {
8727                    continue; // missing-ckpt error surfaces in the main loop
8728                };
8729                let (c, st) = &cols[j - 1];
8730                if conv_pairs.is_empty() {
8731                    conv_words = c.len();
8732                    ssm_words = st.len();
8733                } else if c.len() != conv_words || st.len() != ssm_words {
8734                    uniform = false;
8735                    break;
8736                }
8737                let (pc, _g0) = c.device_ptr(s);
8738                let (dc, _g1) = rl.conv_state.device_ptr(s);
8739                let (ps, _g2) = st.device_ptr(s);
8740                let (ds, _g3) = rl.ssm_state.device_ptr(s);
8741                conv_pairs.push((pc as u64, dc as u64));
8742                ssm_pairs.push((ps as u64, ds as u64));
8743            }
8744            if uniform && !conv_pairs.is_empty() {
8745                let n = conv_pairs.len();
8746                let mut t = vec![0u64; 2 * n];
8747                for (k, &(src, dst)) in conv_pairs.iter().enumerate() {
8748                    t[k] = src;
8749                    t[n + k] = dst;
8750                }
8751                let conv_t = e.htod_u64(&t)?;
8752                for (k, &(src, dst)) in ssm_pairs.iter().enumerate() {
8753                    t[k] = src;
8754                    t[n + k] = dst;
8755                }
8756                let ssm_t = e.htod_u64(&t)?;
8757                e.copy_batch_uniform_f32(&conv_t, n, conv_words)?;
8758                e.copy_batch_uniform_f32(&ssm_t, n, ssm_words)?;
8759                batched_cols = true;
8760            }
8761        }
8762        rewind_tp_kv_verified_prefix(&mut cache.tp_kv, &snap.tp_kv_len, j)?;
8763        for il in 0..self.layers.len() {
8764            if let (Some(kvl), Some(saved)) = (cache.kv[il].as_mut(), snap.kv_len[il]) {
8765                kvl.len = saved + j;
8766                // devacc 3a: spec_rollback_kv already wrote len_d on-device (same value).
8767                if !kv_lens_done {
8768                    e.set_i32_one(&mut kvl.len_d, kvl.len as i32)?;
8769                }
8770            }
8771            if let Some(rl) = cache.recur[il].as_mut() {
8772                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8773                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8774                };
8775                let geometry = linear.geometry;
8776                let d_state = geometry.key_head_dim as usize;
8777                let num_k = geometry.key_heads as usize;
8778                let num_v = geometry.value_heads as usize;
8779                let d_conv = geometry.conv_kernel as usize;
8780                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8781                let scale = 1.0 / (d_state as f32).sqrt();
8782                if let Some(st) = &ckpt.gdn[il] {
8783                    let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8784                    let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8785                    if let Some((acc, base, t_v)) = dev_j {
8786                        // 3b: j read on-device (_dc twins, same bodies; full accept early-exits).
8787                        e.ssm_conv_ring_rebuild_dc(
8788                            &st.qkv_mixed,
8789                            ring_old,
8790                            &mut rl.conv_state,
8791                            conv_dim,
8792                            acc,
8793                            base,
8794                            t_v,
8795                            d_conv,
8796                        )?;
8797                        let mut o = e.uninit(d_state * num_v * j.max(1))?;
8798                        e.gdn_scan_s128_dc(
8799                            &st.q_l2,
8800                            &st.k_l2,
8801                            &st.v_g,
8802                            &st.g_log,
8803                            &st.beta,
8804                            state_in,
8805                            &mut rl.ssm_state,
8806                            &mut o,
8807                            num_v,
8808                            acc,
8809                            base,
8810                            t_v,
8811                            scale,
8812                        )?;
8813                    } else {
8814                        e.ssm_conv_ring_rebuild(
8815                            &st.qkv_mixed,
8816                            ring_old,
8817                            &mut rl.conv_state,
8818                            conv_dim,
8819                            j,
8820                            d_conv,
8821                        )?;
8822                        let mut o = e.uninit(d_state * num_v * j)?; // scan output, discarded
8823                        e.gdn_scan_s128(
8824                            &st.q_l2,
8825                            &st.k_l2,
8826                            &st.v_g,
8827                            &st.g_log,
8828                            &st.beta,
8829                            state_in,
8830                            &mut rl.ssm_state,
8831                            &mut o,
8832                            num_v,
8833                            j,
8834                            scale,
8835                        )?;
8836                    }
8837                } else if let Some(cols) = &ckpt.cols[il] {
8838                    if !batched_cols {
8839                        let (c, s) = &cols[j - 1];
8840                        e.copy_into(&mut rl.conv_state, 0, c, c.len())?;
8841                        e.copy_into(&mut rl.ssm_state, 0, s, s.len())?;
8842                    }
8843                } else {
8844                    return Err(
8845                        "commit_verified_prefix: verify ckpt missing for linear layer".into(),
8846                    );
8847                }
8848            }
8849        }
8850        cache.pos = snap.pos + j;
8851        Ok(())
8852    }
8853
8854    /// ROUND-STREAM: recur restore with device-j (the _dc twins; full accept early-exits
8855    /// in-kernel). Requires the batched-linear stash on every linear layer (stream gate).
8856    fn commit_verified_prefix_stream(
8857        &self,
8858        e: &Engine,
8859        cache: &mut Cache,
8860        snap: &crate::cache::CacheSnapshot,
8861        ckpt: &VerifyCkpt,
8862        acc: &CudaSlice<u32>,
8863        base: usize,
8864        t_v: usize,
8865    ) -> Result<(), Box<dyn std::error::Error>> {
8866        for il in 0..self.layers.len() {
8867            if let Some(rl) = cache.recur[il].as_mut() {
8868                let Mixer::Linear(linear) = &self.layers[il].mixer else {
8869                    return Err(format!("recurrent cache layer {il} has no GDN plan").into());
8870                };
8871                let geometry = linear.geometry;
8872                let d_state = geometry.key_head_dim as usize;
8873                let num_k = geometry.key_heads as usize;
8874                let num_v = geometry.value_heads as usize;
8875                let d_conv = geometry.conv_kernel as usize;
8876                let conv_dim = d_state * num_k * 2 + geometry.value_head_dim as usize * num_v;
8877                let scale = 1.0 / (d_state as f32).sqrt();
8878                let st = ckpt.gdn[il]
8879                    .as_ref()
8880                    .ok_or("stream restore: batched-linear stash missing")?;
8881                let ring_old = snap.conv[il].as_ref().expect("snapshot missing conv");
8882                let state_in = snap.ssm[il].as_ref().expect("snapshot missing ssm");
8883                e.ssm_conv_ring_rebuild_dc(
8884                    &st.qkv_mixed,
8885                    ring_old,
8886                    &mut rl.conv_state,
8887                    conv_dim,
8888                    acc,
8889                    base,
8890                    t_v,
8891                    d_conv,
8892                )?;
8893                let mut o = e.uninit(d_state * num_v * t_v)?;
8894                e.gdn_scan_s128_dc(
8895                    &st.q_l2,
8896                    &st.k_l2,
8897                    &st.v_g,
8898                    &st.g_log,
8899                    &st.beta,
8900                    state_in,
8901                    &mut rl.ssm_state,
8902                    &mut o,
8903                    num_v,
8904                    acc,
8905                    base,
8906                    t_v,
8907                    scale,
8908                )?;
8909            }
8910        }
8911        Ok(())
8912    }
8913
8914    /// EAGLE3 aux-capturing verify forward over `tokens` (T) — mirrors `decode_step_t_h` exactly
8915    /// (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual-
8916    /// stream hiddens (blocks in `aux_layers`) for TWO columns: the LAST column (always) and the
8917    /// optional `pred_col` (the EAGLE seed = bonus's predecessor). Returns
8918    /// (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator's commit.
8919    pub fn decode_step_t_aux2(
8920        &self,
8921        e: &Engine,
8922        tokens: &[u32],
8923        pos0: usize,
8924        cache: &mut Cache,
8925        aux_layers: &[usize],
8926        pred_col: Option<usize>,
8927    ) -> Result<
8928        (Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>),
8929        Box<dyn std::error::Error>,
8930    > {
8931        let cfg = &self.cfg;
8932        let n_embd = cfg.n_embd as usize;
8933        let eps = cfg.rms_eps;
8934        let t = tokens.len();
8935        let pos_vec: Vec<i32> = (0..t).map(|i| (pos0 + i) as i32).collect();
8936        let pos_d = e.htod_i32(&pos_vec)?;
8937        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
8938        let mut aux_last: Vec<CudaSlice<f32>> = Vec::with_capacity(aux_layers.len());
8939        let mut aux_pred: Vec<CudaSlice<f32>> = Vec::new();
8940        let want_pred = pred_col.is_some();
8941
8942        for (il, layer) in self.layers.iter().enumerate() {
8943            // DISPATCH-MIRRORED norms (FP-order lesson #8) — see decode_step_t_h_emb.
8944            let mixer_fast = self.mixer_in_q8_1_fast(e, &layer.mixer);
8945            let norm_fused = std::env::var("MEMRA_NO_FUSE_NORMQ").is_err() && mixer_fast;
8946            let mut h = vbuf(e, t * n_embd)?; // fully written by either rms_norm arm
8947            if norm_fused {
8948                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8949            } else {
8950                e.rms_norm(&x, layer.attn_norm.float_data(), &mut h, n_embd, t, eps)?;
8951            }
8952            let mixed = match &layer.mixer {
8953                Mixer::Full(fa) => {
8954                    self.full_attn_verify(e, fa, &h, None, &pos_d, t, cache, il, None)?
8955                }
8956                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
8957                Mixer::Linear(la) => {
8958                    let mut out = e.zeros(t * n_embd)?;
8959                    for col in 0..t {
8960                        let mut h_col = e.zeros(n_embd)?;
8961                        let src = h.slice(col * n_embd..(col + 1) * n_embd);
8962                        e.copy_view_into(&mut h_col, 0, &src, n_embd)?;
8963                        let m_col = self.linear_attn_decode(e, la, &h_col, cache, il)?;
8964                        e.copy_into(&mut out, col * n_embd, &m_col, n_embd)?;
8965                    }
8966                    out
8967                }
8968            };
8969            let ffn_fuse = match &layer.ffn {
8970                crate::hybrid::Ffn::Dense {
8971                    ffn_gate, ffn_up, ..
8972                } => {
8973                    std::env::var("MEMRA_NO_FUSE_NORMQ").is_err()
8974                        && e.uses_q8_1_fast(ffn_gate)
8975                        && e.uses_q8_1_fast(ffn_up)
8976                }
8977                crate::hybrid::Ffn::Moe(_) => false,
8978            };
8979            let mut x1 = vbuf(e, t * n_embd)?; // fully written by add / add_rms_norm
8980            let mut z = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode / add_rms_norm
8981            if ffn_fuse {
8982                e.add(&x, &mixed, &mut x1, t * n_embd)?;
8983                e.rms_norm_decode(
8984                    &x1,
8985                    layer.post_attn_norm.float_data(),
8986                    &mut z,
8987                    n_embd,
8988                    t,
8989                    eps,
8990                )?;
8991            } else {
8992                e.add_rms_norm(
8993                    &x,
8994                    &mixed,
8995                    layer.post_attn_norm.float_data(),
8996                    &mut x1,
8997                    &mut z,
8998                    n_embd,
8999                    t,
9000                    eps,
9001                )?;
9002            }
9003            let ffn_out = match &layer.ffn {
9004                crate::hybrid::Ffn::Dense {
9005                    ffn_gate,
9006                    ffn_up,
9007                    ffn_down,
9008                } => {
9009                    let n_ff = ffn_gate.out_features();
9010                    let gate = e.matmul_decode_exact(ffn_gate, &z, t)?;
9011                    let up = e.matmul_decode_exact(ffn_up, &z, t)?;
9012                    let mut act = vbuf(e, t * n_ff)?; // fully written by ffn_act_lim
9013                    // dense FFN clamp = the SHEXP array (upstream build_ffn serves both).
9014                    Self::ffn_act_lim(
9015                        e,
9016                        &self.cfg,
9017                        &gate,
9018                        &up,
9019                        1.0,
9020                        1.0,
9021                        self.cfg.clamp_shexp_at(il as u32),
9022                        &mut act,
9023                        t * n_ff,
9024                    )?;
9025                    e.matmul_decode_exact(ffn_down, &act, t)?
9026                }
9027                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il(e, m, &z, t, il as u16)?,
9028            };
9029            let mut x2 = vbuf(e, t * n_embd)?; // fully written by add
9030            e.add(&x1, &ffn_out, &mut x2, t * n_embd)?;
9031            if aux_layers.contains(&il) {
9032                let mut a = e.zeros(n_embd)?;
9033                e.copy_view_into(&mut a, 0, &x2.slice((t - 1) * n_embd..t * n_embd), n_embd)?;
9034                aux_last.push(a);
9035                if let Some(pc) = pred_col {
9036                    let mut ap = e.zeros(n_embd)?;
9037                    e.copy_view_into(
9038                        &mut ap,
9039                        0,
9040                        &x2.slice(pc * n_embd..(pc + 1) * n_embd),
9041                        n_embd,
9042                    )?;
9043                    aux_pred.push(ap);
9044                }
9045            }
9046            x = x2;
9047        }
9048        let mut hn = vbuf(e, t * n_embd)?; // fully written by rms_norm_decode
9049        e.rms_norm_decode(&x, self.output_norm.float_data(), &mut hn, n_embd, t, eps)?;
9050        let logits = e.matmul_decode_exact(&self.output, &hn, t)?;
9051        let host = e.dtoh(&logits)?;
9052        cache.pos += t;
9053        Ok((
9054            host,
9055            aux_last,
9056            if want_pred { Some(aux_pred) } else { None },
9057        ))
9058    }
9059
9060    /// step35 SPEC-VERIFY attention over T query tokens — a per-row REPLAY of the eager
9061    /// `step35_decode_attn`.
9062    ///
9063    /// WHY A REPLAY AND NOT A BATCHED TWIN. The verify's whole job is to be bit-identical to what
9064    /// the eager decode would have computed for the same tokens; that is what makes greedy spec
9065    /// decode exact (run-spec asserts token identity for K=1..8). Every other verify arm in this
9066    /// file earns that identity by carefully mirroring dispatch (`matmul_decode_exact` to force
9067    /// MMVQ at any m, per-layer `ffn_fuse` mirroring, per-row `fa_decode` key bounds). step35
9068    /// stacks FOUR more per-layer degrees of freedom on top of that — per-layer `n_head`
9069    /// (64 full / 96 SWA), per-layer rotary width (64 full / 128 SWA), per-layer rope base, and a
9070    /// SEPARATE `attn_gate` tensor whose projection shares the attn-normed input — and its SWA
9071    /// layers attend through a token-OFFSET view whose offset is a function of the ABSOLUTE
9072    /// position of each query row. A batched twin would have to reproduce all of that AND the
9073    /// per-row offset in one launch; the offset alone rules out the existing rows kernels (they
9074    /// take one `base_len`, not a per-row offset).
9075    ///
9076    /// So this arm calls the eager path itself, once per row, on the same cache. Identity is then
9077    /// true BY CONSTRUCTION rather than by mirroring: row r runs exactly the kernel sequence that
9078    /// eager decode step r runs (same projections, same q8_1 fusion decision, same append, same
9079    /// view arithmetic, same `fa_decode_kvmod`, same gate), because it IS that code. Cost: T x the
9080    /// eager decode mixer instead of one batched pass — the same trade the generic arm's `else`
9081    /// per-row loop already accepts when `fa_rows_eligible` says no. Correctness first; a batched
9082    /// step35 twin is a perf lane's job and must be gated against this arm.
9083    ///
9084    /// The `h_q8` pre-quantized pair from the caller's fused norm is NOT forwarded: it is a
9085    /// T-row buffer and `step35_decode_attn`'s `pre_q` contract is one row. Instead each row's
9086    /// f32 `h` slice is handed over and the callee re-derives its own q8_1 exactly as eager decode
9087    /// does (`quantize_q8_1(h, 1, n_embd)`) — which is the dispatch being mirrored. Callers that
9088    /// took the fused arm therefore MUST still pass a live `h`; `step35_verify` asserts that.
9089    #[allow(clippy::too_many_arguments)]
9090    fn step35_verify(
9091        &self,
9092        e: &Engine,
9093        fa: &FullAttnLayer,
9094        h: &CudaSlice<f32>,
9095        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9096        t: usize,
9097        cache: &mut Cache,
9098        il: usize,
9099    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9100        let n_embd = self.cfg.n_embd as usize;
9101        // The fused (h-less) attn-norm arm hands `h` as a zero-length placeholder. This arm needs
9102        // the f32 rows, so the caller must not take that lever for step35 — enforced at the call
9103        // site by the sliding-gated-MoE `Mixer::Full(_) => false` arm of
9104        // `lin_q8_only` in `decode_step_t_core_stream`, and asserted here so a future caller
9105        // cannot regress it into silently reading an empty buffer.
9106        assert_eq!(
9107            h.len(),
9108            t * n_embd,
9109            "step35_verify needs the f32 attn-normed rows ([t*n_embd]); the caller took the \
9110             fused q8-only norm arm (h_q8={}) — step35 must stay on the unfused arm",
9111            h_q8.is_some()
9112        );
9113        // ROW WIDTH IS n_embd, NOT n_head*head_dim: `step35_decode_attn` returns the mixer output
9114        // AFTER `wo`, so a row is [n_embd] — the same contract the generic arm's
9115        // `matmul_decode_exact(&fa.wo, &attn_g, t)` return has. Sizing this buffer from the
9116        // per-layer head geometry (8192 on full-attn, 12288 on SWA) instead overran the row on the
9117        // FIRST copy and panicked inside `copy_into`'s `CudaView::slice` unwrap
9118        // (raw/mtp-bt-20260806T212127Z.log frames 12-13).
9119        let mut out = vbuf(e, t * n_embd)?; // each row fully written by the copy below
9120        for r in 0..t {
9121            // Absolute position of this query row. `cache.pos` is the committed length at round
9122            // start and every row before r has already been appended by this loop, so the r-th
9123            // verify token sits at cache.pos + r — the same position eager decode would give it.
9124            let pos_d = e.htod_i32(&[(cache.pos + r) as i32])?;
9125            let mut h_row = vbuf(e, n_embd)?; // fully written by copy_view_into
9126            e.copy_view_into(
9127                &mut h_row,
9128                0,
9129                &h.slice(r * n_embd..(r + 1) * n_embd),
9130                n_embd,
9131            )?;
9132            // THE eager decode mixer: appends this row's K/V at kvl.len, advances it, then
9133            // attends over the (SWA-offset) view. Post-`wo`, same contract as this fn returns.
9134            let o = self.step35_decode_attn(e, fa, il, &h_row, None, &pos_d, cache)?;
9135            debug_assert_eq!(
9136                o.len(),
9137                n_embd,
9138                "step35_decode_attn returns post-wo [n_embd]"
9139            );
9140            e.copy_into(&mut out, r * n_embd, &o, n_embd)?;
9141        }
9142        Ok(out)
9143    }
9144
9145    /// Full-attention mixer over T query tokens with a GROWING resident KV (verify path, §D.3).
9146    /// Appends the T new K/V columns to cache.kv[il] then attends causally over [0..len) via
9147    /// fa_prefill. Token-major [T, kv_dim] projection layout == cache row layout (single copy).
9148    #[allow(clippy::too_many_arguments)]
9149    fn full_attn_verify(
9150        &self,
9151        e: &Engine,
9152        fa: &FullAttnLayer,
9153        h: &CudaSlice<f32>,
9154        h_q8: Option<(&CudaSlice<i8>, &CudaSlice<f32>)>,
9155        pos_d: &CudaSlice<i32>,
9156        t: usize,
9157        cache: &mut Cache,
9158        il: usize,
9159        stream_ctr: Option<&CudaSlice<i32>>,
9160    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
9161        // step35: the generic geometry below is wrong for this arch (per-layer n_head, partial
9162        // per-layer rope, the SWA offset view, and a SEPARATE head-wise gate tensor), so it takes
9163        // its own arm. A verify that silently computes different attention than decode defeats the
9164        // whole self-consistency gate, so the arm is a per-row REPLAY of `step35_decode_attn`
9165        // rather than a batched twin — see `step35_verify` for why that is the exactness-correct
9166        // shape and not laziness.
9167        if self.sliding_gated_moe_batch_program() {
9168            if stream_ctr.is_some() {
9169                return Err(
9170                    "step35 has no ROUND-STREAM verify arm (the device-counter _dc twins \
9171                            cannot express the SWA offset KV view; same root cause as the dc \
9172                            decode refusal) — run spec without the stream arm"
9173                        .into(),
9174                );
9175            }
9176            return self.step35_verify(e, fa, h, h_q8, t, cache, il);
9177        }
9178        let cfg = &self.cfg;
9179        let geometry = cfg.full_attention_geometry_at(il as u32);
9180        let n_head = geometry.n_head as usize;
9181        let n_head_kv = geometry.n_head_kv as usize;
9182        let head_dim = geometry.head_dim_k as usize;
9183        let eps = cfg.rms_eps;
9184        let scale = geometry.attention_scale();
9185        let n_embd = cfg.n_embd as usize;
9186
9187        // DECODE-EXACT Q/K/V projections: matmul_decode_exact forces the MMVQ (warp-per-row) path
9188        // for every m, matching the T=1 decode's FP accumulation order. matmul_pre at m>=5 would
9189        // fall to dp4a (128-thread, two-level reduce) with a different FP sum order.
9190        // Q8 TRUNK-FUSION at T=1: DISPATCH-MIRRORS the eager decode's fused3 (bit-identical body).
9191        // BATCHED EPILOGUE RE-FUSE (lane/vt-fixes fix 2): `h_q8` = the attn-input norm's q8_1
9192        // form emitted by the fused rms_norm_q8_1 (bit-identical to rms_norm_decode ->
9193        // quantize_q8_1, kernel-check-pinned). When present (caller checked mixer q8_1-fast),
9194        // every projection consumes it — `h` may be a zero-len placeholder and must not be read.
9195        let (qf, mut k, v) = {
9196            let mut fused = None;
9197            let qkv_fast =
9198                e.uses_q8_1_fast(&fa.wq) && e.uses_q8_1_fast(&fa.wk) && e.uses_q8_1_fast(&fa.wv);
9199            if t == 1 && qkv_fast {
9200                let (hq_o, hd_o);
9201                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9202                    Some(p) => p,
9203                    None => {
9204                        (hq_o, hd_o) = e.quantize_q8_1(h, 1, n_embd)?;
9205                        (&hq_o, &hd_o)
9206                    }
9207                };
9208                fused = e.matmul_q8_fused3(&fa.wq, &fa.wk, &fa.wv, hq, hd)?;
9209            } else if spec_fused_t() && (2..=4).contains(&t) && qkv_fast {
9210                // VERIFY-TIER TRUNK FUSION (MEMRA_SPEC_FUSED_T): one shared quantize + one
9211                // fused3 batched launch replaces three decode-exact calls (3 re-quantizes of
9212                // the same h + 3 _b2/_b4 launches). Bit-identical per (tensor,token,row).
9213                let (hq_o, hd_o);
9214                let (hq, hd): (&CudaSlice<i8>, &CudaSlice<f32>) = match h_q8 {
9215                    Some(p) => p,
9216                    None => {
9217                        (hq_o, hd_o) = e.quantize_q8_1(h, t, n_embd)?;
9218                        (&hq_o, &hd_o)
9219                    }
9220                };
9221                fused = e.matmul_q8_fused3_t(&fa.wq, &fa.wk, &fa.wv, hq, hd, t)?;
9222            }
9223            match (fused, h_q8) {
9224                (Some(triple), _) => triple,
9225                // shared pre-quantized activation (q8_1-fast guaranteed by the caller): the
9226                // decode-exact dispatch consumes (hq, hd) instead of re-quantizing 3x.
9227                (None, Some((hq, hd))) if qkv_fast => (
9228                    e.matmul_decode_exact_pre(&fa.wq, hq, hd, t)?,
9229                    e.matmul_decode_exact_pre(&fa.wk, hq, hd, t)?,
9230                    e.matmul_decode_exact_pre(&fa.wv, hq, hd, t)?,
9231                ),
9232                (None, _) => (
9233                    e.matmul_decode_exact(&fa.wq, h, t)?,
9234                    e.matmul_decode_exact(&fa.wk, h, t)?,
9235                    e.matmul_decode_exact(&fa.wv, h, t)?,
9236                ),
9237            }
9238        };
9239        // M3/Hy3 have no attention output gate — wq out is exactly q; skip the split.
9240        let gated = geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
9241        let (mut q, gate) = if gated {
9242            let mut q = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9243            let mut gate = vbuf(e, t * n_head * head_dim)?; // fully written by q_gate_split
9244            e.q_gate_split(&qf, &mut q, &mut gate, head_dim, n_head, t)?;
9245            (q, Some(gate))
9246        } else {
9247            (qf, None)
9248        };
9249
9250        let mut qn = vbuf(e, t * n_head * head_dim)?; // fully written by rms_norm
9251        e.rms_norm(
9252            &q,
9253            fa.q_norm.float_data(),
9254            &mut qn,
9255            head_dim,
9256            n_head * t,
9257            eps,
9258        )?;
9259        q = qn;
9260        let mut kn = vbuf(e, t * n_head_kv * head_dim)?; // fully written by rms_norm
9261        e.rms_norm(
9262            &k,
9263            fa.k_norm.float_data(),
9264            &mut kn,
9265            head_dim,
9266            n_head_kv * t,
9267            eps,
9268        )?;
9269        k = kn;
9270        let rope_dims = geometry.n_rot as usize;
9271        e.rope_neox(
9272            &mut q,
9273            pos_d,
9274            head_dim,
9275            rope_dims,
9276            n_head,
9277            t,
9278            geometry.rope_base,
9279            1.0,
9280        )?;
9281        e.rope_neox(
9282            &mut k,
9283            pos_d,
9284            head_dim,
9285            rope_dims,
9286            n_head_kv,
9287            t,
9288            geometry.rope_base,
9289            1.0,
9290        )?;
9291
9292        // append T new K/V columns to the resident QUANTIZED cache. k/v are token-major [T, kv_dim]
9293        // f32; append-quantize each of the T token rows into the byte cache (q8_0 K / q5_1 V).
9294        let kvl = cache.kv[il].as_mut().unwrap();
9295        let (kv_dim_k, kv_dim_v, ktb, vtb) =
9296            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes);
9297        if let Some(ctr) = stream_ctr {
9298            // stream: ONE batched append at the device counter (rows kernel = the per-view warp
9299            // math on a (block, token) grid, documented byte-identical); host len is a stale
9300            // LOWER BOUND under pre-issue (drain reconciles it).
9301            e.append_kv_quantized_rows_dc(
9302                &k,
9303                &v,
9304                &mut kvl.k,
9305                &mut kvl.v,
9306                ctr,
9307                t,
9308                kv_dim_k,
9309                kv_dim_v,
9310                ktb,
9311                vtb,
9312                crate::Engine::kv_fp8_on(),
9313            )?;
9314        } else {
9315            for i in 0..t {
9316                let k_row = k.slice(i * kv_dim_k..(i + 1) * kv_dim_k);
9317                let v_row = v.slice(i * kv_dim_v..(i + 1) * kv_dim_v);
9318                e.append_kv_quantized_view(
9319                    &k_row,
9320                    &v_row,
9321                    &mut kvl.k,
9322                    &mut kvl.v,
9323                    kvl.len + i,
9324                    kv_dim_k,
9325                    kv_dim_v,
9326                    ktb,
9327                    vtb,
9328                    crate::Engine::kv_fp8_on(),
9329                )?;
9330            }
9331            kvl.len += t;
9332        }
9333
9334        // BIT-IDENTICAL VERIFY ATTENTION (spec-exactness fix): the FP accumulation order must be
9335        // byte-for-byte identical to the eager decode path. fa_prefill uses a different tile size
9336        // (BLOCK_Q=64, BK=32) and online-softmax structure than fa_decode's split-K + combine,
9337        // which changes FP summation order and can flip argmax at tight logit margins. Query row r
9338        // attends to keys [0..base_len+r+1) — each successive row sees one more key (the causal
9339        // property). This matches eager: decode appends k at len, then fa_decode sees t_kv = len+1
9340        // keys. The verify appends all T tokens first but bounds the key range per row.
9341        //
9342        // MULTI-ROW FUSED PATH (the long-ctx spec fix, 2026-07-03): when every row takes the vec
9343        // kernel (base_len+1 >= FA_VEC_MIN_TKV), ONE fa_decode_rows launch executes the exact
9344        // per-row program for all T rows (grid.z = row, per-row n_splits from the same
9345        // fa_split_keys formula) — replacing T x (2 launches + 2 dtod copies + 5 partial allocs)
9346        // and multiplying resident CTAs by T on a latency-bound kernel. Bit-identical per row by
9347        // construction; kernel-check pins rows-vs-loop byte identity, run-spec is the end gate.
9348        // Short ctx (any row below the vec crossover) and MEMRA_NO_FA_VEC/MEMRA_FA_ROWS_OFF keep the
9349        // per-row loop (whose fa_decode picks scalar/vec per row exactly like eager decode).
9350        let mut attn = vbuf(e, t * n_head * head_dim)?; // fully written by every FA arm below
9351        let base_len = kvl.len - t; // KV len BEFORE this round's T tokens were appended
9352        // T=1 INCLUDED (2026-07-05): p-min cuts the draft to 1 in ~75% of rounds on hard
9353        // (agentic) content — the old t>1 gate sent those rounds to the per-row loop (262us/row
9354        // + q-row copy + per-row allocs vs 93us/row through the fused kernel at grid.z=1, same
9355        // program). nsys accounting: 1088 of 1456 verify FA launches were T=1 escapees.
9356        // LEAN T=1 ARM (MEMRA_SPEC_LEAN, close35): at t==1, q IS one row and fa_decode on it is
9357        // the EXACT eager decode dispatch (vec_q_v2 + combine_f32; the rows pair measured +50us
9358        // at m=1). Byte-identical: kernel-check pins rows-vs-loop identity, and the per-row loop
9359        // at t=1 is fa_decode on the same q with zero-offset copies. Gates arbitrate.
9360        if let Some(ctr) = stream_ctr {
9361            // STREAM ARM: causal base from the device counter; host kvl.len is a stale lower
9362            // bound used only for the split-sizing upper bound (+64 slack covers M pre-issued
9363            // rounds at K<=8). Views span the bound; per-row limits derive in-kernel.
9364            let upper = kvl.len + t + 64;
9365            let k_view = e.view_u8(&kvl.k, (upper.min(cache.max_ctx)) * ktb);
9366            let v_view = e.view_u8(&kvl.v, (upper.min(cache.max_ctx)) * vtb);
9367            e.fa_decode_rows_dc(
9368                &q,
9369                &k_view,
9370                &v_view,
9371                &mut attn,
9372                head_dim,
9373                n_head,
9374                n_head_kv,
9375                ctr,
9376                upper.min(cache.max_ctx),
9377                t,
9378                scale,
9379                ktb,
9380                vtb,
9381                0,
9382                false,
9383            )?;
9384        } else if spec_lean() && t == 1 {
9385            let t_kv = base_len + 1;
9386            let k_view = e.view_u8(&kvl.k, t_kv * ktb);
9387            let v_view = e.view_u8(&kvl.v, t_kv * vtb);
9388            e.fa_decode_kvmod(
9389                &q,
9390                &k_view,
9391                &v_view,
9392                &mut attn,
9393                head_dim,
9394                n_head,
9395                n_head_kv,
9396                t_kv,
9397                scale,
9398                ktb,
9399                vtb,
9400                crate::Engine::kv_fp8_on(),
9401            )?;
9402        } else if e.fa_rows_eligible(base_len, head_dim) {
9403            let k_view = e.view_u8(&kvl.k, (base_len + t) * ktb);
9404            let v_view = e.view_u8(&kvl.v, (base_len + t) * vtb);
9405            e.fa_decode_rows(
9406                &q,
9407                &k_view,
9408                &v_view,
9409                &mut attn,
9410                head_dim,
9411                n_head,
9412                n_head_kv,
9413                base_len,
9414                t,
9415                scale,
9416                ktb,
9417                vtb,
9418                None,
9419                false,
9420                crate::Engine::kv_fp8_on(),
9421                None,
9422            )?;
9423        } else {
9424            for r in 0..t {
9425                let t_kv_r = base_len + r + 1; // this row sees keys [0..t_kv_r)
9426                let k_view_r = e.view_u8(&kvl.k, t_kv_r * ktb);
9427                let v_view_r = e.view_u8(&kvl.v, t_kv_r * vtb);
9428                // copy q row into an owned buffer (fa_decode takes &CudaSlice, not CudaView)
9429                let mut q_row = vbuf(e, n_head * head_dim)?; // fully written by copy_view_into
9430                let q_src = q.slice(r * n_head * head_dim..(r + 1) * n_head * head_dim);
9431                e.copy_view_into(&mut q_row, 0, &q_src, n_head * head_dim)?;
9432                let mut attn_row = vbuf(e, n_head * head_dim)?; // fully written by fa_decode
9433                e.fa_decode_kvmod(
9434                    &q_row,
9435                    &k_view_r,
9436                    &v_view_r,
9437                    &mut attn_row,
9438                    head_dim,
9439                    n_head,
9440                    n_head_kv,
9441                    t_kv_r,
9442                    scale,
9443                    ktb,
9444                    vtb,
9445                    crate::Engine::kv_fp8_on(),
9446                )?;
9447                e.copy_into(
9448                    &mut attn,
9449                    r * n_head * head_dim,
9450                    &attn_row,
9451                    n_head * head_dim,
9452                )?;
9453            }
9454        }
9455
9456        let attn_g = match &gate {
9457            Some(gate) => {
9458                let mut gsig = vbuf(e, t * n_head * head_dim)?; // fully written by sigmoid
9459                e.sigmoid(gate, &mut gsig, t * n_head * head_dim)?;
9460                let mut ag = vbuf(e, t * n_head * head_dim)?; // fully written by mul
9461                e.mul(&attn, &gsig, &mut ag, t * n_head * head_dim)?;
9462                ag
9463            }
9464            None => attn,
9465        };
9466        // DECODE-EXACT wo projection: at m>=5 (K=4+ with pending) the generic matmul would use dp4a
9467        // (128-thread, different FP sum order than MMVQ). Force MMVQ for bit-identity with decode.
9468        Ok(e.matmul_decode_exact(&fa.wo, &attn_g, t)?)
9469    }
9470
9471    /// Context-linear bytes for a plain serving session's trunk cache.
9472    pub fn plain_session_kv_bytes_per_token(&self) -> usize {
9473        crate::cache::cache_bytes_per_token_for_plan(
9474            &self.cfg,
9475            &self.plan,
9476            0,
9477            self.plan.layers.len(),
9478        )
9479    }
9480
9481    /// `(logical bytes/token, ring-capped bytes/token, ring row cap)` for exact admission.
9482    pub fn plain_session_kv_shape(&self) -> (usize, usize, usize) {
9483        (
9484            self.plain_session_kv_bytes_per_token(),
9485            crate::cache::cache_ring_bytes_per_token_for_plan(
9486                &self.cfg,
9487                &self.plan,
9488                0,
9489                self.plan.layers.len(),
9490            ),
9491            crate::cache::cache_ring_row_cap_for_plan(&self.plan),
9492        )
9493    }
9494
9495    /// Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP
9496    /// scratch. With no MTP head this equals the plain coefficient.
9497    pub fn spec_session_kv_bytes_per_token(&self) -> usize {
9498        let scratch = self
9499            .mtp
9500            .iter()
9501            .chain(self.mtp_extra.iter())
9502            .map(|mtp| {
9503                let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9504                k + v
9505            })
9506            .sum::<usize>();
9507        self.plain_session_kv_bytes_per_token()
9508            .saturating_add(scratch)
9509    }
9510
9511    /// Spec twin of [`HybridModel::plain_session_kv_shape`]; Step35's persistent MTP scratch is
9512    /// capped by the same SWA ring rows as the trunk.
9513    pub fn spec_session_kv_shape(&self) -> (usize, usize, usize) {
9514        let total = self.spec_session_kv_bytes_per_token();
9515        let (_, mut ring, rows) = self.plain_session_kv_shape();
9516        if rows > 0 {
9517            ring = ring.saturating_add(
9518                self.mtp
9519                    .iter()
9520                    .chain(self.mtp_extra.iter())
9521                    .map(|mtp| {
9522                        let (_, _, k, v) = mtp_scratch_layout(&self.cfg, mtp.geom.as_ref());
9523                        k + v
9524                    })
9525                    .sum::<usize>(),
9526            );
9527        }
9528        (total, ring, rows)
9529    }
9530
9531    /// Greedy MTP speculative decode (§B). Token-identical to `generate(prompt, max_new)` but uses
9532    /// the NextN head to draft K tokens then verifies them in one batched target forward.
9533    /// Returns (generated tokens, total_drafted, total_accepted) so the caller can report
9534    /// acceptance rate. `k` = draft length per round.
9535    ///
9536    /// GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is
9537    /// Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE
9538    /// and replayed per draft step — the ~40 eager launches per drafted token collapse into one
9539    /// graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step.
9540    /// Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the
9541    /// captured graph references is event-free; the spec loop is strictly single-stream.
9542    /// MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain.
9543    /// SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax,
9544    /// device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are
9545    /// bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in
9546    /// generate_spec_inner2.
9547    /// Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so
9548    /// turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a
9549    /// 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the
9550    /// hybrid linear-attn states are in-place (no position index), so a session can extend but
9551    /// never rewind — `committed` is the exact token list whose state the caches hold (includes
9552    /// any overshoot tokens past max_new; the caller renders from `committed`, not its own echo).
9553    pub fn new_session(
9554        &self,
9555        e: &Engine,
9556        max_ctx: usize,
9557    ) -> Result<SpecSession, Box<dyn std::error::Error>> {
9558        Ok(SpecSession {
9559            // STAGE-OWNED KV (lane/pp2-spec 2026-08-06): `pp::new_cache`, not `Cache::new`. This
9560            // is the SERVING spec-session path, and with the ppN door open across two cards a
9561            // primary-homed cache makes every remote stage peer-read its OWN KV on every verify
9562            // round — the wrong-card class already fixed on the two batched serving paths
9563            // (worker.rs 2483 / 2837). With the door shut `new_cache` IS `Cache::new` (same
9564            // branch, same allocations), so single-device behavior is byte-unchanged.
9565            cache: crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?,
9566            scratch: self.new_mtp_scratch(e, max_ctx)?,
9567            committed: Vec::new(),
9568            last_h: None,
9569            next_pred: None,
9570            sctr: 0,
9571            uctr: 0,
9572            draft_ctx: None,
9573            pending_tok: None,
9574            turn_ckpt: None,
9575            telem: SpecTelemetryCounters::default(),
9576            capture_at: None,
9577            boundary_captures: Vec::new(),
9578            ckpt_at: None,
9579        })
9580    }
9581
9582    /// SPEC-ON-CACHE-HIT restore (lane/spec-on-cache-hit, 2026-08-18 — PORT-PLAN item 3,
9583    /// research/cache-spec-design-20260814, scoped to WHOLE-ENTRY restores only): build a
9584    /// SpecSession around a trunk cache the worker already restored from a prefix-cache
9585    /// entry, re-installing the entry's published draft plane as the MTP scratch rows
9586    /// `[0..prefix.len())` and the entry's boundary hidden as `last_h`, then feeding the
9587    /// prompt SUFFIX here — through EXACTLY the plain path's program selection — so the
9588    /// worker always receives a fully-warm continuation session (committed = whole
9589    /// prompt, `next_pred` + `last_h` set; caller sets `next_pred` from the entry's
9590    /// boundary logits on the empty-suffix shape).
9591    ///
9592    /// PROGRAM LAW (the splitiso two-programs class, learned AGAIN in this lane's own
9593    /// gate): the identity target for a converted hit is the PLAIN hit serving the same
9594    /// request, and plain feeds a carried suffix via eager `decode_step` below
9595    /// PRIME_MIN_T and via `prime_cache` at/above it (prefill_tick's arms). The generate
9596    /// path's tokenwise arm routes qwen35-class through the BATCHED T=1 program
9597    /// (`spec_target_step_h`) instead — ULP-different suffix rows, and the gate measured
9598    /// the near-tie flip at generated token ~8 (research/spec-cache-20260818, qwen r3).
9599    /// So the suffix is fed HERE, mirroring prefill_tick arm-for-arm, not handed to the
9600    /// burst prime.
9601    ///
9602    /// SEED RULE (both sampling regimes; lane/sampled-hit-spec 2026-08-19, sampled draw
9603    /// added by lane/sampled-spec-quality 2026-08-19). The boundary token is produced by
9604    /// EXACTLY the rule the cold burst entry applies to its own first token from the same
9605    /// logits row: `argmax` when greedy, and a `sample_boundary_token` draw at Philox
9606    /// counter 0 when sampled. Both shapes are covered — the entry's boundary logits on a
9607    /// full-cover (empty-suffix) hit, this feed's own boundary logits on a suffix hit.
9608    /// That is what keeps a restored session seed-identical to a cold one PER SEED: the
9609    /// cold session draws from the identical row at counter 0 and then runs its rounds from
9610    /// counter 1, so the restored session admits with `sctr = 1` after its own draw.
9611    /// The WORKER owns the one refusal this constructor cannot see — a constrained request.
9612    /// (The penalized-sampled refusal was LIFTED once the burst's penalty window learned to
9613    /// span the session: `committed` here is the WHOLE prompt, so the restored session's
9614    /// window is the cold session's window. It comes back if `MEMRA_SPEC_PEN_SESSION=0`.)
9615    ///
9616    /// NOT the rolled-back partial-restore hazard: the caller restores at exactly the
9617    /// entry's captured endpoint (`e.pos`) through the shipping whole-entry path;
9618    /// mid-entry (`at < e.pos`) trunk restores stay behind MEMRA_PREFIX_PARTIAL_RESTORE
9619    /// and are never routed here.
9620    ///
9621    /// Failure contract: `Err((Some(cache), why))` before any trunk mutation — the
9622    /// worker rebuilds the plain carrier and the hit serves plain, byte-unchanged.
9623    /// `Err((None, why))` after the suffix feed began — the carrier is part-fed and
9624    /// UNUSABLE; the worker serves the request cold-plain (correct, slower) and the
9625    /// entry stays published for the next request.
9626    #[allow(clippy::too_many_arguments)]
9627    pub fn spec_session_from_restored(
9628        &self,
9629        e: &Engine,
9630        mut cache: Cache,
9631        prefix: Vec<u32>,
9632        suffix: &[u32],
9633        draft_k: &CudaSlice<u8>,
9634        draft_v: &CudaSlice<u8>,
9635        draft_k_tok_bytes: usize,
9636        draft_v_tok_bytes: usize,
9637        draft_len: usize,
9638        last_h: &[f32],
9639        // The ENTRY's boundary logits row (the full-cover shape's seed source). May be empty
9640        // when a suffix follows — the feed's own logits are the boundary then.
9641        boundary_logits: &[f32],
9642        // The request's sampler, or None for greedy. Owned here so the seed rule lives in
9643        // ONE place instead of being half-applied by the worker.
9644        sampling: Option<SpecSampling>,
9645        require_anchor: bool,
9646        max_ctx: usize,
9647        // STABLE-BOUNDARY REPUBLICATION (lane/frspec-multiturn-cache, 2026-08-21): ABSOLUTE
9648        // prompt position to split the suffix feed at and capture the extended-entry
9649        // publication + this session's `turn_ckpt` — the worker's stable pre-generation
9650        // boundary (`plain_checkpoint_boundary`). None = legacy prompt-end republication.
9651        // WHY: the prompt-end capture below includes the template's live generation header
9652        // (`<|im_start|>assistant\n<think>\n`), which the next turn's re-render replaces, so
9653        // for a hybrid (whole-entry restores only) every extended entry's last ~2 tokens
9654        // diverged from every future prompt and the hit boundary FROZE at the first
9655        // lcp-split entry forever (measured: cached 6811 of 38228 by turn 8, B4).
9656        republish_at: Option<usize>,
9657    ) -> Result<SpecSession, (Option<Cache>, String)> {
9658        let pos = prefix.len();
9659        let fail = |cache: Cache, msg: String| -> Result<SpecSession, (Option<Cache>, String)> {
9660            Err((Some(cache), msg))
9661        };
9662        if self.mtp.is_none() {
9663            return fail(cache, "no MTP head attached (nothing to draft with)".into());
9664        }
9665        if pos == 0 {
9666            return fail(cache, "empty committed prefix".into());
9667        }
9668        if cache.pos != pos {
9669            let msg = format!(
9670                "restored cache pos {} != restored prefix len {pos}",
9671                cache.pos
9672            );
9673            return fail(cache, msg);
9674        }
9675        if draft_len != pos {
9676            return fail(
9677                cache,
9678                format!("draft plane len {draft_len} != restored prefix len {pos}"),
9679            );
9680        }
9681        if pos + suffix.len() >= max_ctx {
9682            return fail(
9683                cache,
9684                format!(
9685                    "prompt {} + suffix would not leave generation room in ctx {max_ctx}",
9686                    pos + suffix.len(),
9687                ),
9688            );
9689        }
9690        let mut scratch = match MtpScratch::new(
9691            e,
9692            &self.cfg,
9693            &self.plan,
9694            max_ctx,
9695            self.mtp.as_ref().and_then(|m| m.geom.as_ref()),
9696        ) {
9697            Ok(s) => s,
9698            Err(err) => return fail(cache, format!("draft scratch alloc failed: {err}")),
9699        };
9700        if scratch.kv.ring.is_some() {
9701            return fail(
9702                cache,
9703                "ring-backed draft scratch (Step35 SWA) cannot take a flat prefix restore".into(),
9704            );
9705        }
9706        if scratch.kv.k_tok_bytes != draft_k_tok_bytes
9707            || scratch.kv.v_tok_bytes != draft_v_tok_bytes
9708        {
9709            return fail(
9710                cache,
9711                format!(
9712                    "draft plane layout {draft_k_tok_bytes}/{draft_v_tok_bytes} != scratch \
9713                     {}/{} bytes/token (stale entry across a format change)",
9714                    scratch.kv.k_tok_bytes, scratch.kv.v_tok_bytes,
9715                ),
9716            );
9717        }
9718        if pos > scratch.cap {
9719            return fail(
9720                cache,
9721                format!(
9722                    "draft plane rows {pos} exceed scratch capacity {}",
9723                    scratch.cap
9724                ),
9725            );
9726        }
9727        let kb = pos * draft_k_tok_bytes;
9728        let vb = pos * draft_v_tok_bytes;
9729        if draft_k.len() < kb || draft_v.len() < vb {
9730            return fail(
9731                cache,
9732                format!(
9733                    "truncated draft plane: K {} < {kb} or V {} < {vb} bytes",
9734                    draft_k.len(),
9735                    draft_v.len(),
9736                ),
9737            );
9738        }
9739        if kb > 0 {
9740            if let Err(err) = e.copy_u8_into(&mut scratch.kv.k, 0, draft_k, kb) {
9741                return fail(cache, format!("draft K restore copy failed: {err}"));
9742            }
9743        }
9744        if vb > 0 {
9745            if let Err(err) = e.copy_u8_into(&mut scratch.kv.v, 0, draft_v, vb) {
9746                return fail(cache, format!("draft V restore copy failed: {err}"));
9747            }
9748        }
9749        if let Err(err) = scratch.set_len(e, pos) {
9750            return fail(cache, format!("draft scratch len set failed: {err}"));
9751        }
9752        let mut last_h_dev = if last_h.len() == self.cfg.n_embd as usize {
9753            // anchor upload failure is acceptance-only when a suffix feed follows (fill
9754            // row-0 falls back to zeros) but FATAL for an empty-suffix continuation (the
9755            // burst entry asserts committed + last_h + next_pred) — the caller says which.
9756            e.htod(last_h).ok()
9757        } else {
9758            None
9759        };
9760        if require_anchor && last_h_dev.is_none() {
9761            return fail(
9762                cache,
9763                "empty-suffix continuation requires the entry's boundary hidden anchor".into(),
9764            );
9765        }
9766        let mut committed = prefix;
9767        // Set on BOTH shapes below (suffix-fed and full-cover) — never left None, which is
9768        // what the empty-suffix continuation assert in the burst entry requires.
9769        let next_pred;
9770        // Philox: (0,0) at admit exactly like a fresh session; a sampled boundary draw below
9771        // consumes counter 0 and leaves 1, which is the state a cold session reaches after
9772        // drawing its own first token from the same row.
9773        let mut sctr = 0u32;
9774        let sampled = sampling.is_some_and(|s| s.temp > 0.0) && spec_sampled_boundary_on();
9775        // Penalty window for the boundary draw: the last `penalty_last_n` tokens of the WHOLE
9776        // prompt, which is what the cold session's own burst sees (Item 2's window). Built
9777        // after the suffix joins `committed` below.
9778        let mut boundary_captures: Vec<SpecBoundaryCapture> = Vec::new();
9779        let mut restored_turn_ckpt: Option<SpecCheckpoint> = None;
9780        if !suffix.is_empty() {
9781            // ---- SUFFIX FEED, mirroring prefill_tick's program selection exactly ----
9782            // From here on the trunk cache mutates: failures return Err((None, _)) and
9783            // the worker serves the request cold-plain instead of reusing the carrier.
9784            let dirty =
9785                |msg: String| -> Result<SpecSession, (Option<Cache>, String)> { Err((None, msg)) };
9786            let n_embd = self.cfg.n_embd as usize;
9787            let t = suffix.len();
9788            let mut h_rows = match e.uninit(t * n_embd) {
9789                Ok(b) => b,
9790                Err(err) => return fail(cache, format!("suffix hidden buffer alloc: {err}")),
9791            };
9792            // STABLE-BOUNDARY split (see `republish_at`): feed stops at the boundary so the
9793            // in-place GDN conv/ssm state can be snapshotted there — the only moment it
9794            // exists (the cold prime-split law). suffix-relative; None = one-segment legacy.
9795            let b_rel = republish_at
9796                .and_then(|abs| abs.checked_sub(pos))
9797                .filter(|&r| r > 0 && r < t);
9798            let mut feed_logits = Vec::new();
9799            let tokenwise_env = std::env::var("MEMRA_PRIME_TOKENWISE").is_ok()
9800                || e.frozen_cpu_experts_prefer_tokenwise_prime();
9801            let mut fed = 0usize;
9802            for seg_end in [b_rel, Some(t)].into_iter().flatten() {
9803                if seg_end <= fed {
9804                    continue;
9805                }
9806                let seg = &suffix[fed..seg_end];
9807                let batched = seg.len() >= crate::hybrid_forward::PRIME_MIN_T && !tokenwise_env;
9808                if batched {
9809                    // prefill_tick's prime arm: request-level prime_cache call; tokens still
9810                    // queued after this segment ride `queued_after` so Step35 arm selection
9811                    // stays keyed to the request's end (tick-seg law).
9812                    match self.prime_cache(e, seg, &mut cache, t - seg_end) {
9813                        Ok((l, _h_seed, hiddens)) => {
9814                            if let Err(err) =
9815                                e.copy_into(&mut h_rows, fed * n_embd, &hiddens, seg.len() * n_embd)
9816                            {
9817                                return dirty(format!("suffix hidden copy: {err}"));
9818                            }
9819                            feed_logits = l;
9820                        }
9821                        Err(err) => return dirty(format!("suffix prime failed: {err}")),
9822                    }
9823                } else {
9824                    // prefill_tick's tokenwise arm: eager decode_step, one token at a time.
9825                    for (i, &tok) in seg.iter().enumerate() {
9826                        match self.decode_step_h(e, tok, &mut cache) {
9827                            Ok((l, h)) => {
9828                                if let Err(err) =
9829                                    e.copy_into(&mut h_rows, (fed + i) * n_embd, &h, n_embd)
9830                                {
9831                                    return dirty(format!("suffix hidden copy: {err}"));
9832                                }
9833                                feed_logits = l;
9834                            }
9835                            Err(err) => return dirty(format!("suffix decode_step failed: {err}")),
9836                        }
9837                    }
9838                }
9839                fed = seg_end;
9840                if Some(seg_end) == b_rel {
9841                    // The stable pre-generation boundary: capture the extended-entry
9842                    // publication AND this session's own turn checkpoint here instead of at
9843                    // prompt-end (both would otherwise carry the volatile live-header tail
9844                    // the next re-render replaces). Failure silent, turn_ckpt convention.
9845                    debug_assert_eq!(
9846                        cache.pos,
9847                        pos + seg_end,
9848                        "stable-boundary capture off the feed split"
9849                    );
9850                    if spec_restore_republish_on() {
9851                        if let Ok(snap) = cache.snapshot(e) {
9852                            boundary_captures.push(SpecBoundaryCapture {
9853                                snap,
9854                                pos: pos + seg_end,
9855                                logits: feed_logits.clone(),
9856                                last_h: capture_boundary_hidden(e, &h_rows, seg_end, n_embd),
9857                            });
9858                        }
9859                    }
9860                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
9861                        e.uninit(n_embd).and_then(|mut a| {
9862                            e.copy_view_into(
9863                                &mut a,
9864                                0,
9865                                &h_rows.slice((seg_end - 1) * n_embd..seg_end * n_embd),
9866                                n_embd,
9867                            )?;
9868                            Ok(a)
9869                        });
9870                    if let (Ok(snap), Ok(last_h)) = (cache.snapshot(e), anchor) {
9871                        restored_turn_ckpt = Some(SpecCheckpoint {
9872                            snap,
9873                            pos: pos + seg_end,
9874                            last_h,
9875                        });
9876                    }
9877                }
9878            }
9879            // Draft-scratch fill for the suffix rows, predecessor-paired: row `pos` reads
9880            // the entry's boundary anchor (zeros fallback — acceptance-only), row `pos+i`
9881            // reads h_rows[i-1]. Chunked like the generate path's fill (transients scale
9882            // with T). Fill failures are acceptance-only — truncate to the restored rows
9883            // and continue; the burst's own set_len keeps the invariant.
9884            let mtp = self.mtp.as_ref().expect("mtp checked above");
9885            let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
9886            let embd_gpu = if spec_host_embd() {
9887                None
9888            } else {
9889                Some(
9890                    self.embd_gpu
9891                        .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
9892                )
9893            };
9894            let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
9895            let fill_chunk = 4096usize;
9896            let mut filled = true;
9897            let mut start = 0usize;
9898            'fill: while start < t {
9899                let end = (start + fill_chunk).min(t);
9900                let tc = end - start;
9901                let Ok(mut phs) = e.zeros(tc * n_embd) else {
9902                    filled = false;
9903                    break 'fill;
9904                };
9905                let (src_lo, dst_off, n_copy) = if start == 0 {
9906                    (0, n_embd, (tc - 1) * n_embd)
9907                } else {
9908                    ((start - 1) * n_embd, 0, tc * n_embd)
9909                };
9910                if start == 0 {
9911                    if let Some(lh) = last_h_dev.as_ref() {
9912                        if e.copy_into(&mut phs, 0, lh, n_embd).is_err() {
9913                            filled = false;
9914                            break 'fill;
9915                        }
9916                    }
9917                }
9918                if n_copy > 0
9919                    && e.copy_view_into(
9920                        &mut phs,
9921                        dst_off,
9922                        &h_rows.slice(src_lo..src_lo + n_copy),
9923                        n_copy,
9924                    )
9925                    .is_err()
9926                {
9927                    filled = false;
9928                    break 'fill;
9929                }
9930                if self
9931                    .mtp_kv_fill_all(
9932                        e,
9933                        &suffix[start..end],
9934                        &phs,
9935                        pos + start,
9936                        &mut scratch,
9937                        embd_dev,
9938                    )
9939                    .is_err()
9940                {
9941                    filled = false;
9942                    break 'fill;
9943                }
9944                start = end;
9945            }
9946            if !filled {
9947                // acceptance-only: drafts over missing suffix rows are cheap and wrong,
9948                // so keep only the restored rows resident and let verify arbitrate.
9949                if let Err(err) = scratch.set_len(e, pos) {
9950                    return dirty(format!("scratch truncation after failed fill: {err}"));
9951                }
9952            }
9953            // EXTENDED-ENTRY PUBLICATION (lane/sampled-spec-quality, Item 3 — the fix for
9954            // "a restored spec session never publishes an extended entry", SAMPLED-HIT.md
9955            // finding (d)). Pre-lane, publication was armed only for COLD sessions
9956            // (`spec_resumed == 0` in the worker) and both engine capture sites require a
9957            // non-continuation burst — but a converted hit's first burst IS a continuation,
9958            // so a growing conversation learned exactly ONE boundary and turn 3 could never
9959            // hit a longer prefix than turn 2 did.
9960            //
9961            // WHERE, and why it is safe here: `cache.pos == prefix + suffix` at this exact
9962            // line — the trunk is primed over the whole prompt, nothing is generated, and the
9963            // draft plane rows [0..prompt) are filled just above. That is a complete
9964            // whole-entry boundary (`pos == fed_len`), the same shape the cold seed capture
9965            // publishes; the worker's existing publication sweep picks it up because it is
9966            // keyed on non-empty `boundary_captures` and is sampler- and resume-independent.
9967            // NOT the partial-restore hazard: the boundary is this session's own prompt END,
9968            // never mid-entry, so `entry_pos != fed_len` still refuses on the way back in.
9969            // Failure is SILENT by design (the turn_ckpt / boundary-capture convention):
9970            // publication is an optimization, never a correctness dependency.
9971            //
9972            // SUPERSEDED WHEN `republish_at` FIRED (lane/frspec-multiturn-cache): a prompt-end
9973            // entry's tail is the live generation header the next re-render replaces, so on a
9974            // hybrid (whole-entry restores) it can never serve the conversation's next turn —
9975            // the stable-boundary capture above IS this publication, minus the poisoned tail.
9976            if spec_restore_republish_on() && boundary_captures.is_empty() {
9977                debug_assert_eq!(
9978                    cache.pos,
9979                    pos + t,
9980                    "extended-entry capture must sit at the restored session's prompt end",
9981                );
9982                if let Ok(snap) = cache.snapshot(e) {
9983                    boundary_captures.push(SpecBoundaryCapture {
9984                        snap,
9985                        pos: pos + t,
9986                        logits: feed_logits.clone(),
9987                        last_h: capture_boundary_hidden(e, &h_rows, t, n_embd),
9988                    });
9989                }
9990            }
9991            // continuation seed: the feed's boundary logits ARE the plain path's boundary
9992            // logits (same program), so greedy's argmax here is plain's first emitted token,
9993            // and the sampled draw is the cold sampled session's own first token.
9994            next_pred = Some(if sampled {
9995                let sp = sampling.expect("sampled implies a sampler");
9996                // `committed` is still the restored prefix here; the suffix joins it below —
9997                // so this is the last-N window over the WHOLE prompt, exactly the cold
9998                // session's own window at its first token.
9999                let hist = pen_window_seed(&committed, suffix, sp.penalty_last_n);
10000                match sample_boundary_token(
10001                    e,
10002                    &feed_logits,
10003                    &sp,
10004                    &hist,
10005                    &mut sctr,
10006                    "restore-suffix-feed",
10007                ) {
10008                    Ok(t) => t,
10009                    // the trunk is already fed: hand nothing back, the worker serves the
10010                    // request cold-plain. Never fall back to an argmax — that would put a
10011                    // greedy token in a sampled stream to save a slow path.
10012                    Err(err) => {
10013                        return dirty(format!("boundary token draw failed: {err}"));
10014                    }
10015                }
10016            } else {
10017                argmax(&feed_logits) as u32
10018            });
10019            let mut lh = match e.uninit(n_embd) {
10020                Ok(b) => b,
10021                Err(err) => return dirty(format!("boundary hidden alloc: {err}")),
10022            };
10023            if let Err(err) = e.copy_view_into(
10024                &mut lh,
10025                0,
10026                &h_rows.slice((t - 1) * n_embd..t * n_embd),
10027                n_embd,
10028            ) {
10029                return dirty(format!("boundary hidden copy: {err}"));
10030            }
10031            last_h_dev = Some(lh);
10032            committed.extend_from_slice(suffix);
10033        } else {
10034            // FULL-COVER shape (empty suffix — the identical-repeat / agent-loop shape): the
10035            // ENTRY's boundary logits are the boundary row, and this is the token the cold
10036            // session emits from that same row. Owned here rather than in the worker so the
10037            // sampled draw cannot be half-applied on one shape (the worker used to argmax it).
10038            if boundary_logits.is_empty() {
10039                return fail(
10040                    cache,
10041                    "full-cover restore without the entry's boundary logits".into(),
10042                );
10043            }
10044            next_pred = Some(if sampled {
10045                let sp = sampling.expect("sampled implies a sampler");
10046                let hist = pen_window_seed(&committed, &[], sp.penalty_last_n);
10047                match sample_boundary_token(
10048                    e,
10049                    boundary_logits,
10050                    &sp,
10051                    &hist,
10052                    &mut sctr,
10053                    "restore-full-cover",
10054                ) {
10055                    Ok(t) => t,
10056                    // nothing has been mutated on this shape — hand the carrier back and let
10057                    // the hit serve PLAIN (the banked pre-lane path).
10058                    Err(err) => {
10059                        return fail(cache, format!("boundary token draw failed: {err}"));
10060                    }
10061                }
10062            } else {
10063                argmax(boundary_logits) as u32
10064            });
10065        }
10066        Ok(SpecSession {
10067            cache,
10068            scratch,
10069            committed,
10070            last_h: last_h_dev,
10071            next_pred,
10072            sctr,
10073            uctr: 0,
10074            draft_ctx: None,
10075            pending_tok: None,
10076            // Stable-boundary capture from the split feed above (None on the legacy shape):
10077            // a restored session previously parked WITHOUT a checkpoint, so the next turn's
10078            // affinity probe declined ("no turn checkpoint retained") and the conversation
10079            // fell back to the frozen prefix entry forever.
10080            turn_ckpt: restored_turn_ckpt,
10081            telem: SpecTelemetryCounters::default(),
10082            capture_at: None,
10083            boundary_captures,
10084            ckpt_at: None,
10085        })
10086    }
10087
10088    /// Forced-gate exact state comparison. This intentionally reads the real live prefixes from
10089    /// their owning PP devices: matching emitted ids alone would miss a stale `len_d`, recurrent
10090    /// snapshot, or draft-KV row that only corrupts the following round.
10091    pub fn optipipe_compare_session_state(
10092        &self,
10093        e: &Engine,
10094        reference: &SpecSession,
10095        candidate: &SpecSession,
10096    ) -> Result<OptiForkStateIdentity, Box<dyn std::error::Error>> {
10097        fn fail(what: &str) -> Box<dyn std::error::Error> {
10098            format!("optipipe state mismatch: {what}").into()
10099        }
10100        fn same_f32(a: &[f32], b: &[f32]) -> bool {
10101            a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
10102        }
10103        fn compare_layers(
10104            es: &Engine,
10105            range: std::ops::Range<usize>,
10106            reference: &SpecSession,
10107            candidate: &SpecSession,
10108            report: &mut OptiForkStateIdentity,
10109        ) -> Result<(), Box<dyn std::error::Error>> {
10110            for il in range {
10111                match (&reference.cache.kv[il], &candidate.cache.kv[il]) {
10112                    (Some(a), Some(b)) => {
10113                        if a.len != b.len {
10114                            return Err(fail(&format!(
10115                                "layer {il} host KV len {} != {}",
10116                                a.len, b.len
10117                            )));
10118                        }
10119                        let ad = es.dtoh_i32(&a.len_d)?;
10120                        let bd = es.dtoh_i32(&b.len_d)?;
10121                        if ad != bd || ad.first().copied() != Some(a.len as i32) {
10122                            return Err(fail(&format!(
10123                                "layer {il} device KV len {ad:?} != {bd:?} (host={})",
10124                                a.len,
10125                            )));
10126                        }
10127                        let kb = a.len * a.k_tok_bytes;
10128                        let vb = a.len * a.v_tok_bytes;
10129                        if kb > 0 {
10130                            let ak = es.dtoh_u8_view(&a.k.slice(0..kb))?;
10131                            let bk = es.dtoh_u8_view(&b.k.slice(0..kb))?;
10132                            if ak != bk {
10133                                let at = ak.iter().zip(&bk).position(|(x, y)| x != y).unwrap();
10134                                return Err(fail(&format!(
10135                                    "layer {il} K bytes at byte {at} row {} offset {}: {} != {}",
10136                                    at / a.k_tok_bytes,
10137                                    at % a.k_tok_bytes,
10138                                    ak[at],
10139                                    bk[at],
10140                                )));
10141                            }
10142                        }
10143                        if vb > 0 {
10144                            let av = es.dtoh_u8_view(&a.v.slice(0..vb))?;
10145                            let bv = es.dtoh_u8_view(&b.v.slice(0..vb))?;
10146                            if av != bv {
10147                                let at = av.iter().zip(&bv).position(|(x, y)| x != y).unwrap();
10148                                return Err(fail(&format!(
10149                                    "layer {il} V bytes at byte {at} row {} offset {}: {} != {}",
10150                                    at / a.v_tok_bytes,
10151                                    at % a.v_tok_bytes,
10152                                    av[at],
10153                                    bv[at],
10154                                )));
10155                            }
10156                        }
10157                        report.trunk_kv_bytes += kb + vb;
10158                    }
10159                    (None, None) => {}
10160                    _ => return Err(fail(&format!("layer {il} KV presence"))),
10161                }
10162                match (&reference.cache.recur[il], &candidate.cache.recur[il]) {
10163                    (Some(a), Some(b)) => {
10164                        let ac = es.dtoh(&a.conv_state)?;
10165                        let bc = es.dtoh(&b.conv_state)?;
10166                        if !same_f32(&ac, &bc) {
10167                            return Err(fail(&format!("layer {il} conv state")));
10168                        }
10169                        let as_ = es.dtoh(&a.ssm_state)?;
10170                        let bs = es.dtoh(&b.ssm_state)?;
10171                        if !same_f32(&as_, &bs) {
10172                            return Err(fail(&format!("layer {il} SSM state")));
10173                        }
10174                        report.recurrent_bytes += (ac.len() + as_.len()) * 4;
10175                    }
10176                    (None, None) => {}
10177                    _ => return Err(fail(&format!("layer {il} recurrent presence"))),
10178                }
10179            }
10180            Ok(())
10181        }
10182
10183        if reference.committed != candidate.committed {
10184            return Err(fail("committed token ids"));
10185        }
10186        if reference.cache.pos != candidate.cache.pos
10187            || reference.cache.max_ctx != candidate.cache.max_ctx
10188        {
10189            return Err(fail("cache pos/capacity"));
10190        }
10191        if reference.pending_tok != candidate.pending_tok
10192            || reference.next_pred != candidate.next_pred
10193            || reference.sctr != candidate.sctr
10194            || reference.uctr != candidate.uctr
10195        {
10196            return Err(fail("pending/prediction/counter tail"));
10197        }
10198
10199        let mut report = OptiForkStateIdentity::default();
10200        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
10201            let rt = crate::pp::PpNRt::get(e)?;
10202            for stage in 0..rt.n_stages() {
10203                let _scope = rt.enter(stage);
10204                compare_layers(
10205                    rt.engine(stage, e),
10206                    fence[stage]..fence[stage + 1],
10207                    reference,
10208                    candidate,
10209                    &mut report,
10210                )?;
10211            }
10212        } else {
10213            compare_layers(e, 0..self.layers.len(), reference, candidate, &mut report)?;
10214        }
10215
10216        if reference.scratch.plane_count() != candidate.scratch.plane_count() {
10217            return Err(fail("draft scratch plane count"));
10218        }
10219        for index in 0..reference.scratch.plane_count() {
10220            let (a, _) = reference.scratch.plane(index);
10221            let (b, _) = candidate.scratch.plane(index);
10222            if a.len != b.len
10223                || a.kv_dim_k != b.kv_dim_k
10224                || a.kv_dim_v != b.kv_dim_v
10225                || a.k_tok_bytes != b.k_tok_bytes
10226                || a.v_tok_bytes != b.v_tok_bytes
10227                || e.dtoh_i32(&a.len_d)? != e.dtoh_i32(&b.len_d)?
10228            {
10229                return Err(fail(&format!("draft scratch plane {index} length/layout")));
10230            }
10231            let kb = a.len * a.k_tok_bytes;
10232            let vb = a.len * a.v_tok_bytes;
10233            if kb > 0 && e.dtoh_u8_view(&a.k.slice(0..kb))? != e.dtoh_u8_view(&b.k.slice(0..kb))? {
10234                return Err(fail(&format!("draft scratch plane {index} K bytes")));
10235            }
10236            if vb > 0 && e.dtoh_u8_view(&a.v.slice(0..vb))? != e.dtoh_u8_view(&b.v.slice(0..vb))? {
10237                return Err(fail(&format!("draft scratch plane {index} V bytes")));
10238            }
10239            report.scratch_kv_bytes += kb + vb;
10240        }
10241
10242        match (&reference.last_h, &candidate.last_h) {
10243            (Some(a), Some(b)) => {
10244                let ah = e.dtoh(a)?;
10245                let bh = e.dtoh(b)?;
10246                if !same_f32(&ah, &bh) {
10247                    return Err(fail("last hidden/seed bytes"));
10248                }
10249                report.hidden_bytes = ah.len() * 4;
10250            }
10251            (None, None) => {}
10252            _ => return Err(fail("last hidden/seed presence")),
10253        }
10254        Ok(report)
10255    }
10256
10257    /// SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll `sess` back to its
10258    /// retained prompt-end checkpoint, so a request whose prompt matches
10259    /// `committed[..rewind_pos()]` exactly can resume there and prime only its own delta.
10260    ///
10261    /// EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that
10262    /// boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm
10263    /// restored from the device copy taken there, draft scratch length reset, `committed`
10264    /// truncated, `last_h` = the boundary's predecessor anchor. That is precisely the state a
10265    /// fresh prime of `committed[..pos]` would have produced, so the following suffix prime and
10266    /// every burst after it are identical to a cold run of the same token stream — the
10267    /// committed-tokens-authoritative contract.
10268    ///
10269    /// `next_pred` and `pending_tok` are CLEARED: both describe generation past the boundary,
10270    /// which the rewind discards. The caller therefore must supply a non-empty suffix (a
10271    /// rewound session cannot serve an empty-suffix continuation burst — there is nothing to
10272    /// continue). The persistent draft graph survives: it bakes only session-stable pointers
10273    /// (the scratch KV, the resident embedding), none of which the rewind moves.
10274    ///
10275    /// The checkpoint is CONSUMED (`turn_ckpt` taken): its snapshot buffers are freed here, and
10276    /// this turn's own prime installs a fresh one at the new prompt end. Returns the position
10277    /// rewound to, or `None` when the session holds no checkpoint (caller: full re-prime).
10278    pub fn spec_rewind_to_checkpoint(
10279        &self,
10280        e: &Engine,
10281        sess: &mut SpecSession,
10282    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10283        if sess.turn_ckpt.as_ref().is_some_and(|ckpt| {
10284            !sess.cache.can_rollback(&ckpt.snap, 0) || !sess.scratch.can_rewind_to(ckpt.pos)
10285        }) {
10286            return Err(
10287                "SWA ring rewind checkpoint has been lapped; full re-prime required".into(),
10288            );
10289        }
10290        let Some(ckpt) = sess.turn_ckpt.take() else {
10291            return Ok(None);
10292        };
10293        assert!(
10294            ckpt.pos <= sess.committed.len(),
10295            "checkpoint past committed ({} > {})",
10296            ckpt.pos,
10297            sess.committed.len()
10298        );
10299        // Restore through each layer's owning engine. A single primary-engine rollback is not
10300        // sufficient when the serving cache is stage-owned under cross-device PP.
10301        crate::pp::restore_cache_checkpoint(e, self, None, &mut sess.cache, &ckpt.snap)?;
10302        debug_assert_eq!(
10303            sess.cache.pos, ckpt.pos,
10304            "rollback landed off the checkpoint"
10305        );
10306        sess.scratch.set_len(e, ckpt.pos)?;
10307        sess.committed.truncate(ckpt.pos);
10308        sess.last_h = Some(ckpt.last_h);
10309        sess.next_pred = None;
10310        sess.pending_tok = None;
10311        Ok(Some(ckpt.pos))
10312    }
10313
10314    /// Grow a parked speculative session to `target_cap` and rewind it to its retained turn
10315    /// checkpoint without re-priming the checkpoint prefix.
10316    ///
10317    /// The trunk cache is restored exactly like a plain grown cache: append-only full-attention
10318    /// KV rows come from the parked cache, while recurrent state comes from the checkpoint's
10319    /// owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint
10320    /// remain authoritative, so they are copied into a fresh larger scratch before its length is
10321    /// truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.
10322    ///
10323    /// All fallible work completes before `sess` is mutated. A failed allocation or copy leaves
10324    /// the parked session intact, allowing the caller one reclaim-and-retry attempt.
10325    pub fn spec_grow_and_rewind_to_checkpoint(
10326        &self,
10327        e: &Engine,
10328        sess: &mut SpecSession,
10329        target_cap: usize,
10330    ) -> Result<Option<usize>, Box<dyn std::error::Error>> {
10331        if target_cap <= sess.cache.max_ctx {
10332            return self.spec_rewind_to_checkpoint(e, sess);
10333        }
10334        let Some(ckpt) = sess.turn_ckpt.as_ref() else {
10335            return Ok(None);
10336        };
10337        if ckpt.pos == 0 || ckpt.pos > sess.committed.len() {
10338            return Err(format!(
10339                "checkpoint pos {} outside committed length {}",
10340                ckpt.pos,
10341                sess.committed.len(),
10342            )
10343            .into());
10344        }
10345        if ckpt.pos > target_cap {
10346            return Err(format!(
10347                "checkpoint pos {} exceeds grown capacity {target_cap}",
10348                ckpt.pos,
10349            )
10350            .into());
10351        }
10352
10353        let mut grown_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, target_cap)?;
10354        let mut grown_scratch = self.new_mtp_scratch(e, target_cap)?;
10355        crate::pp::restore_cache_checkpoint(
10356            e,
10357            self,
10358            Some(&sess.cache),
10359            &mut grown_cache,
10360            &ckpt.snap,
10361        )?;
10362
10363        if sess.scratch.plane_count() != grown_scratch.plane_count() {
10364            return Err("checkpoint draft plane count mismatch".into());
10365        }
10366        for index in 0..sess.scratch.plane_count() {
10367            let (src, _) = sess.scratch.plane(index);
10368            let (dst, _) = grown_scratch.plane_mut(index);
10369            if ckpt.pos > src.len
10370                || src.kv_dim_k != dst.kv_dim_k
10371                || src.kv_dim_v != dst.kv_dim_v
10372                || src.k_tok_bytes != dst.k_tok_bytes
10373                || src.v_tok_bytes != dst.v_tok_bytes
10374            {
10375                return Err(format!(
10376                    "checkpoint draft plane {index} layout mismatch (pos {}, source len {})",
10377                    ckpt.pos, src.len,
10378                )
10379                .into());
10380            }
10381            match (&src.ring, dst.ring.as_ref()) {
10382                (Some(sring), Some(_)) => {
10383                    // Ring-backed draft plane (step35): `ckpt.pos` is absolute and exceeds the
10384                    // physical rows once lapped — same class as the trunk-KV restore panic
10385                    // (2026-08-29 warm-turn-at-40k). Copy the aligned live window, rebase.
10386                    let (new_base, phys) = sring.restore_plan(ckpt.pos).map_err(|err| {
10387                        format!("checkpoint draft plane {index} SWA restore refused: {err}")
10388                    })?;
10389                    let rows = phys.len();
10390                    let kb = rows * src.k_tok_bytes;
10391                    let vb = rows * src.v_tok_bytes;
10392                    if kb > 0 {
10393                        e.copy_u8_range_into(
10394                            &mut dst.k,
10395                            0,
10396                            &src.k,
10397                            phys.start * src.k_tok_bytes,
10398                            kb,
10399                        )?;
10400                    }
10401                    if vb > 0 {
10402                        e.copy_u8_range_into(
10403                            &mut dst.v,
10404                            0,
10405                            &src.v,
10406                            phys.start * src.v_tok_bytes,
10407                            vb,
10408                        )?;
10409                    }
10410                    dst.ring
10411                        .as_mut()
10412                        .expect("ring presence checked above")
10413                        .apply_rebase(new_base);
10414                    if let Some(base_d) = dst.base_d.as_mut() {
10415                        e.set_i32_one(base_d, new_base as i32)?;
10416                    }
10417                }
10418                (None, None) => {
10419                    let kb = ckpt.pos * src.k_tok_bytes;
10420                    let vb = ckpt.pos * src.v_tok_bytes;
10421                    if kb > 0 {
10422                        e.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
10423                    }
10424                    if vb > 0 {
10425                        e.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
10426                    }
10427                }
10428                _ => {
10429                    return Err(format!("checkpoint draft plane {index} ring/flat mismatch").into());
10430                }
10431            }
10432        }
10433        grown_scratch.set_len(e, ckpt.pos)?;
10434        // The old scratch is dropped immediately after publication below. Bound its D2D reads
10435        // first; growth happens once per rewritten turn, outside the decode hot loop.
10436        e.stream().synchronize()?;
10437
10438        let ckpt = sess
10439            .turn_ckpt
10440            .take()
10441            .expect("checkpoint remained present through transactional grow");
10442        let pos = ckpt.pos;
10443        sess.cache = grown_cache;
10444        sess.scratch = grown_scratch;
10445        sess.committed.truncate(pos);
10446        sess.last_h = Some(ckpt.last_h);
10447        sess.next_pred = None;
10448        sess.pending_tok = None;
10449        sess.draft_ctx = None;
10450        debug_assert_eq!(sess.cache.pos, pos, "grown rewind landed off checkpoint");
10451        debug_assert!(
10452            (0..sess.scratch.plane_count()).all(|index| sess.scratch.plane(index).0.len == pos),
10453            "grown draft rewind landed off checkpoint"
10454        );
10455        Ok(Some(pos))
10456    }
10457
10458    /// Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass
10459    /// (its logits' argmax becomes next_pred) + the draft-KV fill at the carried anchor —
10460    /// byte-identical to the pre-carry session tail. Required before a non-empty-suffix
10461    /// prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.
10462    /// `sampling` is the sampler of the request that will CONSUME the resulting `next_pred`
10463    /// (lane/sampled-spec-quality): this is a boundary site like any other, so a sampled
10464    /// consumer must get a DRAWN token, not an argmax. Pass `None` from the park/demote
10465    /// callers — a pending only ever exists on the GREEDY tail, and the consumer of a
10466    /// park-time flush is a future request whose sampler is not knowable here (residual
10467    /// named at the pool-resume probe in worker.rs and in SAMPLED-QUALITY.md).
10468    pub fn spec_flush_pending(
10469        &self,
10470        e: &Engine,
10471        sess: &mut SpecSession,
10472        sampling: Option<SpecSampling>,
10473    ) -> Result<(), Box<dyn std::error::Error>> {
10474        let Some(b) = sess.pending_tok.take() else {
10475            return Ok(());
10476        };
10477        if self.mtp.is_none() {
10478            return Err("pending carry requires an MTP head".into());
10479        }
10480        let n_embd = self.cfg.n_embd as usize;
10481        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
10482        let embd_gpu = if spec_host_embd() {
10483            None
10484        } else {
10485            Some(
10486                self.embd_gpu
10487                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
10488            )
10489        };
10490        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
10491        let pos_b = sess.cache.pos;
10492        sess.scratch.set_len(e, pos_b)?;
10493        let (lg_b, hb) = self.spec_target_step_h(e, b, &mut sess.cache)?;
10494        sess.next_pred = Some(match sampling {
10495            Some(sp) if sp.temp > 0.0 && spec_sampled_boundary_on() => {
10496                // window includes `b` itself: it is committed by this pass, and the pre-lane
10497                // code never counted a boundary token in the penalty history at all.
10498                let hist = pen_window_seed(&sess.committed, &[b], sp.penalty_last_n);
10499                sample_boundary_token(e, &lg_b, &sp, &hist, &mut sess.sctr, "flush-pending")?
10500            }
10501            _ => argmax(&lg_b) as u32,
10502        });
10503        let anchor = sess
10504            .last_h
10505            .as_ref()
10506            .expect("pending carry requires last_h (the predecessor-row anchor)");
10507        self.mtp_kv_fill_all(e, &[b], anchor, pos_b, &mut sess.scratch, embd_dev)?;
10508        sess.last_h = Some(hb);
10509        sess.committed.push(b);
10510        Ok(())
10511    }
10512
10513    /// Solo target feed used only at speculative round boundaries. Step35 serving made its
10514    /// staged batched B=1 graph authoritative, so a speculative session must enter and leave
10515    /// rounds through that same graph. Other model families keep their eager T=1 contract.
10516    fn spec_target_step_h(
10517        &self,
10518        e: &Engine,
10519        token: u32,
10520        cache: &mut Cache,
10521    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
10522        if !self.sliding_gated_moe_batch_program() && !self.batched_serving_numeric_class() {
10523            return self.decode_step_h(e, token, cache);
10524        }
10525        let pos0 = cache.pos;
10526        let (logits, hidden) = self.decode_step_t_core(e, &[token], pos0, cache, None, None)?;
10527        Ok((e.dtoh(&logits)?, hidden))
10528    }
10529
10530    /// The archs whose LIVE B=1 serving runs the generic BATCHED numeric class (decode_step_batch
10531    /// walk + batched head), so their spec verify must run the SAME class. MoE learned this
10532    /// 2026-08-14 AM (4b777ccc5); the dense hybrid reproduced the identical near-tie flip class
10533    /// the same day on Qwen3.8-27B — eager-class verify logits drift from batched-class serving
10534    /// logits ("1 ULP at layer 2 → 2.3e-1 logit maxdiff at the head"), and the GDN recurrence
10535    /// carries the drift until a near-tie flips deep in generation. One predicate so the five
10536    /// dispatch sites cannot drift apart again.
10537    /// Draft-graph head admissibility (lane/draftcost-moe, 2026-08-20): the capture body
10538    /// (`mtp_head_forward_cap`) supports Dense heads AND resident-MoE heads
10539    /// (`Ffn::Moe(m) if m.dev_exps.is_some()`); non-resident MoE still refuses inside the
10540    /// capture and the caller falls back to the eager chain by design. Trunk FFN class is
10541    /// irrelevant — the graph body is the HEAD forward only. One predicate for all three
10542    /// eligibility sites so they cannot drift (the serving numeric-class lesson).
10543    fn mtp_graph_capturable(&self) -> bool {
10544        // EVERY loaded head must be capturable: the multi-head chain graphs capture each
10545        // head's forward (lane/step37-draft-graph-serving-20260830), so one SLRU-locked MoE
10546        // head anywhere in the chain refuses capture for the whole chain (loudly, via the
10547        // capture-site WARN) rather than capturing a subset the launch order cannot honor.
10548        let head_ok = |m: &MtpHead| match &m.ffn {
10549            crate::hybrid::Ffn::Dense { .. } => true,
10550            crate::hybrid::Ffn::Moe(mo) => mo.dev_exps.is_some(),
10551        };
10552        self.mtp.as_ref().map(&head_ok).unwrap_or(false) && self.mtp_extra.iter().all(head_ok)
10553    }
10554
10555    fn batched_serving_numeric_class(&self) -> bool {
10556        self.plan
10557            .trunk_operations()
10558            .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
10559    }
10560
10561    /// The family the MTP verify-graph default was measured on: GatedDeltaNet state layers
10562    /// (a `recur` mixer) together with a routed-MoE FFN — Ornith-1.5-35B-A3B and its kin. The
10563    /// server-side twin of this test is `model_forces_spec_replay` (GatedDeltaNet + MoeMlp);
10564    /// keeping the engine's own version structural rather than name-based means a new
10565    /// checkpoint of the same shape inherits the default, and a different shape does not.
10566    fn vgraph_family_default(&self) -> bool {
10567        let has_linear = self
10568            .layers
10569            .iter()
10570            .any(|l| matches!(l.mixer, Mixer::Linear(_)));
10571        let has_moe = self
10572            .layers
10573            .iter()
10574            .any(|l| matches!(l.ffn, crate::hybrid::Ffn::Moe(_)));
10575        has_linear && has_moe
10576    }
10577
10578    fn sliding_gated_moe_batch_program(&self) -> bool {
10579        self.uses_sliding_gated_moe_program()
10580    }
10581
10582    fn gemma_batch_program(&self) -> bool {
10583        self.uses_gemma_program()
10584    }
10585
10586    /// Reduced-matrix admission for increment 1. This deliberately does not change the PP-2
10587    /// serving policy: the worker calls it only after `MEMRA_SPEC_PIPE=1` and an explicit spec
10588    /// session already exist.
10589    pub fn spec_pipe_available(&self, e: &Engine) -> bool {
10590        if std::env::var("MEMRA_SPEC_PIPE").as_deref() != Ok("1")
10591            || !spec_devacc()
10592            || spec_replay_env_enabled()
10593            || spec_stream()
10594            || std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1")
10595            || std::env::var("MEMRA_SPEC_PMIN0").as_deref() == Ok("1")
10596            || std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1")
10597            || std::env::var("MEMRA_SPEC_PMIN")
10598                .ok()
10599                .and_then(|v| v.parse::<f32>().ok())
10600                .unwrap_or(0.0)
10601                > 0.0
10602            || self.is_gemma4_e4b()
10603            || self.gemma_batch_program()
10604            || self.mtp.is_none()
10605            || !self.mtp_extra.is_empty()
10606        {
10607            return false;
10608        }
10609        let Some(cuts) = crate::pp::pp_cuts(self.layers.len()) else {
10610            return false;
10611        };
10612        if cuts.len() != 3 || crate::pp::pp2_streams_off() || !crate::pp::spec_pp_on() {
10613            return false;
10614        }
10615        crate::pp::PpNRt::get(e)
10616            .map(|rt| rt.n_stages() == 2 && rt.cross_device())
10617            .unwrap_or(false)
10618    }
10619
10620    /// Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing
10621    /// `generate_spec_inner2` call stacks own all per-session round locals; only phase issue order
10622    /// changes. No callback is accepted in increment 1 — the worker publishes each completed burst.
10623    #[allow(clippy::too_many_arguments)]
10624    pub fn generate_spec_session_pair(
10625        &self,
10626        e: &Engine,
10627        sess_a: &mut SpecSession,
10628        max_new_a: usize,
10629        k_a: usize,
10630        sess_b: &mut SpecSession,
10631        max_new_b: usize,
10632        k_b: usize,
10633    ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn std::error::Error>>
10634    {
10635        if !self.spec_pipe_available(e) {
10636            return Err("two-session speculative pipeline is outside its reduced matrix".into());
10637        }
10638        if max_new_a == 0 || max_new_b == 0 || k_a == 0 || k_b == 0 {
10639            return Err(
10640                "two-session speculative pipeline requires non-empty positive-K bursts".into(),
10641            );
10642        }
10643        for sess in [&*sess_a, &*sess_b] {
10644            if sess.committed.is_empty()
10645                || sess.last_h.is_none()
10646                || (sess.next_pred.is_none() && sess.pending_tok.is_none())
10647            {
10648                return Err("two-session speculative pipeline requires warm continuations".into());
10649            }
10650        }
10651
10652        let graph_ok = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10653            && !spec_host_embd()
10654            && self.mtp_graph_capturable()
10655            && self.mtp_extra.is_empty()
10656            && !crate::model::full_prec_enabled();
10657        let graph_a = graph_ok && k_a + 2 < 96;
10658        let graph_b = graph_ok && k_b + 2 < 96;
10659        let was_tracking = e.ctx().is_event_tracking();
10660        if (graph_a || graph_b) && was_tracking {
10661            unsafe {
10662                e.ctx().disable_event_tracking();
10663            }
10664        }
10665
10666        static LOGGED: std::sync::Once = std::sync::Once::new();
10667        LOGGED.call_once(|| {
10668            eprintln!("[spec-pipe] two-session PP-2 continuation pipeline engaged");
10669        });
10670        let sync = std::sync::Arc::new(SpecPipeSync::new());
10671        let lane_a = SpecPipeLane {
10672            sync: sync.clone(),
10673            lane: 0,
10674        };
10675        let lane_b = SpecPipeLane { sync, lane: 1 };
10676        let mut sess_b_ptr = SpecPipeSessionPtr(sess_b as *mut SpecSession);
10677        let (result_a, result_b) = std::thread::scope(|scope| {
10678            let b = scope.spawn(move || {
10679                let mut finish = SpecPipeFinish::new(&lane_b);
10680                let sess_b = unsafe { sess_b_ptr.get_mut() };
10681                let result = e
10682                    .ctx()
10683                    .bind_to_thread()
10684                    .map_err(|err| err.to_string())
10685                    .and_then(|_| {
10686                        self.generate_spec_inner2(
10687                            e,
10688                            &[],
10689                            max_new_b,
10690                            k_b,
10691                            graph_b,
10692                            Some(sess_b),
10693                            None,
10694                            None,
10695                            None,
10696                            None,
10697                            Some(&lane_b),
10698                        )
10699                        .map_err(|err| err.to_string())
10700                    });
10701                finish.close(result.is_err());
10702                result
10703            });
10704            let mut finish = SpecPipeFinish::new(&lane_a);
10705            let result_a = self.generate_spec_inner2(
10706                e,
10707                &[],
10708                max_new_a,
10709                k_a,
10710                graph_a,
10711                Some(sess_a),
10712                None,
10713                None,
10714                None,
10715                None,
10716                Some(&lane_a),
10717            );
10718            finish.close(result_a.is_err());
10719            let result_b = b
10720                .join()
10721                .map_err(|_| "paired speculative session B panicked".to_string())
10722                .and_then(|r| r);
10723            (result_a, result_b)
10724        });
10725
10726        if (graph_a || graph_b) && was_tracking {
10727            unsafe {
10728                e.ctx().enable_event_tracking();
10729            }
10730        }
10731        let result_a = result_a?;
10732        let result_b = result_b.map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
10733        Ok((result_a, result_b))
10734    }
10735
10736    /// One spec-decode turn on a live session. `suffix` = the NEW tokens only (turn N+1's user
10737    /// message rendered through the chat template continuation). Returns (new tokens emitted,
10738    /// drafted, accepted); session.committed grows by suffix + emitted.
10739    pub fn generate_spec_session(
10740        &self,
10741        e: &Engine,
10742        sess: &mut SpecSession,
10743        suffix: &[u32],
10744        max_new: usize,
10745        k: usize,
10746    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10747        self.generate_spec_session_sampled(e, sess, suffix, max_new, k, None, None)
10748    }
10749
10750    /// Serve-path sampled spec: routes the burst through the rejection-sampling verify with
10751    /// per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy.
10752    /// Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact
10753    /// for the filtered target (feat/filtered-spec).
10754    ///
10755    /// `on_commit` (sse-cadence, 2026-08-05): called with each newly-emitted slice of the
10756    /// output — once right after the prime's first token, then once per round commit — so a
10757    /// streaming caller can flush text at round cadence instead of once per burst. The slices
10758    /// are disjoint, in order, and concatenate to exactly the returned token vec. Emission-
10759    /// timing only: token bytes, session state, and exactness are untouched.
10760    ///
10761    /// The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): `false` ends
10762    /// the burst at the current round boundary, exactly as if `max_new` had been reached —
10763    /// the caller's scheduler regains control without waiting the burst out. Burst size is
10764    /// content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns,
10765    /// never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream
10766    /// drains and the defensive tail flush can land with nothing new committed).
10767    #[allow(clippy::too_many_arguments)]
10768    pub fn generate_spec_session_sampled(
10769        &self,
10770        e: &Engine,
10771        sess: &mut SpecSession,
10772        suffix: &[u32],
10773        max_new: usize,
10774        k: usize,
10775        sampling: Option<SpecSampling>,
10776        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10777    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10778        self.generate_spec_session_sampled_prime_split(
10779            e, sess, suffix, max_new, k, sampling, None, on_commit,
10780        )
10781    }
10782
10783    /// Serve-only cold-prime segmentation twin. `prime_split` is the same stable boundary the
10784    /// plain worker would honor before entering its sub-floor tokenwise tail; warm continuations
10785    /// pass `None` and stay on the existing zero-prime path.
10786    #[allow(clippy::too_many_arguments)]
10787    pub fn generate_spec_session_sampled_prime_split(
10788        &self,
10789        e: &Engine,
10790        sess: &mut SpecSession,
10791        suffix: &[u32],
10792        max_new: usize,
10793        k: usize,
10794        sampling: Option<SpecSampling>,
10795        prime_split: Option<usize>,
10796        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10797    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10798        self.generate_spec_session_constrained_prime_split(
10799            e,
10800            sess,
10801            suffix,
10802            max_new,
10803            k,
10804            sampling,
10805            None,
10806            prime_split,
10807            on_commit,
10808        )
10809    }
10810
10811    /// `generate_spec_session_sampled` + GRAMMAR (constrained decoding, 2026-08-03): the
10812    /// hook truncates acceptance at the first grammar-illegal token AFTER the exactness
10813    /// verify (grammar is an extra rejection rule, ordering like the batched-verify twins)
10814    /// and replaces an illegal bonus with the MASKED argmax of the target's own verify
10815    /// column — token-identical to constrained plain greedy decode. GREEDY only (the
10816    /// worker routes sampled constrained to plain decode). Acceptance under tight grammars
10817    /// may drop (drafter is unconstrained); that is measured, not hidden.
10818    #[allow(clippy::too_many_arguments)]
10819    pub fn generate_spec_session_constrained(
10820        &self,
10821        e: &Engine,
10822        sess: &mut SpecSession,
10823        suffix: &[u32],
10824        max_new: usize,
10825        k: usize,
10826        sampling: Option<SpecSampling>,
10827        constraint: Option<&mut dyn SpecConstraint>,
10828        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10829    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10830        self.generate_spec_session_constrained_prime_split(
10831            e, sess, suffix, max_new, k, sampling, constraint, None, on_commit,
10832        )
10833    }
10834
10835    #[allow(clippy::too_many_arguments)]
10836    pub fn generate_spec_session_constrained_prime_split(
10837        &self,
10838        e: &Engine,
10839        sess: &mut SpecSession,
10840        suffix: &[u32],
10841        max_new: usize,
10842        k: usize,
10843        sampling: Option<SpecSampling>,
10844        constraint: Option<&mut dyn SpecConstraint>,
10845        prime_split: Option<usize>,
10846        on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10847    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10848        if constraint.is_some() && sampling.is_some_and(|s| s.temp > 0.0) {
10849            return Err(
10850                "constrained spec decode is greedy-only (worker routes sampled \
10851                        constrained to plain decode)"
10852                    .into(),
10853            );
10854        }
10855        // PENDING-CARRY entry flush: a carried bonus precedes any new suffix in the sequence,
10856        // so it must commit BEFORE the suffix primes; the sampled path doesn't carry (its
10857        // round-0 accept needs the commit pass's logits). Empty-suffix greedy bursts — the
10858        // serve continuation case — consume the carry in-loop with zero solo passes.
10859        if sess.pending_tok.is_some()
10860            && (!suffix.is_empty() || sampling.map_or(false, |s| s.temp > 0.0))
10861        {
10862            self.spec_flush_pending(e, sess, sampling)?;
10863        }
10864
10865        // FULL_PREC forces the EAGER draft: the graph capture would enclose cuBLASLt f32 GEMV
10866        // (the FloatBf16 else-branches) and a bf16_to_f32 dequant alloc — neither is stream-capture
10867        // safe. Eager rides matmul/matmul_decode_exact, which dequant FloatBf16 on use. (§item 2.)
10868        // Multi-head MTP (mtp_extra non-empty) no longer disqualifies: the chain captures
10869        // per-head graphs (lane/step37-draft-graph-serving-20260830, MEMRA_MTP_CHAIN_GRAPH).
10870        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10871            && !spec_host_embd()
10872            && self.mtp_graph_capturable()
10873            && k + 2 < 96
10874            && !crate::model::full_prec_enabled();
10875        let was_tracking = e.ctx().is_event_tracking();
10876        if graph_draft && was_tracking {
10877            unsafe {
10878                e.ctx().disable_event_tracking();
10879            }
10880        }
10881        let r = self.generate_spec_inner2(
10882            e,
10883            suffix,
10884            max_new,
10885            k,
10886            graph_draft,
10887            Some(sess),
10888            sampling,
10889            constraint,
10890            on_commit,
10891            prime_split,
10892            None,
10893        );
10894        if graph_draft && was_tracking {
10895            unsafe {
10896                e.ctx().enable_event_tracking();
10897            }
10898        }
10899        let (out, d, a) = r?;
10900        Ok((out, d, a))
10901    }
10902
10903    pub fn generate_spec(
10904        &self,
10905        e: &Engine,
10906        prompt: &[u32],
10907        max_new: usize,
10908        k: usize,
10909    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10910        if crate::pp::pp_cuts(self.layers.len()).is_some()
10911            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
10912        {
10913            return Err("pipeline rewrite is not qualified for speculative decode".into());
10914        }
10915        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::MtpSpec) {
10916            return Err("speculative rewrite is not qualified for this ModelPlan".into());
10917        }
10918        // FULL_PREC forces eager (see generate_spec_session note): CUDA graph capture cannot
10919        // enclose cuBLASLt f32 GEMV or the bf16_to_f32 dequant alloc the FloatBf16 path needs.
10920        // Multi-head MTP no longer disqualifies (chain graphs; see generate_spec_session).
10921        let graph_draft = std::env::var("MEMRA_SPEC_NOGRAPH").is_err()
10922            && !spec_host_embd()
10923            && self.mtp_graph_capturable()
10924            && k + 2 < 96
10925            && !crate::model::full_prec_enabled();
10926        if !graph_draft {
10927            return self.generate_spec_inner2(
10928                e, prompt, max_new, k, false, None, None, None, None, None, None,
10929            );
10930        }
10931        let was_tracking = e.ctx().is_event_tracking();
10932        if was_tracking {
10933            unsafe {
10934                e.ctx().disable_event_tracking();
10935            }
10936        }
10937        let r = self.generate_spec_inner2(
10938            e, prompt, max_new, k, true, None, None, None, None, None, None,
10939        );
10940        if was_tracking {
10941            unsafe {
10942                e.ctx().enable_event_tracking();
10943            }
10944        }
10945        r
10946    }
10947
10948    fn generate_spec_inner2(
10949        &self,
10950        e: &Engine,
10951        prompt: &[u32],
10952        max_new: usize,
10953        k: usize,
10954        graph_draft: bool,
10955        mut sess: Option<&mut SpecSession>,
10956        sampling: Option<SpecSampling>,
10957        mut constraint: Option<&mut dyn SpecConstraint>,
10958        mut on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>,
10959        prime_split: Option<usize>,
10960        pipe: Option<&SpecPipeLane>,
10961    ) -> Result<(Vec<u32>, usize, usize), Box<dyn std::error::Error>> {
10962        assert!(k >= 1, "k must be >= 1");
10963        if let Some(p) = pipe {
10964            p.setup_begin()?;
10965        }
10966        // sse-cadence flush cursor: everything in out[..flushed] has been handed to on_commit.
10967        let mut flushed = 0usize;
10968        // admission yield (2026-08-06): on_commit's continue-verdict; false = end the burst
10969        // at the next round boundary (same exit as max_new reached — the session tail runs).
10970        // Initialized by the unconditional post-prime flush below.
10971        let mut keep_going;
10972        let mtp = self
10973            .mtp
10974            .as_ref()
10975            .expect("generate_spec requires an MTP head (nextn_predict_layers>0)");
10976        let n_vocab = self.output.out_features();
10977        // FR-Spec: the draft head may be TRIMMED (fewer rows than n_vocab); the draft argmax runs
10978        // over the draft vocab and the winning index maps through d2t to a TARGET token id.
10979        // Everything downstream (verify/accept/commit) sees target ids only — exactness unchanged.
10980        let d_vocab = mtp
10981            .shared_head_head
10982            .as_ref()
10983            .unwrap_or(&self.output)
10984            .out_features();
10985        if !self.mtp_extra.is_empty() {
10986            if self.plan.draft_source != memra_gguf::model_plan::DraftSourcePlan::Embedded
10987                || self.plan.mtp_blocks.len() != self.mtp_head_count()
10988            {
10989                return Err(
10990                    "multi-head MTP requires one embedded canonical block per loaded head".into(),
10991                );
10992            }
10993            // TRIMMED chains (2026-08-27): every head must carry the SAME d2t — the ranking is
10994            // token-frequency and head-independent, and every downstream remap (per-step argmax,
10995            // stream pack, sampled d2t_dev) reads head 0's map, so equality is what makes that
10996            // single map correct for the whole chain. Mixed trimmed/untrimmed is refused.
10997            for (offset, head) in self.mtp_extra.iter().enumerate() {
10998                if head.d2t != mtp.d2t
10999                    || head
11000                        .shared_head_head
11001                        .as_ref()
11002                        .unwrap_or(&self.output)
11003                        .out_features()
11004                        != d_vocab
11005                {
11006                    return Err(format!(
11007                        "embedded MTP head {} has incompatible draft vocabulary",
11008                        offset + 1
11009                    )
11010                    .into());
11011                }
11012            }
11013            eprintln!(
11014                "[mtp-chain] heads={} policy=step-modulo prefix-replay kv=per-head",
11015                self.mtp_head_count()
11016            );
11017        }
11018        let n_embd = self.cfg.n_embd as usize;
11019        // SESSION MODE: reuse the live cache/scratch, prime only the suffix. `base` = tokens
11020        // already committed (their state is in the caches); 0 = fresh single-shot call.
11021        let session_mode = sess.is_some();
11022        let max_ctx = match sess.as_ref() {
11023            Some(s) => s.cache.max_ctx,
11024            None => prompt.len() + max_new + k + 8,
11025        };
11026        let mut own_cache;
11027        let mut own_scratch;
11028        // PREFIX-CACHE capture request threaded out of the session (lane/spec-prefix-cache):
11029        // (requested split, destination list). Single-shot per burst; fresh calls have none.
11030        let mut sess_capture: Option<(Option<usize>, &mut Vec<SpecBoundaryCapture>)> = None;
11031        // STABLE-BOUNDARY turn-checkpoint request (lane/frspec-multiturn-cache): ABSOLUTE
11032        // committed-length position; consumed one-shot like `capture_at`. None = legacy
11033        // prompt-end capture below.
11034        let mut ckpt_req: Option<usize> = None;
11035        let (
11036            cache,
11037            scratch,
11038            mut sess_tail,
11039            mut sess_draft_slot,
11040            mut sess_pending_slot,
11041            sess_ckpt_slot,
11042            sess_telem,
11043        ): (
11044            &mut Cache,
11045            &mut MtpScratch,
11046            Option<(
11047                &mut Vec<u32>,
11048                &mut Option<CudaSlice<f32>>,
11049                &mut Option<u32>,
11050                &mut u32,
11051                &mut u32,
11052            )>,
11053            Option<&mut Option<DraftGraphCtx>>,
11054            Option<&mut Option<u32>>,
11055            Option<&mut Option<SpecCheckpoint>>,
11056            Option<&SpecTelemetryCounters>,
11057        ) = match sess.take() {
11058            Some(sr) => {
11059                let SpecSession {
11060                    cache,
11061                    scratch,
11062                    committed,
11063                    last_h,
11064                    next_pred,
11065                    sctr: s_sctr,
11066                    uctr: s_uctr,
11067                    draft_ctx,
11068                    pending_tok,
11069                    turn_ckpt,
11070                    telem,
11071                    capture_at,
11072                    boundary_captures,
11073                    ckpt_at,
11074                } = sr;
11075                sess_capture = Some((capture_at.take(), boundary_captures));
11076                ckpt_req = ckpt_at.take();
11077                (
11078                    cache,
11079                    scratch,
11080                    Some((committed, last_h, next_pred, s_sctr, s_uctr)),
11081                    Some(draft_ctx),
11082                    Some(pending_tok),
11083                    Some(turn_ckpt),
11084                    Some(telem),
11085                )
11086            }
11087            None => {
11088                // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut =
11089                // `Cache::new` verbatim.
11090                own_cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, max_ctx)?;
11091                // Persistent scratch = max_ctx rows (~2KB/token quantized).
11092                own_scratch = self.new_mtp_scratch(e, max_ctx)?;
11093                (
11094                    &mut own_cache,
11095                    &mut own_scratch,
11096                    None,
11097                    None,
11098                    None,
11099                    None,
11100                    None,
11101                )
11102            }
11103        };
11104        if scratch.plane_count() != self.mtp_head_count() {
11105            return Err(format!(
11106                "MTP scratch/head count mismatch ({}/{})",
11107                scratch.plane_count(),
11108                self.mtp_head_count()
11109            )
11110            .into());
11111        }
11112        let base = cache.pos;
11113        // PENDING-CARRY consume (2026-08-01): a carried bonus reaches here only on the
11114        // empty-suffix GREEDY continuation path (generate_spec_session_sampled flushed every
11115        // other case). It enters the round loop as round-0's pending — verify col 0 — exactly
11116        // like a mid-burst full-accept boundary: no init feed, no tail commit pass.
11117        let carried_pending: Option<u32> = sess_pending_slot.as_mut().and_then(|s| s.take());
11118        // PERSISTENT DRAFT KV (the only mode since 2026-07-08 — the legacy round-local scratch,
11119        // MEMRA_SPEC_KVLOCAL, measured -35 acceptance pts on the 27B p3 sweep and was removed;
11120        // acceptance-only — exactness is verify's job either way).
11121        // HIDDEN-PAIRING CONVENTION (DEFAULT = predecessor-row, 2026-07-04 — the 27B acceptance
11122        // unlock, +16pts): the MTP head is TRAINED on rows pairing token x_p with the trunk
11123        // hidden of its PREDECESSOR h_{p-1} (the reference engine's mtp_update shifts the target
11124        // hiddens right by one; its draft step 0 feeds (id_last, TRUE hidden of the row id_last
11125        // was sampled from)). memra's historical convention paired SAME-ROW (x_p, h_p) in the fill
11126        // and seeded chain step 0 through an extra MTP pass on a duplicated token (the
11127        // pseudo-seed) — measured 27B p2 K=3 acceptance 0.569 vs 0.731, p3 0.445 vs 0.63+, and
11128        // the chain steps j>=1 were already predecessor-shaped, so ONLY the fill + step-0 seed
11129        // move. The fill shifts by one and the chain seeds from the predecessor's true hidden
11130        // DIRECTLY (vh_seed / vx[j-1]) — the pseudo pass disappears (one MTP-block pass saved
11131        // per round on top of the acceptance win). Draft-quality-only: exactness stays the
11132        // verify's job either way. (The legacy same-row pairing seam, MEMRA_SPEC_HSAME, and its
11133        // pseudo-seed passes were removed 2026-07-08 — predecessor pairing won by +16 acc pts;
11134        // the legacy round-local scratch, MEMRA_SPEC_KVLOCAL, went with it.)
11135        // REPLAY-FREE PARTIAL ACCEPT (default, 2026-07-03): partial rounds keep the verify's own
11136        // bit-identical committed-prefix state (KV truncate + recur rebuild from the VerifyCkpt)
11137        // and leave the bonus PENDING — no duplicate trunk pass (profiled ~0.54 extra full weight
11138        // reads/round at long ctx). MEMRA_SPEC_REPLAY=1 restores the legacy rollback+replay (A/B
11139        // + fallback seam).
11140        // Qwen35-MoE replay pin LIFTED (lane/draftcost-moe, 2026-08-20). The pin's stated
11141        // bar — the retained verify-state commit proven equivalent to sequential serving —
11142        // was waiting on this arch running the serving batched verify class, which the
11143        // t-parallel admission (this lane, increment 1) provided: the VerifyCkpt the
11144        // replay-free commit consumes is now produced by the SAME serving-class verify that
11145        // qualified dense qwen35 on 2026-08-15 (where the per-round duplicate replay
11146        // measured 69 -> 30 tok/s). Qualification receipts (run-spec K=1..8 both arms,
11147        // 8-prompt replay-vs-replay-free canary, long-prompt cell):
11148        // research/draftcost-moe-20260820/RECEIPTS.md. MEMRA_SPEC_REPLAY=1 stays the
11149        // rollback + A/B seam.
11150        let spec_replay = spec_replay_env_enabled();
11151        if constraint.is_some() && spec_replay {
11152            return Err(
11153                "constrained spec decode does not support MEMRA_SPEC_REPLAY=1 \
11154                        (legacy replay commits an unmasked bonus)"
11155                    .into(),
11156            );
11157        }
11158        // TRUE-HIDDEN REFRESH (default in persistent-draft-KV mode): every round overwrites the
11159        // committed positions' scratch entries from the verify's exact hiddens (mtp_kv_fill batch)
11160        // instead of keeping chain-approximate entries. MEMRA_SPEC_NOREFRESH=1 = legacy (A/B seam).
11161        let refresh = std::env::var("MEMRA_SPEC_NOREFRESH").is_err();
11162        if !refresh && !self.mtp_extra.is_empty() {
11163            return Err("multi-head MTP requires exact accepted-prefix refresh".into());
11164        }
11165
11166        // prime: BATCHED cache prime (prime_cache — the measured #1 e2e gap: tokenwise primed at
11167        // ~102/38 tok/s vs the engine's ~2000-5900 tok/s batched prefill). prime_cache returns the
11168        // full pre-output_norm hidden stack [T, n_embd], which IS prompt_h (the persistent-draft-KV
11169        // mtp_kv_fill input) — no per-token collection needed. Prompts below PRIME_MIN_T, and
11170        // MEMRA_PRIME_TOKENWISE=1, and frozen Hy3 CPU/GPU expert splits take the tokenwise
11171        // decode_step_h loop. The latter avoids transient GPU staging of the spilled expert bank.
11172        // EMPTY-SUFFIX CONTINUATION (serve bursts): a session turn with NO new tokens resumes
11173        // generation exactly where the last turn stopped — no prime at all. The stashed
11174        // `next_pred` plays prime_logits' role: it is the token produced from the logits after
11175        // committed.last() by the same rule this entry applies to a cold prime's last row —
11176        // an argmax when greedy, a `sample_boundary_token` draw when sampled (the burst tail,
11177        // or `spec_session_from_restored` for a converted prefix-cache hit, did the drawing
11178        // where the sampler and the session's Philox counters were live). `last_h` seeds the
11179        // predecessor pairing below. Fresh calls and non-empty suffixes take the normal path.
11180        let continuation = prompt.is_empty();
11181        if continuation {
11182            assert!(session_mode, "empty prompt requires a session");
11183            assert!(
11184                sess_tail
11185                    .as_ref()
11186                    .map_or(false, |(c, lh, np, _, _)| !c.is_empty()
11187                        && lh.is_some()
11188                        && (np.is_some() || carried_pending.is_some())),
11189                "empty-suffix continuation needs a primed session (committed + last_h + next_pred|pending)"
11190            );
11191        }
11192        let mut prime_logits;
11193        let mut prompt_h: Option<CudaSlice<f32>> = None;
11194        let t_prime = std::time::Instant::now();
11195        let batched_prime = !continuation
11196            && prompt.len() >= crate::hybrid_forward::PRIME_MIN_T
11197            && std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11198            && !e.frozen_cpu_experts_prefer_tokenwise_prime();
11199        let prime_split = prime_split.filter(|&split| split > 0 && split < prompt.len());
11200        if prime_split.is_some() && continuation {
11201            return Err("spec prime split requires a non-empty prime".into());
11202        }
11203        // STABLE-BOUNDARY TURN CHECKPOINT stop (lane/frspec-multiturn-cache, 2026-08-21):
11204        // the worker's `ckpt_at` request, ABSOLUTE -> prompt-relative. On WARM bursts
11205        // (base != 0, an affinity-rewound or pool-resumed session priming its own delta)
11206        // this is the only stop; on COLD bursts it usually coincides with `prime_split`
11207        // (both are the plain tier's stable pre-generation boundary). A boundary the prime
11208        // cannot honor (outside this prime's range) silently drops the capture — the
11209        // turn_ckpt convention: the next turn re-primes in full, never a wrong resume.
11210        let ckpt_rel = if continuation {
11211            None
11212        } else {
11213            ckpt_req
11214                .and_then(|abs| abs.checked_sub(base))
11215                .filter(|&r| r > 0 && r < prompt.len())
11216        };
11217        // Prime stops, ordered: each is a boundary the prime halts at so the in-place GDN
11218        // conv/ssm state can be snapshotted there (the only moment it exists). One stop =
11219        // the legacy single-split program, byte-for-byte.
11220        let mut stops: Vec<usize> = Vec::new();
11221        for b in [prime_split, ckpt_rel].into_iter().flatten() {
11222            if !stops.contains(&b) {
11223                stops.push(b);
11224            }
11225        }
11226        stops.sort_unstable();
11227        // Captured at the ckpt stop, installed into the session slot post-prime (replacing
11228        // the legacy prompt-end capture). Some(None) = capture attempted and failed -> the
11229        // slot is cleared (a stale checkpoint would rewind to the WRONG boundary).
11230        let mut ckpt_early: Option<Option<SpecCheckpoint>> = None;
11231        if continuation {
11232            prime_logits = Vec::new();
11233        } else if !stops.is_empty() {
11234            if let Some(&first) = stops.first() {
11235                if prime_split == Some(first) && first < crate::hybrid_forward::PRIME_MIN_T {
11236                    return Err(format!(
11237                        "spec prime split {first} is below PRIME_MIN_T {}",
11238                        crate::hybrid_forward::PRIME_MIN_T,
11239                    )
11240                    .into());
11241                }
11242            }
11243            // Mirror the plain worker's boundary stops exactly. Each segment is a
11244            // request-level prime (`queued_after` keeps Step35 arm selection independent of
11245            // the stops — tick-seg law); a segment below PRIME_MIN_T (and the final tail
11246            // under MEMRA_PRIME_TOKENWISE) takes the same eager tokenwise continuation as
11247            // prefill_tick. Retain every hidden row so the draft scratch fill remains one
11248            // coherent prompt.
11249            let mut h_all = e.uninit(prompt.len() * n_embd)?;
11250            prime_logits = Vec::new();
11251            let mut prev = 0usize;
11252            for seg_end in stops.iter().copied().chain(std::iter::once(prompt.len())) {
11253                if seg_end <= prev {
11254                    continue;
11255                }
11256                let seg = &prompt[prev..seg_end];
11257                let is_final = seg_end == prompt.len();
11258                let batched_seg = seg.len() >= crate::hybrid_forward::PRIME_MIN_T
11259                    && (!is_final
11260                        || (std::env::var("MEMRA_PRIME_TOKENWISE").is_err()
11261                            && !e.frozen_cpu_experts_prefer_tokenwise_prime()));
11262                if batched_seg {
11263                    let (l, _, h_seg) =
11264                        self.prime_cache(e, seg, &mut *cache, prompt.len() - seg_end)?;
11265                    e.copy_into(&mut h_all, prev * n_embd, &h_seg, seg.len() * n_embd)?;
11266                    prime_logits = l;
11267                } else {
11268                    for (i, &tok) in seg.iter().enumerate() {
11269                        let (l, h) = self.decode_step_h(e, tok, &mut *cache)?;
11270                        e.copy_into(&mut h_all, (prev + i) * n_embd, &h, n_embd)?;
11271                        prime_logits = l;
11272                    }
11273                }
11274                prev = seg_end;
11275                if is_final {
11276                    break;
11277                }
11278                debug_assert_eq!(cache.pos, base + seg_end, "prime stop landed off boundary");
11279                // PREFIX-CACHE BOUNDARY CAPTURE (lane/spec-prefix-cache): the GDN conv/ssm
11280                // states are about to be advanced in place by the next segment, so this is
11281                // the ONLY moment the boundary's recurrent state exists. Capture iff the
11282                // worker requested exactly this stop (cold sessions only — `capture_at` is
11283                // never armed warm). A failed snapshot is silent (turn_ckpt convention) —
11284                // publication is an optimization, never a correctness dependency.
11285                if base == 0 {
11286                    if let Some((requested, slot)) = sess_capture.as_mut() {
11287                        // Publish at the requested miss-LCP stop (the shared-prefix class)
11288                        // AND at the stable-boundary stop (the next-turn re-render class,
11289                        // lane/frspec-multiturn-cache) — the same boundary set the plain
11290                        // prefill tick learns. Without the second entry, the turn after a
11291                        // cold re-park could only hit the OLDER lcp entry (the measured
11292                        // one-turn transient: t3 restored 607 of 24122 while the plain arm
11293                        // rewound to 15222). Dedupe is the worker sweep's has_key.
11294                        if *requested == Some(seg_end) || ckpt_rel == Some(seg_end) {
11295                            if let Ok(snap) = cache.snapshot(e) {
11296                                slot.push(SpecBoundaryCapture {
11297                                    snap,
11298                                    pos: seg_end,
11299                                    logits: prime_logits.clone(),
11300                                    // rows [0..seg_end) of h_all are primed — the following
11301                                    // segments append, never overwrite.
11302                                    last_h: capture_boundary_hidden(e, &h_all, seg_end, n_embd),
11303                                });
11304                            }
11305                        }
11306                    }
11307                }
11308                // SESSION-AFFINITY TURN CHECKPOINT at the STABLE boundary (see `ckpt_at`):
11309                // same snapshot mechanics, installed post-prime in place of the prompt-end
11310                // capture the re-render class always diverged below.
11311                if ckpt_rel == Some(seg_end) {
11312                    let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11313                        e.uninit(n_embd).and_then(|mut a| {
11314                            e.copy_view_into(
11315                                &mut a,
11316                                0,
11317                                &h_all.slice((seg_end - 1) * n_embd..seg_end * n_embd),
11318                                n_embd,
11319                            )?;
11320                            Ok(a)
11321                        });
11322                    ckpt_early = Some(match (cache.snapshot(e), anchor) {
11323                        (Ok(snap), Ok(last_h)) => Some(SpecCheckpoint {
11324                            snap,
11325                            pos: base + seg_end,
11326                            last_h,
11327                        }),
11328                        _ => None,
11329                    });
11330                }
11331            }
11332            if std::env::var("MEMRA_SPEC_STATS").as_deref() == Ok("1") {
11333                eprintln!(
11334                    "[spec-prime] stops={stops:?} tail={}",
11335                    prompt.len() - stops.last().copied().unwrap_or(0)
11336                );
11337            }
11338            prompt_h = Some(h_all);
11339        } else if batched_prime {
11340            let (l, _h_seed, hiddens) = self.prime_cache(e, prompt, &mut *cache, 0)?;
11341            prime_logits = l;
11342            prompt_h = Some(hiddens);
11343        } else {
11344            prime_logits = Vec::new();
11345            prompt_h = Some(e.uninit(prompt.len() * n_embd)?);
11346            for (i, &tok) in prompt.iter().enumerate() {
11347                let (l, h) = self.spec_target_step_h(e, tok, &mut *cache)?;
11348                if let Some(ph) = prompt_h.as_mut() {
11349                    e.copy_into(ph, i * n_embd, &h, n_embd)?;
11350                }
11351                prime_logits = l;
11352            }
11353        }
11354        e.stream().synchronize()?;
11355        // PREFIX-CACHE SEED CAPTURE (lane/spec-prefix-cache): boundary == prompt end (the seed
11356        // case — no shared-prefix split, publish the whole prompt). The prime just finished, so
11357        // cache.pos == base + prompt.len() and the recurrent state IS the boundary state;
11358        // prime_logits are the boundary logits. Cold sessions only (base == 0) — same law as
11359        // prime_split. The mid-prompt capture above already consumed the request if it matched.
11360        if !continuation && base == 0 {
11361            if let Some((requested, slot)) = sess_capture.as_mut() {
11362                if *requested == Some(prompt.len()) && slot.is_empty() {
11363                    debug_assert_eq!(cache.pos, prompt.len(), "seed capture off prompt end");
11364                    if let Ok(snap) = cache.snapshot(e) {
11365                        slot.push(SpecBoundaryCapture {
11366                            snap,
11367                            pos: prompt.len(),
11368                            logits: prime_logits.clone(),
11369                            last_h: prompt_h
11370                                .as_ref()
11371                                .map(|ph| capture_boundary_hidden(e, ph, prompt.len(), n_embd))
11372                                .unwrap_or_default(),
11373                        });
11374                    }
11375                }
11376            }
11377        }
11378        // Harness timing contract (see crate::PRIME_NANOS): gen-only throughput without the
11379        // prime-subtraction hack.
11380        crate::PRIME_NANOS.store(
11381            t_prime.elapsed().as_nanos() as u64,
11382            std::sync::atomic::Ordering::Relaxed,
11383        );
11384
11385        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
11386        // Resident table is fastest when it fits. Large spill deployments can preserve that HBM
11387        // for expert-cache slots and gather only the exact rows needed by MTP/verify from host.
11388        let host_embd = spec_host_embd();
11389        let embd_gpu = if host_embd {
11390            None
11391        } else {
11392            Some(
11393                self.embd_gpu
11394                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
11395            )
11396        };
11397        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
11398        if host_embd {
11399            eprintln!(
11400                "[spec] host-row embedding: {} bytes kept off HBM",
11401                self.embd.raw.len()
11402            );
11403        }
11404        let mut out: Vec<u32> = Vec::with_capacity(max_new);
11405        let mut total_drafted = 0usize;
11406        let mut total_accepted = 0usize;
11407
11408        // --- SAMPLER FIRST (lane/sampled-spec-quality, 2026-08-19) ---
11409        // The sampler config, the session's Philox counters and the penalty window are parsed
11410        // HERE, above the boundary-token selection, because the boundary token must be drawn
11411        // from the sampler the request asked for. Pre-lane this block sat ~50 lines BELOW the
11412        // selection, which is the whole mechanical reason the boundary token was an argmax:
11413        // the sampler state was not in scope yet. Nothing here depends on the round loop, so
11414        // moving it up is a pure reordering for greedy (`sampled == false` ⇒ every branch
11415        // below takes the argmax path it always took).
11416        // --- SAMPLED SPEC (MEMRA_SPEC_TEMP>0, research/sampled-spec-impl-map.md): rejection-
11417        // sampling verify (Leviathan/Chen) — accept draft x at u < p(x)/q(x), resample from
11418        // norm(max(0,p-q)) on reject, bonus sampled from p on full accept. Counter-based Philox
11419        // everywhere (seed, event) -> reproducible. temp==0/unset = the greedy path, untouched.
11420        let sp = sampling.unwrap_or_else(|| SpecSampling {
11421            temp: std::env::var("MEMRA_SPEC_TEMP")
11422                .ok()
11423                .and_then(|v| v.parse().ok())
11424                .unwrap_or(0.0),
11425            seed: std::env::var("MEMRA_SEED")
11426                .ok()
11427                .and_then(|v| v.parse().ok())
11428                .unwrap_or(42),
11429            top_k: std::env::var("MEMRA_TOP_K")
11430                .ok()
11431                .and_then(|v| v.parse().ok())
11432                .unwrap_or(0),
11433            top_p: std::env::var("MEMRA_TOP_P")
11434                .ok()
11435                .and_then(|v| v.parse().ok())
11436                .unwrap_or(1.0),
11437            min_p: std::env::var("MEMRA_MIN_P")
11438                .ok()
11439                .and_then(|v| v.parse().ok())
11440                .unwrap_or(0.0),
11441            penalty_last_n: std::env::var("MEMRA_PENALTY_LAST_N")
11442                .ok()
11443                .and_then(|v| v.parse().ok())
11444                .unwrap_or(0),
11445            penalty_repeat: std::env::var("MEMRA_PENALTY_REPEAT")
11446                .ok()
11447                .and_then(|v| v.parse().ok())
11448                .unwrap_or(1.0),
11449            penalty_freq: std::env::var("MEMRA_PENALTY_FREQ")
11450                .ok()
11451                .and_then(|v| v.parse().ok())
11452                .unwrap_or(0.0),
11453            penalty_present: std::env::var("MEMRA_PENALTY_PRESENT")
11454                .ok()
11455                .and_then(|v| v.parse().ok())
11456                .unwrap_or(0.0),
11457        });
11458        let (sp_temp, sp_seed) = (sp.temp, sp.seed);
11459        let sampled = sp_temp > 0.0;
11460        // Counters resume from the session (burst continuity: randomness must never repeat
11461        // across generate_spec_session calls); one-shot callers start at (0,0). Read through
11462        // sess_tail — `sess` was take()n into it above, so sess.as_ref() here is always None.
11463        let mut sctr: u32 = sess_tail.as_ref().map(|(_, _, _, s, _)| **s).unwrap_or(0);
11464        let mut uctr: u32 = sess_tail.as_ref().map(|(_, _, _, _, u)| **u).unwrap_or(0);
11465        // Penalties (v2.1): applied to COPIES of q rows and p columns symmetrically (exactness
11466        // for the penalized+filtered target). History = generated tokens, host-tracked window.
11467        let pen_on = sampled
11468            && sp.penalty_last_n > 0
11469            && (sp.penalty_repeat != 1.0 || sp.penalty_freq != 0.0 || sp.penalty_present != 0.0);
11470        // SESSION-SPANNING PENALTY WINDOW (Item 2). Pre-lane this was
11471        // `prompt.iter().rev().take(64).rev()` — the BURST's suffix slice — so a continuation
11472        // burst (the majority of a stream's tokens, and ALL of a converted cache hit's) started
11473        // with an EMPTY penalty history and the client's repetition/frequency/presence penalties
11474        // silently reset at every burst boundary. The window now spans `committed ++ prompt`,
11475        // which is what the API contract says and what the plain sampler's own `history` does.
11476        // Byte-identical to the pre-lane seed for a cold turn-1 burst at the default window.
11477        let mut pen_hist: Vec<u32> = if pen_on {
11478            let sess_hist: &[u32] = if spec_pen_session_on() {
11479                sess_tail
11480                    .as_ref()
11481                    .map(|(c, ..)| c.as_slice())
11482                    .unwrap_or(&[])
11483            } else {
11484                &[] // MEMRA_SPEC_PEN_SESSION=0: pre-lane burst-local window
11485            };
11486            pen_window_seed(sess_hist, prompt, sp.penalty_last_n)
11487        } else {
11488            Vec::new()
11489        };
11490        // First generated token = the BOUNDARY token: greedy takes the argmax of the prompt's
11491        // last logits (== greedy's first token, byte-contract); SAMPLED draws it from the
11492        // request's own filtered/penalized target through the session's Philox stream
11493        // (`sample_boundary_token`, lane/sampled-spec-quality Item 1 — pre-lane this was an
11494        // argmax in both regimes, so ~1 token per burst of a sampled stream was greedy).
11495        // Emit it, then FEED it to establish the loop invariant below.
11496        // PENDING-CARRY: the carried bonus was already emitted by the LAST burst — it becomes
11497        // last_token WITHOUT re-emission, and round 0 consumes it as pending (no init feed).
11498        // CONSTRAINED entry rules: the first emitted token is the MASKED argmax of the
11499        // prompt's last logits (plain constrained-greedy identity); a continuation without
11500        // a carried pending would emit an UNMASKED stashed next_pred — refused loudly (the
11501        // worker never resumes constrained sessions from the pool, so this cannot fire).
11502        if let Some(c) = constraint.as_deref_mut() {
11503            if continuation && carried_pending.is_none() {
11504                return Err("constrained spec continuation requires a carried pending \
11505                            (pool resume is unconstrained-only)"
11506                    .into());
11507            }
11508            if !continuation {
11509                c.mask_logits(&mut prime_logits)
11510                    .map_err(|e2| format!("constraint: {e2}"))?;
11511            }
11512        }
11513        let mut last_token = if let Some(b) = carried_pending {
11514            b
11515        } else if continuation {
11516            // A continuation's boundary token was DRAWN by the burst that stashed it (the
11517            // session tail below), or by `spec_session_from_restored` for a converted
11518            // prefix-cache hit — in both cases from the correct logits row with this same
11519            // session's Philox stream, which is why it can be consumed here as-is.
11520            sess_tail.as_ref().unwrap().2.unwrap()
11521        } else if sampled && constraint.is_none() && spec_sampled_boundary_on() {
11522            sample_boundary_token(e, &prime_logits, &sp, &pen_hist, &mut sctr, "cold-prime")?
11523        } else {
11524            // greedy (byte contract), the rollback door, or constrained (masked-argmax
11525            // identity — the worker routes sampled+constrained to the plain path, and this
11526            // function refuses the combination outright above).
11527            argmax(&prime_logits) as u32
11528        };
11529        if pen_on {
11530            // The boundary token is a GENERATED token: the plain sampler `accept()`s every
11531            // emitted token into its penalty history, and pre-lane the burst's first token
11532            // was invisible to penalties forever (never pushed, and never in `committed`
11533            // until this burst's tail). Covers the carry/continuation seeds too — neither is
11534            // in `committed` yet.
11535            pen_hist.push(last_token);
11536        }
11537        if carried_pending.is_none() {
11538            out.push(last_token);
11539            // grammar advances with every emitted token (carried pendings were consumed
11540            // by the burst that emitted them).
11541            if let Some(c) = constraint.as_deref_mut() {
11542                c.consume(last_token)
11543                    .map_err(|e2| format!("constraint: {e2}"))?;
11544            }
11545        }
11546        if continuation {
11547            // draft-KV invariant: entries [0..base) are the session's exact fills; truncate any
11548            // overhang so the chain's first append lands at slot base (== committed.len()).
11549            scratch.set_len(e, base)?;
11550        }
11551        // sse-cadence: hand the caller every not-yet-flushed token (disjoint in-order slices
11552        // concatenating to the full `out`). Called after the prime's first token and after each
11553        // round commit — emission timing only, token bytes untouched. The slice may be EMPTY
11554        // (poll-only boundary: zero-round folds commit nothing new); returns the caller's
11555        // continue-verdict (admission yield, 2026-08-06) — false ends the burst at this round.
11556        fn flush_commit(
11557            cb: &mut Option<&mut dyn FnMut(&[u32]) -> bool>,
11558            out: &[u32],
11559            flushed: &mut usize,
11560        ) -> bool {
11561            if let Some(f) = cb.as_mut() {
11562                let keep = f(&out[*flushed..]);
11563                *flushed = out.len();
11564                keep
11565            } else {
11566                true
11567            }
11568        }
11569        keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
11570        // INVARIANT at loop top: `last_token` is the most-recently-committed/emitted token, its
11571        // KV+recur state IS in `cache` (cache.pos = position right AFTER last_token), `last_pred`
11572        // is the greedy ARGMAX of the logits that predict the token FOLLOWING last_token, and
11573        // `h_seed` = last_token's pre-output_norm hidden. Establish it by feeding last_token once
11574        // (mirrors plain greedy). DEVICE-ARGMAX lever: the accept walk only ever consumes the
11575        // argmax of those logits — never the full vector — so a host u32 replaces the Vec<f32>.
11576        // Trimmed heads: q lives on the trimmed vocab; accept gathers use the TRIMMED index and
11577        // the residual scatters q into target-id space (q=-inf off-trim — the head cannot propose
11578        // those, so their residual mass is p(x), correct by construction).
11579        let d2t_dev: Option<CudaSlice<u32>> = if sampled || crate::spec::spec_stream() {
11580            match &mtp.d2t {
11581                Some(map) => Some(e.htod_u32_v(map)?),
11582                None => None,
11583            }
11584        } else {
11585            None
11586        };
11587        let mut q_full_buf: Option<CudaSlice<f32>> = None;
11588        // host Philox4x32-10 accept-test uniforms: module fn `host_u01` (shared with the
11589        // dspark sampled-admission walk); byte-identical to the closure it replaces.
11590        let mut draft_logits: Vec<CudaSlice<f32>> = Vec::new(); // retained head logits (q), per slot
11591        let mut draft_stats: Vec<(f32, f32, f32)> = Vec::new(); // (row_max, th_e, z_e) per slot
11592        let mut perturb_buf: Option<CudaSlice<f32>> = None; // gumbel scratch (max(n_vocab,d_vocab))
11593        let mut sample_tok = e.alloc_u32_zeroed(1)?; // residual/bonus sample out
11594        let mut col_buf: Option<CudaSlice<f32>> = None; // materialized verify column
11595        let mut pen_hist_d: Option<CudaSlice<u32>> = None;
11596        let mut pcol_buf: Option<CudaSlice<f32>> = None; // penalized p-column scratch
11597        // MEMRA_SPEC_SETUP_TRACE=1 (diagnostics): per-call wall decomposition of the burst
11598        // SETUP + TAIL segments (the round loop's internals are MEMRA_SPEC_PHASE's job) —
11599        // built to pin the serve per-burst fixed cost (research/spec-serving-20260801).
11600        let setup_trace = std::env::var("MEMRA_SPEC_SETUP_TRACE").as_deref() == Ok("1");
11601        let t_ent = std::time::Instant::now();
11602
11603        // SESSION-AFFINITY TURN CHECKPOINT (lane/session-affinity, 2026-08-05): capture the
11604        // PROMPT-END boundary state so a LATER turn can rewind here and re-prime only its own
11605        // delta instead of the whole conversation. See `SpecCheckpoint` for why this boundary is
11606        // the one that matters (a history-rewriting client mutates what the session GENERATED,
11607        // so the next turn's prompt agrees with this one up to exactly here).
11608        //
11609        // WHERE — AND WHY THIS EXACT LINE. Right after the trunk prime, BEFORE the init feed
11610        // (`decode_step_h(last_token)`) and before round 0: the last instant at which the caches
11611        // hold exactly `base + prompt.len()` rows and nothing generated.
11612        //
11613        // This was WRONG in the first cut of this lane: the capture sat after the draft-KV fill,
11614        // which is also after the init feed, so `cache.pos` was `base + prompt.len() + 1` — the
11615        // boundary included the FIRST GENERATED TOKEN. That token is the first thing inside the
11616        // `<think>` block the client strips, so every later turn's diff diverged exactly one
11617        // token below the checkpoint and affinity declined 100% of the time. Measured on the
11618        // owner regime: "history diverged at 12233 of checkpoint 12234". The off-by-one made the
11619        // whole mechanism inert while looking, from the outside, like a working
11620        // correctness-declines-safely path — hence the decline log carries the offsets.
11621        //
11622        // The full-attn planes are `len`-truncatable so the snapshot copies only the GDN conv/ssm
11623        // state (the reason a spec session could not rewind before). The draft scratch needs no
11624        // copy: rows below the boundary are rewritten by the next turn's own fill.
11625        //
11626        // WHEN: non-empty prime only. An empty-suffix continuation burst adds no prompt boundary
11627        // (its "prompt end" IS the previous checkpoint's, already held), so it keeps the existing
11628        // checkpoint rather than replacing it with a strictly worse one.
11629        //
11630        // FAILURE IS SILENT BY DESIGN: on a VRAM-tight rig the snapshot alloc can fail. That
11631        // costs the NEXT turn its rewind (it re-primes fully, today's behavior) and must never
11632        // fail the burst that is already running — so the error is swallowed, loud only under
11633        // MEMRA_DEBUG_SPEC.
11634        //
11635        // STABLE-BOUNDARY OVERRIDE (lane/frspec-multiturn-cache, 2026-08-21): the prompt-end
11636        // posture above was DISPROVED for the think-posture template class — the prompt's own
11637        // tail is the live generation header (`<|im_start|>assistant\n<think>\n`) that the
11638        // next turn's re-render replaces, so the diff diverged a couple tokens BELOW the
11639        // checkpoint and affinity declined 100% of multi-turn agent traffic (the same class
11640        // the plain tier fixed on 2026-08-09 via `plain_checkpoint_boundary`; the port to the
11641        // spec tier is this lane). When the worker armed `ckpt_at`, the capture happened at
11642        // that stop inside the prime above (`ckpt_early`) and is installed here instead;
11643        // capture-attempted-but-failed clears the slot exactly like the legacy arm.
11644        if let Some(slot) = sess_ckpt_slot {
11645            if let Some(early) = ckpt_early {
11646                if early.is_none() && std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11647                    eprintln!(
11648                        "[spec] stable-boundary turn checkpoint skipped; \
11649                               next turn re-primes in full"
11650                    );
11651                }
11652                *slot = early;
11653            } else if !continuation {
11654                let pos = cache.pos;
11655                debug_assert_eq!(
11656                    pos,
11657                    base + prompt.len(),
11658                    "turn checkpoint must sit at the prompt end, before the init feed"
11659                );
11660                let anchor: Result<CudaSlice<f32>, Box<dyn std::error::Error>> =
11661                    if let Some(ph) = &prompt_h {
11662                        // hidden of the LAST primed row = the predecessor anchor at this
11663                        // boundary (exactly what a fresh prime of committed[..pos] leaves in
11664                        // last_h, and what the next prime's fill reads for its first row).
11665                        let np = prompt.len();
11666                        e.uninit(n_embd).and_then(|mut a| {
11667                            e.copy_view_into(
11668                                &mut a,
11669                                0,
11670                                &ph.slice((np - 1) * n_embd..np * n_embd),
11671                                n_embd,
11672                            )?;
11673                            Ok(a)
11674                        })
11675                    } else {
11676                        Err("no prompt hiddens".into())
11677                    };
11678                match (cache.snapshot(e), anchor) {
11679                    (Ok(snap), Ok(last_h)) => {
11680                        *slot = Some(SpecCheckpoint { snap, pos, last_h });
11681                    }
11682                    (s, a) => {
11683                        *slot = None; // a stale checkpoint would rewind to the WRONG boundary
11684                        if std::env::var("MEMRA_DEBUG_SPEC").is_ok() {
11685                            let err = s
11686                                .err()
11687                                .map(|e| e.to_string())
11688                                .or_else(|| a.err().map(|e| e.to_string()))
11689                                .unwrap_or_default();
11690                            eprintln!(
11691                                "[spec] turn checkpoint skipped ({err}); \
11692                                       next turn re-primes in full"
11693                            );
11694                        }
11695                    }
11696                }
11697            }
11698        }
11699        // INIT FEED — skipped on a pending carry: last_token (the carried bonus) is NOT in the
11700        // caches and must NOT be fed solo; round 0's batched verify commits it as col 0. Its
11701        // seed/anchor hidden is the carried last_h (copied below); last_pred is dead in the
11702        // pending path (t_pred reads verify col 0 — the accept walk overwrites it).
11703        let mut last_pred = 0u32;
11704        let mut last_col_logits: Option<CudaSlice<f32>> = None;
11705        // CONSTRAINED: the init feed's logits back the (n_acc==0, base==0) masked-argmax
11706        // recompute in the grammar-truncation walk — retained host-side, round 0 only.
11707        let mut init_logits_host: Option<Vec<f32>> = None;
11708        let h_seed0: CudaSlice<f32> = if carried_pending.is_none() {
11709            let (init_logits, h) = self.spec_target_step_h(e, last_token, &mut *cache)?;
11710            last_pred = argmax(&init_logits) as u32;
11711            if constraint.is_some() {
11712                init_logits_host = Some(init_logits.clone());
11713            }
11714            // sampled mode: p-distribution after last_token, for the j==0/base==0 accept test.
11715            if sampled {
11716                last_col_logits = Some(e.htod(&init_logits)?);
11717            }
11718            h
11719        } else {
11720            // predecessor-row anchor: hidden of the last COMMITTED row (the carry contract).
11721            let lh = sess_tail
11722                .as_ref()
11723                .unwrap()
11724                .1
11725                .as_ref()
11726                .expect("pending carry requires last_h");
11727            e.clone_dtod(lh)?
11728        };
11729        let t_init = t_ent.elapsed();
11730        let mut last_col_stats: Option<(f32, f32, f32)> = None;
11731        // PERSISTENT h_seed buffer (allocated BEFORE any graph capture so no captured scratch can
11732        // alias it): every path that updates the round seed copies INTO it — no per-round allocs,
11733        // stable pointer for the graph-draft round-start copy.
11734        let mut h_seed_buf = e.clone_dtod(&h_seed0)?;
11735        // Predecessor-pairing trackers: `fill_prev` = trunk hidden AT the last COMMITTED row (the
11736        // predecessor of the next verify's col 0 — the reference's carried pending-h analogue;
11737        // also the predecessor-row hidden for the round-0 legacy-replay seed). At round 0 that
11738        // row is last_token's own (h_seed0). The chain step-0 seed under the pairing default =
11739        // hidden of the row BEFORE last_token = the prompt's last row at round 0 (h_seed_buf
11740        // overwritten below).
11741        let mut fill_prev = e.clone_dtod(&h_seed0)?;
11742        {
11743            if let Some(ph) = &prompt_h {
11744                let np = prompt.len();
11745                e.copy_view_into(
11746                    &mut h_seed_buf,
11747                    0,
11748                    &ph.slice((np - 1) * n_embd..np * n_embd),
11749                    n_embd,
11750                )?;
11751            } else if continuation {
11752                if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
11753                    if let Some(lh) = lh.as_ref() {
11754                        e.copy_into(&mut h_seed_buf, 0, lh, n_embd)?;
11755                    }
11756                }
11757            }
11758        }
11759        // Persistent device prediction slots for the accept walk (max k+1 verify columns).
11760        let mut preds_d = e.alloc_u32_zeroed(k + 2)?;
11761
11762        let debug_spec = std::env::var("MEMRA_DEBUG_SPEC").is_ok();
11763        let fork_mode = OptiForkGateMode::configured();
11764        // MEMRA_SPEC_STATS=1: per-slot accept histogram + draft-length histogram, printed once at
11765        // the end. Metric normalization vs the reference engine: BOTH engines count
11766        // accepted/drafted where the chain stopped at p-min and the sub-threshold token is
11767        // discarded uncounted — per-slot decay + chain-length mix are the extra dimensions.
11768        let spec_stats = std::env::var("MEMRA_SPEC_STATS").is_ok();
11769        let mut st_drafted = vec![0usize; k];
11770        let mut st_accepted = vec![0usize; k];
11771        let mut st_len_hist = vec![0usize; k + 1];
11772        let mut st_full = 0usize;
11773        // P-MIN CONFIDENCE GATE (MEMRA_SPEC_PMIN, the serve script's --spec-draft-p-min mechanism):
11774        // stop the draft chain early when the head's softmax confidence in its own pick drops
11775        // below p_min. Hoisted above the loop: the graph capture bakes the prob kernels iff on.
11776        static PMIN: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
11777        let p_min = *PMIN.get_or_init(|| {
11778            std::env::var("MEMRA_SPEC_PMIN")
11779                .ok()
11780                .and_then(|v| v.parse().ok())
11781                .unwrap_or(0.0)
11782        });
11783        // ZERO-DRAFT ROUNDS (MEMRA_SPEC_PMIN0=1, vendored from llama.cpp's draft gating): let the
11784        // p-min gate apply at j==0 too, so a low-confidence round drafts NOTHING and the verify
11785        // batch is just the pending bonus (m=1 = a plain decode step). llama's 35B win rides
11786        // exactly this — draft acceptance 76% at mean len 2.5 because unpredictable stretches
11787        // never pay draft+verify overhead. Only legal when a pending bonus exists (an empty
11788        // verify batch is not); the j==0 exemption stays for pending-less rounds.
11789        let pmin0 = std::env::var("MEMRA_SPEC_PMIN0")
11790            .map(|v| v == "1")
11791            .unwrap_or(false);
11792
11793        // --- GRAPH DRAFT setup: persistent I/O buffers + ONE capture (2 warmups inside). The
11794        // warmups mutate scratch len_d / pos / tok / seed — all reset at every round start, so the
11795        // only restore needed is the scratch counter. Capture failure (e.g. a non-capturable
11796        // cuBLAS path in an exotic head) falls back to the eager draft chain.
11797        // PER-SESSION PERSISTENCE (2026-08-01): session calls reuse the DraftGraphCtx parked on
11798        // the SpecSession — the capture (2 warmup head forwards + instantiate) ran ONCE at the
11799        // session's first burst, not per burst (measured ~16ms/burst fixed cost on H100 q27,
11800        // research/spec-serving-20260801). Reuse is pointer-exact: the graph bakes the session's
11801        // own scratch KV (never realloc'd), the model's resident embedding, the OnceLock p_min,
11802        // and the g_* buffers carried in the ctx — replay dispatch is identical to a fresh
11803        // capture, so draft tokens are bit-identical (drafts never decide exactness anyway; the
11804        // verify arbitrates). Single-shot calls (sess=None) build a fresh ctx and drop it.
11805        let mut dctx: DraftGraphCtx = match sess_draft_slot.as_mut().and_then(|s| s.take()) {
11806            Some(c) => c,
11807            None => DraftGraphCtx::new(e, n_embd, if sampled { d_vocab } else { 1 })?,
11808        };
11809        // A session that ran greedy bursts first sized g_q/g_perturb at 1; a sampled resume
11810        // needs d_vocab. Realloc is legal exactly while graph_s is None (nothing baked them).
11811        if sampled && dctx.g_q.len() < d_vocab {
11812            dctx.g_q = e.zeros(d_vocab)?;
11813            dctx.g_perturb = e.zeros(d_vocab)?;
11814        }
11815        // DRAFT-SIDE GRAMMAR MASK (lane/draft-mask, 2026-08-04): the drafter samples the
11816        // grammar's legal set, so proposals are legal BY CONSTRUCTION and the verify-side
11817        // truncation (the correctness backstop) stops cutting every tight-schema round.
11818        // The mask is one node inside the captured draft chain — presence is a CAPTURE-TIME
11819        // shape, so a parked graph of the other shape is dropped and recaptured.
11820        let dmask_on = constraint
11821            .as_deref()
11822            .is_some_and(|c| c.draft_mask_enabled());
11823        let dmask_words = if dmask_on { d_vocab.div_ceil(32) } else { 0 };
11824        if dmask_on && dctx.g_dmask.len() < dmask_words {
11825            dctx.g_dmask = e.alloc_u32_zeroed(dmask_words)?;
11826            dctx.graph = None; // the old capture baked the old (or no) mask pointer
11827            dctx.chain = None; // chain last-row graphs bake the same pointer
11828            dctx.failed.clear_greedy();
11829            dctx.keeper.clear();
11830        }
11831        if (dctx.graph.is_some() || dctx.chain.is_some()) && dctx.graph_masked != dmask_on {
11832            dctx.graph = None;
11833            dctx.chain = None;
11834            dctx.failed.clear_greedy();
11835            dctx.keeper.clear();
11836        }
11837        // MULTI-HEAD CHAIN mode (mtp_extra non-empty — step37's 3-head shipping shape): the
11838        // step-modulo prefix-replay chain captures PER-HEAD single-row graphs
11839        // (`DraftChainGraphs`) instead of the one self-feeding graph below; the single-head
11840        // capture arms are untouched and unreachable in this mode (the launch arms branch the
11841        // same way). This removes the historical `mtp_extra.is_empty()` capture exclusion —
11842        // and with it the silent no-attempt hole: a chain capture that FAILS now trips the
11843        // same LOUD draft-graph WARN as a single-head failure.
11844        let chain_mode = !self.mtp_extra.is_empty();
11845        if graph_draft
11846            && !sampled
11847            && chain_mode
11848            && dctx.chain.is_none()
11849            && !dctx.failed.greedy_failed()
11850        {
11851            if mtp_chain_graph_on() {
11852                // dcw door: same warmup headroom pre-arm as the single-head capture below —
11853                // every plane, because each head's capture warmups append on its OWN plane.
11854                if step35_draft_dcw_on() {
11855                    scratch.ensure_dcw_headroom(e, k + 2)?;
11856                }
11857                let heads_n = self.mtp_head_count();
11858                let DraftGraphCtx {
11859                    g_tok,
11860                    g_pos,
11861                    g_seed,
11862                    g_p,
11863                    g_dmask,
11864                    ..
11865                } = &mut dctx;
11866                if dmask_on {
11867                    e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
11868                }
11869                let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
11870                let with_prob = p_min > 0.0;
11871                // CAPTURE-RETAIN (#68 fix): one keeper for the whole chain — every graph's
11872                // warmup transients stay pinned as long as any of them replays.
11873                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
11874                    let mut interior = Vec::with_capacity(heads_n);
11875                    let mut last = Vec::with_capacity(heads_n);
11876                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
11877                    for hi in 0..heads_n {
11878                        let head = self.mtp_head_at(hi);
11879                        // interior row: KV append + carrier only (`with_head=false` — the
11880                        // eager chain discards interior logits too, so this is the same
11881                        // consumed-byte program minus the dead full-vocab head matmul).
11882                        let (g, keep) = e.capture_graph_retained(|e| {
11883                            self.mtp_head_forward_cap(
11884                                e,
11885                                head,
11886                                g_tok,
11887                                g_pos,
11888                                g_seed,
11889                                g_p,
11890                                &mut *scratch,
11891                                hi,
11892                                false,
11893                                false,
11894                                embd_gpu.expect("graph draft requires resident embedding"),
11895                                embd_qt,
11896                                embd_rb,
11897                                d_vocab,
11898                                None,
11899                                None,
11900                                None,
11901                            )
11902                        })?;
11903                        // the warmups appended rows on plane hi; rewind before the next
11904                        // capture so successive warmups never outrun the pre-armed headroom.
11905                        scratch.set_plane_len(e, hi, base)?;
11906                        interior.push(g);
11907                        keeper.extend(keep);
11908                        // last row: head matmul + greedy argmax tail (+ p when the policy
11909                        // reads it, + the grammar-mask node when constrained).
11910                        let (g2, keep2) = e.capture_graph_retained(|e| {
11911                            self.mtp_head_forward_cap(
11912                                e,
11913                                head,
11914                                g_tok,
11915                                g_pos,
11916                                g_seed,
11917                                g_p,
11918                                &mut *scratch,
11919                                hi,
11920                                with_prob,
11921                                true,
11922                                embd_gpu.expect("graph draft requires resident embedding"),
11923                                embd_qt,
11924                                embd_rb,
11925                                d_vocab,
11926                                None,
11927                                None,
11928                                if dmask_on {
11929                                    Some((g_dmask_ro, dmask_words))
11930                                } else {
11931                                    None
11932                                },
11933                            )
11934                        })?;
11935                        scratch.set_plane_len(e, hi, base)?;
11936                        last.push(g2);
11937                        keeper.extend(keep2);
11938                    }
11939                    Ok(DraftChainGraphs {
11940                        interior,
11941                        last,
11942                        keeper,
11943                    })
11944                })();
11945                match cap_res {
11946                    Ok(cg) => {
11947                        scratch.set_len(e, base)?;
11948                        // POSITIVE engagement receipt (the 3a lesson: a WARN-free boot is
11949                        // NOT evidence of capture — the captured state must name itself).
11950                        eprintln!(
11951                            "[mtp-chain-graph] captured mode=greedy heads={heads_n} \
11952                             interior={heads_n} last={heads_n} masked={}",
11953                            dmask_on as u8
11954                        );
11955                        dctx.chain = Some(cg);
11956                        dctx.graph_masked = dmask_on;
11957                    }
11958                    Err(err) => {
11959                        scratch.set_len(e, base)?;
11960                        // LOUD flip (audit Q2): a dropped draft graph is a coverage loss,
11961                        // never silent — now including the multi-head shipping shape.
11962                        if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
11963                            eprintln!("{line}");
11964                        }
11965                    }
11966                }
11967            } else {
11968                // Disarmed by MEMRA_MTP_CHAIN_GRAPH=0: say so once per process — the OFF arm
11969                // must be attributable in a boot log, never inferable from silence.
11970                static NOTE: std::sync::Once = std::sync::Once::new();
11971                NOTE.call_once(|| {
11972                    eprintln!(
11973                        "[spec] multi-head draft-chain capture disarmed \
11974                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
11975                    );
11976                });
11977            }
11978        }
11979        if graph_draft
11980            && !sampled
11981            && !chain_mode
11982            && dctx.graph.is_none()
11983            && !dctx.failed.greedy_failed()
11984        {
11985            // dcw door: the capture warmups append device-counter rows the capture body cannot
11986            // rebase for; pre-arm ring headroom host-side (no-op on flat planes / room-enough
11987            // rings, and the door-off path is untouched).
11988            if step35_draft_dcw_on() {
11989                scratch.ensure_dcw_headroom(e, k + 2)?;
11990            }
11991            let DraftGraphCtx {
11992                g_tok,
11993                g_pos,
11994                g_seed,
11995                g_p,
11996                g_dmask,
11997                ..
11998            } = &mut dctx;
11999            // capture-time contents: ALL-ONES (ban nothing). A replay only ever runs after the
12000            // host uploads the position's real words, so the warmups stay grammar-free.
12001            if dmask_on {
12002                e.htod_u32_into(g_dmask, &vec![u32::MAX; dmask_words])?;
12003            }
12004            let g_dmask_ro: &CudaSlice<u32> = &*g_dmask;
12005            // CAPTURE-RETAIN (#68 fix): the warmup transients' pool addresses are baked into the
12006            // captured graph; the keeper pins them for the graph's lifetime. capture_graph (non-
12007            // retained) freed them at exit — safe for one-shot generate_spec (nothing else touches
12008            // the pool between replays) but WRONG for sessions: burst-boundary prime/fill/commit
12009            // passes (and, in serve, other sessions) recycle those addresses and the replay then
12010            // clobbers live buffers — the ST serve-spec corruption (research/serve-st-20260803).
12011            let cap_res = e.capture_graph_retained(|e| {
12012                self.mtp_head_forward_cap(
12013                    e,
12014                    mtp,
12015                    g_tok,
12016                    g_pos,
12017                    g_seed,
12018                    g_p,
12019                    &mut *scratch,
12020                    0,
12021                    p_min > 0.0 || fork_mode == OptiForkGateMode::Controller,
12022                    true,
12023                    embd_gpu.expect("graph draft requires resident embedding"),
12024                    embd_qt,
12025                    embd_rb,
12026                    d_vocab,
12027                    None,
12028                    None,
12029                    if dmask_on {
12030                        Some((g_dmask_ro, dmask_words))
12031                    } else {
12032                        None
12033                    },
12034                )
12035            });
12036            match cap_res {
12037                Ok((g, keep)) => {
12038                    scratch.set_len(e, base)?;
12039                    dctx.graph = Some(g);
12040                    dctx.graph_masked = dmask_on;
12041                    dctx.keeper = keep;
12042                }
12043                Err(err) => {
12044                    scratch.set_len(e, base)?;
12045                    // LOUD flip (audit Q2): a dropped draft graph is a coverage loss, never
12046                    // silent. Once per flip — mark returns None on an already-failed ctx.
12047                    if let Some(line) = dctx.failed.mark_greedy(&err.to_string()) {
12048                        eprintln!("{line}");
12049                    }
12050                }
12051            }
12052        }
12053        // --- SAMPLED GRAPH DRAFT setup (step 3 of the sampled-spec arc): a SECOND capture, own
12054        // graph object, built only when sampled && graph-eligible — the greedy capture above is
12055        // untouched (and skipped when sampled: its graph would never be launched). Same head
12056        // forward, but the in-graph argmax reads GUMBEL-PERTURBED logits; the Philox event
12057        // counter lives in the persistent device g_ctr (bumped in-graph, host-seeded from sctr
12058        // once per round); the raw head logits land in the persistent g_q for the host's
12059        // per-replay async D2D into the round's q slot (q_slots, K x d_vocab, allocated once).
12060        // seed/temp are capture-time constants — baked into graph_s, so a pool-resumed request
12061        // with a different (seed, temp, k) drops the parked sampled graph and recaptures.
12062        // COST OF THE FRESH-SEED SERVE DEFAULT (dogfood F4, 2026-08-04): omitting `seed` on a
12063        // serve request now draws fresh per-request entropy (it used to default to a pinned 0),
12064        // so a seed-omitting request that RESUMES a parked spec session finds an s_key baked
12065        // with the PREVIOUS request's seed and pays one recapture. Bounded, and it does not
12066        // reopen the ~16ms/burst regression the persistent ctx exists to fix: a session's seed
12067        // is fixed for its whole lifetime (worker.rs reads s.sampler.seed() per burst), so
12068        // this compare misses at most ONCE per resumed request — the first burst recaptures
12069        // and every later burst in that request replays. A client that wants the parked graph
12070        // AND reproducibility supplies an explicit `seed`, honored exactly, which keeps s_key
12071        // stable across its whole conversation.
12072        // COMPOSITION RULE (fspec x gsd merge): the in-graph chain samples from the RAW
12073        // softmax — it can hold neither per-row filter stats nor the varying penalty history.
12074        // The sampled graph therefore engages only in the PURE-TEMP regime; filters/penalties
12075        // force the eager draft (which computes stats/penalties per row).
12076        // KEY THE WHOLE REGIME, not just the baked constants (lane/graph-s-key-exactness-
12077        // 20260819). `s_key` used to be `(seed, temp, k)`; the filters and penalties were left
12078        // out, so a filtered request resuming a session that parked a PURE-TEMP graph kept it —
12079        // and the launch site never re-asked `pure_temp`. See [`SampledGraphKey`] for what that
12080        // costs (an unconditional accept of out-of-head draft tokens, i.e. an exactness bug on
12081        // the request shape the vendor-default flip makes the majority).
12082        let s_key = SampledGraphKey::new(sp_seed, sp_temp, k, sp.top_k, sp.top_p, sp.min_p, pen_on);
12083        let pure_temp = s_key.pure_temp();
12084        // The regime the sampled graph may be captured/launched in: pure-temp always;
12085        // truncation-filtered when the filtered-capture door is on (the filter runs
12086        // IN-GRAPH — lane/step37-draft-graph-serving-20260830); penalties never.
12087        let s_capturable = s_key.graph_capturable();
12088        if sampled && dctx.s_key.is_some_and(|old| old != s_key) {
12089            dctx.graph_s = None;
12090            dctx.chain_s = None;
12091            dctx.failed.clear_sampled();
12092            dctx.s_key = None;
12093            dctx.q_slots.clear();
12094            dctx.keeper_s.clear();
12095        }
12096        // FILTERED capture nodes need q slots sized d_vocab AND the stat slots; the pure-temp
12097        // body leaves g_th/g_z/g_mx untouched (they exist from ctx creation either way).
12098        if graph_draft
12099            && sampled
12100            && s_capturable
12101            && chain_mode
12102            && dctx.chain_s.is_none()
12103            && !dctx.failed.sampled_failed()
12104        {
12105            if mtp_chain_graph_on() {
12106                if step35_draft_dcw_on() {
12107                    scratch.ensure_dcw_headroom(e, k + 2)?;
12108                }
12109                let heads_n = self.mtp_head_count();
12110                let filtered = s_key.filtered();
12111                let DraftGraphCtx {
12112                    g_tok,
12113                    g_pos,
12114                    g_seed,
12115                    g_p,
12116                    g_ctr,
12117                    g_perturb,
12118                    g_q,
12119                    g_rows0,
12120                    g_th,
12121                    g_z,
12122                    g_mx,
12123                    ..
12124                } = &mut dctx;
12125                let with_prob = p_min > 0.0;
12126                let cap_res = (|| -> Result<DraftChainGraphs, Box<dyn std::error::Error>> {
12127                    let mut interior = Vec::with_capacity(heads_n);
12128                    let mut last = Vec::with_capacity(heads_n);
12129                    let mut keeper: Vec<Box<dyn std::any::Any + Send>> = Vec::new();
12130                    for hi in 0..heads_n {
12131                        let head = self.mtp_head_at(hi);
12132                        // interior row: no head, no draw — shared shape with the greedy
12133                        // chain's interior, captured per mode for keeper-lifetime hygiene.
12134                        let (g, keep) = e.capture_graph_retained(|e| {
12135                            self.mtp_head_forward_cap(
12136                                e,
12137                                head,
12138                                g_tok,
12139                                g_pos,
12140                                g_seed,
12141                                g_p,
12142                                &mut *scratch,
12143                                hi,
12144                                false,
12145                                false,
12146                                embd_gpu.expect("graph draft requires resident embedding"),
12147                                embd_qt,
12148                                embd_rb,
12149                                d_vocab,
12150                                None,
12151                                None,
12152                                None,
12153                            )
12154                        })?;
12155                        scratch.set_plane_len(e, hi, base)?;
12156                        interior.push(g);
12157                        keeper.extend(keep);
12158                        // last row: head matmul + the in-graph categorical draw (filtered
12159                        // nodes when the request carries filters).
12160                        let (g2, keep2) = e.capture_graph_retained(|e| {
12161                            self.mtp_head_forward_cap(
12162                                e,
12163                                head,
12164                                g_tok,
12165                                g_pos,
12166                                g_seed,
12167                                g_p,
12168                                &mut *scratch,
12169                                hi,
12170                                with_prob,
12171                                true,
12172                                embd_gpu.expect("graph draft requires resident embedding"),
12173                                embd_qt,
12174                                embd_rb,
12175                                d_vocab,
12176                                Some(SampledCapArgs {
12177                                    ctr: &mut *g_ctr,
12178                                    perturb: &mut *g_perturb,
12179                                    q_out: &mut *g_q,
12180                                    seed: sp_seed,
12181                                    temp: sp_temp,
12182                                    filt: if filtered {
12183                                        Some(SampledCapFilter {
12184                                            rows0: &*g_rows0,
12185                                            th: &mut *g_th,
12186                                            z: &mut *g_z,
12187                                            mx: &mut *g_mx,
12188                                            top_k: sp.top_k,
12189                                            top_p: sp.top_p,
12190                                            min_p: sp.min_p,
12191                                        })
12192                                    } else {
12193                                        None
12194                                    },
12195                                }),
12196                                None,
12197                                None, // constrained spec is greedy-only
12198                            )
12199                        })?;
12200                        scratch.set_plane_len(e, hi, base)?;
12201                        last.push(g2);
12202                        keeper.extend(keep2);
12203                    }
12204                    Ok(DraftChainGraphs {
12205                        interior,
12206                        last,
12207                        keeper,
12208                    })
12209                })();
12210                match cap_res {
12211                    Ok(cg) => {
12212                        scratch.set_len(e, base)?;
12213                        for _ in 0..k {
12214                            dctx.q_slots.push(e.zeros(d_vocab)?);
12215                        }
12216                        eprintln!(
12217                            "[mtp-chain-graph] captured mode=sampled heads={heads_n} \
12218                             interior={heads_n} last={heads_n} filtered={} key={s_key:?}",
12219                            s_key.filtered() as u8
12220                        );
12221                        dctx.chain_s = Some(cg);
12222                        dctx.s_key = Some(s_key);
12223                    }
12224                    Err(err) => {
12225                        scratch.set_len(e, base)?;
12226                        if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
12227                            eprintln!("{line}");
12228                        }
12229                    }
12230                }
12231            } else {
12232                static NOTE_S: std::sync::Once = std::sync::Once::new();
12233                NOTE_S.call_once(|| {
12234                    eprintln!(
12235                        "[spec] multi-head draft-chain capture disarmed \
12236                         (MEMRA_MTP_CHAIN_GRAPH=0); eager chain serves this shape"
12237                    );
12238                });
12239            }
12240        }
12241        if graph_draft
12242            && sampled
12243            && s_capturable
12244            && !chain_mode
12245            && dctx.graph_s.is_none()
12246            && !dctx.failed.sampled_failed()
12247        {
12248            // dcw door: same warmup headroom pre-arm as the greedy capture above.
12249            if step35_draft_dcw_on() {
12250                scratch.ensure_dcw_headroom(e, k + 2)?;
12251            }
12252            let filtered = s_key.filtered();
12253            let DraftGraphCtx {
12254                g_tok,
12255                g_pos,
12256                g_seed,
12257                g_p,
12258                g_ctr,
12259                g_perturb,
12260                g_q,
12261                g_rows0,
12262                g_th,
12263                g_z,
12264                g_mx,
12265                ..
12266            } = &mut dctx;
12267            // CAPTURE-RETAIN (#68 fix): same keeper contract as the greedy capture above.
12268            let cap_res = e.capture_graph_retained(|e| {
12269                self.mtp_head_forward_cap(
12270                    e,
12271                    mtp,
12272                    g_tok,
12273                    g_pos,
12274                    g_seed,
12275                    g_p,
12276                    &mut *scratch,
12277                    0,
12278                    p_min > 0.0,
12279                    true,
12280                    embd_gpu.expect("graph draft requires resident embedding"),
12281                    embd_qt,
12282                    embd_rb,
12283                    d_vocab,
12284                    Some(SampledCapArgs {
12285                        ctr: &mut *g_ctr,
12286                        perturb: &mut *g_perturb,
12287                        q_out: &mut *g_q,
12288                        seed: sp_seed,
12289                        temp: sp_temp,
12290                        filt: if filtered {
12291                            Some(SampledCapFilter {
12292                                rows0: &*g_rows0,
12293                                th: &mut *g_th,
12294                                z: &mut *g_z,
12295                                mx: &mut *g_mx,
12296                                top_k: sp.top_k,
12297                                top_p: sp.top_p,
12298                                min_p: sp.min_p,
12299                            })
12300                        } else {
12301                            None
12302                        },
12303                    }),
12304                    None,
12305                    None, // constrained spec is greedy-only — sampled never carries a hook
12306                )
12307            });
12308            match cap_res {
12309                Ok((g, keep)) => {
12310                    scratch.set_len(e, base)?;
12311                    for _ in 0..k {
12312                        dctx.q_slots.push(e.zeros(d_vocab)?);
12313                    }
12314                    dctx.graph_s = Some(g);
12315                    dctx.s_key = Some(s_key);
12316                    dctx.keeper_s = keep;
12317                }
12318                Err(err) => {
12319                    scratch.set_len(e, base)?;
12320                    // LOUD flip (audit Q2): same contract as the greedy capture above.
12321                    if let Some(line) = dctx.failed.mark_sampled(&err.to_string()) {
12322                        eprintln!("{line}");
12323                    }
12324                }
12325            }
12326        }
12327        // ---- EXACTNESS GUARD, the enforceable half (lane/graph-s-key-exactness-20260819,
12328        // widened by lane/step37-draft-graph-serving-20260830) ----
12329        // With the filters and penalties in `s_key`, a graph that SURVIVED the drop above was
12330        // captured under THIS request's exact regime, and capture requires `graph_capturable`
12331        // (pure-temp, or filtered with the in-graph filter nodes; never penalties) — so a
12332        // parked graph implies both. That implication is the whole exactness argument for the
12333        // graph arm, so it is asserted here rather than assumed: a future change that widens
12334        // the capture condition, narrows the key, or copies a `DraftGraphCtx` across regimes
12335        // fails LOUDLY at this line instead of silently drafting from a distribution the
12336        // verify never reconstructs. Release builds refuse the graph (drop it, draft eager)
12337        // rather than launching it; the launch site re-tests the regime independently.
12338        if sampled
12339            && (dctx.graph_s.is_some() || dctx.chain_s.is_some())
12340            && (!s_capturable || dctx.s_key != Some(s_key))
12341        {
12342            debug_assert!(
12343                false,
12344                "sampled draft graph parked under {:?} survived into a request outside its \
12345                 capture regime (top_k={} top_p={} min_p={} pen_on={} capturable={}): the \
12346                 in-graph draw and the verify's accept test would see different distributions",
12347                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12348            );
12349            eprintln!(
12350                "[spec] BUG: dropping a parked sampled draft graph that outlived its capture \
12351                 regime (s_key={:?}, request top_k={} top_p={} min_p={} pen_on={} \
12352                 capturable={}); drafting EAGER — the key must carry every field that shapes q",
12353                dctx.s_key, sp.top_k, sp.top_p, sp.min_p, pen_on, s_capturable,
12354            );
12355            dctx.graph_s = None;
12356            dctx.chain_s = None;
12357            dctx.s_key = None;
12358            dctx.q_slots.clear();
12359            dctx.keeper_s.clear();
12360        }
12361        // SKEY PROBE (MEMRA_SKEY_PROBE=1): the burst-entry facts the reachability question turns
12362        // on — is this request sampled, is it in a regime the sampled graph is legal in, and is
12363        // a graph PARKED from an earlier request of the same session? The launch arms below
12364        // print which chain actually ran, so the probe never restates the condition.
12365        if skey_probe() {
12366            eprintln!(
12367                "[skey] burst sampled={} pure_temp={} capturable={} temp={} top_k={} top_p={} \
12368                 min_p={} pen_on={} k={} graph_draft={} graph_s_parked={} chain_s_parked={} \
12369                 s_key_parked={:?}",
12370                sampled as u8,
12371                pure_temp as u8,
12372                s_capturable as u8,
12373                sp_temp,
12374                sp.top_k,
12375                sp.top_p,
12376                sp.min_p,
12377                pen_on as u8,
12378                k,
12379                graph_draft as u8,
12380                dctx.graph_s.is_some() as u8,
12381                dctx.chain_s.is_some() as u8,
12382                dctx.s_key,
12383            );
12384        }
12385        let t_cap = t_ent.elapsed();
12386        // PERSISTENT DRAFT KV: fill the MTP block's K/V for every prompt position from the exact
12387        // trunk hiddens collected during prime — ONE batched K/V-only pass (overwrites any
12388        // capture-warmup garbage; capture left len at 0). last_token (the init feed) needs no
12389        // fill: the first chain step processes it and appends its entry at slot prompt.len().
12390        if let Some(ph) = &prompt_h {
12391            // SESSION: rows [0..base) are the previous turns' exact fills (refresh overwrote them
12392            // with true verify hiddens) — truncate any draft overhang, fill ONLY the suffix at
12393            // global positions [base..base+tp). Fresh call: base==0, identical to before.
12394            scratch.set_len(e, base)?;
12395            // CHUNKED FILL (long-ctx OOM fix, 2026-07-05): mtp_kv_fill's transients scale with its
12396            // T (concat = T*2*n_embd*4B — 1.5GB at 40k) and its concat loop is 2*T launches. The
12397            // fill is a pure sequential append, so chunking is exact: each chunk appends its rows
12398            // at pos0=base+start with the identical per-row math. Same knob as the trunk prime.
12399            let tp = prompt.len();
12400            let fill_chunk: usize = if crate::cache::swa_ring_on() {
12401                crate::hybrid_forward::prime_chunk_tokens(tp, self.layers.len())
12402            } else {
12403                // Preserve the flag-OFF schedule byte-for-byte, including the legacy zero value
12404                // meaning one monolithic fill.
12405                std::env::var("MEMRA_PRIME_CHUNK")
12406                    .ok()
12407                    .and_then(|v| v.parse().ok())
12408                    .unwrap_or(4096)
12409            };
12410            let fill_chunk = if fill_chunk == 0 { tp } else { fill_chunk };
12411            let mut start = 0usize;
12412            while start < tp {
12413                let end = (start + fill_chunk).min(tp);
12414                let tc = end - start;
12415                {
12416                    // PREDECESSOR pairing: row i gets h[i-1]; global row 0 a zeros row (the
12417                    // reference engine's initial pending-h is zeroed too); a session turn's row 0
12418                    // gets the PREVIOUS turn's last committed hidden (sess.last_h). Per chunk:
12419                    // rows start..end read h[start-1..end-1] — one dtod into a chunk buffer.
12420                    let mut phs = e.zeros(tc * n_embd)?;
12421                    let (src_lo, dst_off) = if start == 0 {
12422                        (0, n_embd)
12423                    } else {
12424                        ((start - 1) * n_embd, 0)
12425                    };
12426                    let n_copy = if start == 0 {
12427                        (tc - 1) * n_embd
12428                    } else {
12429                        tc * n_embd
12430                    };
12431                    if start == 0 {
12432                        if let Some((_, lh, _, _, _)) = sess_tail.as_ref() {
12433                            if let Some(lh) = lh.as_ref() {
12434                                e.copy_into(&mut phs, 0, lh, n_embd)?;
12435                            }
12436                        }
12437                    }
12438                    if n_copy > 0 {
12439                        e.copy_view_into(
12440                            &mut phs,
12441                            dst_off,
12442                            &ph.slice(src_lo..src_lo + n_copy),
12443                            n_copy,
12444                        )?;
12445                    }
12446                    self.mtp_kv_fill_all(
12447                        e,
12448                        &prompt[start..end],
12449                        &phs,
12450                        base + start,
12451                        &mut *scratch,
12452                        embd_dev,
12453                    )?;
12454                }
12455                start = end;
12456            }
12457        }
12458        // MEMRA_PROFILE_SPEC=2: profiler capture starts HERE — after the prime, so an
12459        // `nsys -c cudaProfilerApi` capture contains ONLY the round loop (draft/verify/commit).
12460        // (=1 brackets the whole call in run_spec.rs, prime included.)
12461        if std::env::var("MEMRA_PROFILE_SPEC").as_deref() == Ok("2") {
12462            unsafe extern "C" {
12463                fn cudaProfilerStart() -> i32;
12464            }
12465            unsafe {
12466                cudaProfilerStart();
12467            }
12468        }
12469        // ROUND-STREAM stage (c) 4 (MEMRA_SPEC_STREAM=1, experimental): pre-issued M-round
12470        // bursts with ZERO per-round host readbacks — the accept/seed/rollback/ring kernels
12471        // consume each other's device outputs; the host drains the ring every M rounds. v1
12472        // constraints: greedy, !spec_replay, single-shot, batched-linear layers, no refresh
12473        // fills (acceptance effect A/B-arbitrated), enters from round 1 (pending guaranteed).
12474        // NOTE: not gated on the caller's graph_draft (its trunk_dense conjunct turns the 35B
12475        // MoE off) — the stream capture encloses ONLY the dense MTP head; the head-dense /
12476        // full-prec / k gates are re-derived here and a failed capture degrades to stream-off.
12477        let stream_on = crate::spec::spec_stream()
12478            && !sampled
12479            && !spec_replay
12480            && self.mtp_extra.is_empty()
12481            && constraint.is_none()
12482            && !session_mode
12483            && embd_gpu.is_some()
12484            && !crate::model::full_prec_enabled()
12485            && k + 2 < 96;
12486        let mut stream_graph: Option<cudarc::driver::CudaGraph> = None;
12487        let mut g_tokp2k = e.alloc_u32_zeroed(2 * k.max(1))?;
12488        if stream_on {
12489            let cap = e.capture_graph(|e| {
12490                for j in 0..k.max(1) {
12491                    self.mtp_head_forward_cap(
12492                        e,
12493                        mtp,
12494                        &mut dctx.g_tok,
12495                        &mut dctx.g_pos,
12496                        &mut dctx.g_seed,
12497                        &mut dctx.g_p,
12498                        &mut *scratch,
12499                        0,
12500                        true,
12501                        true,
12502                        embd_gpu.expect("round stream requires resident embedding"),
12503                        embd_qt,
12504                        embd_rb,
12505                        d_vocab,
12506                        None,
12507                        Some((&mut g_tokp2k, j, d2t_dev.as_ref())),
12508                        None, // round-stream requires constraint.is_none() (see stream_on)
12509                    )?;
12510                }
12511                Ok(())
12512            });
12513            match cap {
12514                Ok(g) => {
12515                    scratch.set_len(e, 0)?;
12516                    stream_graph = Some(g);
12517                }
12518                Err(err) => {
12519                    scratch.set_len(e, 0)?;
12520                    if debug_spec {
12521                        eprintln!("[spec] stream-graph capture failed ({err}); stream off");
12522                    }
12523                }
12524            }
12525        }
12526        let stream_active = stream_on && stream_graph.is_some();
12527        if debug_spec {
12528            eprintln!(
12529                "[spec] stream_on={stream_on} env={} samp={sampled} dg={} captured={} active={stream_active} session={session_mode} replay={spec_replay}",
12530                crate::spec::spec_stream(),
12531                dctx.graph.is_some(),
12532                stream_graph.is_some()
12533            );
12534        }
12535        let t_v_s = k + 1;
12536        // ROUND-STREAM buffers + ptr tables now live in the model-generic round_stream
12537        // module (extracted 2026-07-12; the gemma burst reuses them).
12538        let sb = crate::round_stream::StreamBufs::new(e, k, crate::spec::spec_stream_m())?;
12539        let crate::round_stream::StreamBufs {
12540            mut vtok_d,
12541            mut brk_d,
12542            mut pend_d,
12543            last_pred_d,
12544            mut pos_ctr,
12545            mut pos_start_d,
12546            mut ring_d,
12547            acc_d: mut stream_acc,
12548            m_rounds,
12549            k: _,
12550        } = sb;
12551        let stream_ptrs: Option<CudaSlice<u64>> = if stream_active {
12552            Some(crate::round_stream::kv_len_ptr_table(
12553                e,
12554                cache,
12555                Some(&pos_ctr),
12556            )?)
12557        } else {
12558            None
12559        };
12560
12561        let t_fill = t_ent.elapsed();
12562        let mut round = 0usize;
12563        // ADAPTIVE DRAFT LENGTH (MEMRA_SPEC_ADAPT=1, opt-in — the gemma_spec accepted-run law,
12564        // ported 2026-08-01): next round's draft depth = last round's accepted run + 1, clamped
12565        // to [floor(pos), k_cap] — a miss shrinks the next draft to the miss point + 1,
12566        // full-accept streaks re-deepen one step per round. NOT the 2026-07-07 acceptance-EMA
12567        // (that arm measured an HONEST LOSS to static per-class optima — 115.0/85.8/73.4 vs
12568        // 121.6/92.7/75.6, EMA lag — and was removed 2026-07-08; rig5090.jsonl has the record).
12569        // The gemma law has no lag class: it reacts within one round, and was worth +7-20% on
12570        // the gemma cells at unchanged exactness (2026-07-10 flip; floor sweep 2026-07-25;
12571        // position key 2026-07-26). Signal = n_acc from the round's EXISTING accept readback —
12572        // zero new syncs; the draft graph is a SINGLE-STEP capture replayed per drafted token,
12573        // so a per-round depth needs no re-capture (unlike gemma's whole-chain graphs). qwen's
12574        // in-round p-min cut already shortens chains mid-round, so gemma's one-round-late p-min
12575        // fold into kc is unnecessary here — the accepted-run law sees the cut via n_acc.
12576        // Exactness is the verify's job at ANY depth (same contract as p-min variable rounds).
12577        // DEFAULT OFF on the qwen path until its cells gate a flip (gemma's is default-on).
12578        // MEASURED 2026-08-01 (H100 GPU-3, interleaved x3, NGEN=256, same-invocation plain
12579        // denominators; research/qwen-adaptive-k-20260801/): REFUTED on the tuned qwen configs.
12580        // q27 K=3+HPOST+PMIN=0.3: short +0.8% (noise; law ~idles, len_hist identical), board
12581        // -1.9%, agentic -0.5%; board PMIN=0 -2.1% (not p-min shadowing — the law itself);
12582        // floor=1 -2.8% (gemma's floor-collapse, reproduced). q35 K=2 board: -6.4% (52/136
12583        // rounds shrink to depth 1; no depth to reclaim at K=2). The gemma direction DOES
12584        // appear at untuned depth-K — q27 K=6 floor=4 +1.5% over fixed K=6 — but stays -3.7%
12585        // below fixed K=3: same verdict class as the retired EMA arm (honest loss to static
12586        // per-class optima). Acceptance-rate rises under the law while tokens/round falls —
12587        // it buys accept-% by adding rounds, and a round's fixed draft+verify cost wins.
12588        // K=1..8 self-consistency PASS both models with the law ON (exactness held).
12589        let adapt = std::env::var("MEMRA_SPEC_ADAPT").as_deref() == Ok("1");
12590        // floor: per-model default keyed on n_embd (gemma's tiering — models with an expensive
12591        // verify keep deep drafts after a miss); MEMRA_SPEC_ADAPT_FLOOR pins it everywhere.
12592        let adapt_floor_env: Option<usize> = std::env::var("MEMRA_SPEC_ADAPT_FLOOR")
12593            .ok()
12594            .and_then(|v| v.parse().ok());
12595        let adapt_floor_default: usize = if self.cfg.n_embd as usize >= 3500 {
12596            4
12597        } else if self.cfg.n_embd as usize >= 2500 {
12598            2
12599        } else {
12600            1
12601        };
12602        let adapt_floor: usize = adapt_floor_env.unwrap_or(adapt_floor_default);
12603        // position key: past floor_ctx a HIGH floor (>=4) relaxes to 1 — forced-deep drafts
12604        // turn net-negative at depth (gemma 31B d1736 evidence); MEMRA_SPEC_FLOOR_CTX moves
12605        // the boundary, an explicit MEMRA_SPEC_ADAPT_FLOOR pins the floor everywhere.
12606        let floor_ctx: usize = std::env::var("MEMRA_SPEC_FLOOR_CTX")
12607            .ok()
12608            .and_then(|v| v.parse().ok())
12609            .unwrap_or(1024);
12610        let floor_at = |pos: usize| -> usize {
12611            if adapt_floor_env.is_some() || pos < floor_ctx {
12612                adapt_floor
12613            } else if adapt_floor >= 4 {
12614                1
12615            } else {
12616                adapt_floor
12617            }
12618        };
12619        // cap: MEMRA_SPEC_CAPMAX (gemma semantics, default 7). Binds only under adapt — the
12620        // fixed-K default path is untouched by this whole block.
12621        let cap_max: usize = std::env::var("MEMRA_SPEC_CAPMAX")
12622            .ok()
12623            .and_then(|v| v.parse().ok())
12624            .unwrap_or(7);
12625        let k_cap = k.min(cap_max).max(1);
12626        let mut kc = k_cap;
12627        let mut opti_fork: Option<OptiForkState> = None;
12628        let mut fork_snapshot: Option<crate::cache::CacheSnapshot> = None;
12629        if fork_mode != OptiForkGateMode::Disabled {
12630            let fence = crate::pp::pp_cuts(self.layers.len());
12631            let refusal = if !session_mode {
12632                Some("not-session")
12633            } else if k != 1 || adapt {
12634                Some("requires-fixed-k1")
12635            } else if sampled || constraint.is_some() || spec_replay {
12636                Some("sampled-constrained-or-replay")
12637            } else if pipe.is_some() {
12638                Some("two-session-pipeline")
12639            } else if !spec_devacc() {
12640                Some("requires-device-accept")
12641            } else if stream_active || crate::spec::spec_stream() {
12642                Some("round-stream")
12643            } else if !self.mtp_extra.is_empty() {
12644                Some("multi-head-mtp")
12645            } else if crate::cache::swa_ring_on() || cache.has_swa_ring() {
12646                Some("swa-ring")
12647            } else if crate::pp::pp_host_bounce_active() {
12648                Some("host-bounce")
12649            } else if fork_mode == OptiForkGateMode::Controller
12650                && cache.recur.iter().any(Option::is_some)
12651            {
12652                Some("controller-requires-zero-recurrent-state")
12653            } else if fence.as_ref().is_none_or(|f| f.len() != 3) {
12654                Some("requires-pp2")
12655            } else {
12656                None
12657            };
12658            if let Some(reason) = refusal {
12659                OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12660                eprintln!("[opti-fork] refused reason={reason}");
12661            } else {
12662                let fence = fence.expect("validated PP-2 fence");
12663                let rt = crate::pp::PpNRt::get(e)?;
12664                let primary_stage0 = rt.engine(0, e).ctx().ordinal() == e.ctx().ordinal();
12665                let primary_stage1 = rt.engine(1, e).ctx().ordinal() == e.ctx().ordinal();
12666                let primary_supported =
12667                    primary_stage0 || (fork_mode == OptiForkGateMode::Controller && primary_stage1);
12668                if !rt.cross_device() || !primary_supported {
12669                    OPTI_FORK_REFUSALS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
12670                    eprintln!("[opti-fork] refused reason=requires-supported-primary-cross-device");
12671                } else {
12672                    // Both recurrent snapshots and both seed generations are allocated before
12673                    // the first fork, each through its owning PP stage. Allocation failure
12674                    // therefore happens before any optimistic state mutation can occur.
12675                    let current_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12676                    let alternate_snapshot = opti_snapshot_stage_owned(e, cache, rt, &fence)?;
12677                    let fork = OptiForkState::new(
12678                        e,
12679                        cache,
12680                        fork_mode,
12681                        alternate_snapshot,
12682                        &h_seed_buf,
12683                        &fill_prev,
12684                        rt,
12685                        fence[1],
12686                        self.layers.len(),
12687                    )?;
12688                    eprintln!(
12689                        "[opti-fork] armed mode={fork_mode:?} snapshots=2 seeds=2 split={} \
12690                         payload_dev0={} payload_dev1={} q_threshold={:.3}",
12691                        fence[1],
12692                        fork.logical_payload_bytes[0],
12693                        fork.logical_payload_bytes[1],
12694                        fork.controller.map_or(0.0, |policy| policy.threshold),
12695                    );
12696                    fork_snapshot = Some(current_snapshot);
12697                    opti_fork = Some(fork);
12698                }
12699            }
12700        }
12701        // Persistent snapshot buffers are allocated once and refreshed in place. The fork arm
12702        // uses stage-owned snapshots; refused/disabled arms retain the existing generic helper.
12703        let mut snap = match fork_snapshot {
12704            Some(snapshot) => snapshot,
12705            None => cache.snapshot(e)?,
12706        };
12707        let mut carried_opti: Option<OptiControllerTicket> = None;
12708        // ROUND-STREAM stage (b) 3a: device table of per-layer kvl.len_d pointers (stable — the
12709        // cache never reallocates len_d; see cache.rs "stable pointer" note). 0 = no KV layer.
12710        let kv_len_ptrs: Option<CudaSlice<u64>> = if spec_devacc() && !spec_replay {
12711            Some(crate::round_stream::kv_len_ptr_table(e, cache, None)?)
12712        } else {
12713            None
12714        };
12715        // BONUS FOLD (2026-07-04): after a FULL accept the bonus token is NOT committed with a
12716        // separate T=1 trunk pass (a full weight read per round). It stays PENDING and rides as
12717        // column 0 of the NEXT round's verify batch. Under predecessor pairing the next chain
12718        // seeds from the bonus's predecessor's TRUE verify hidden (free — no extra
12719        // pass of any kind). Verify still
12720        // checks every emitted token against the target -> exactness holds by construction; only
12721        // DRAFT QUALITY can shift, which the acceptance numbers arbitrate.
12722        // bonus emitted but not yet committed to cache. A carried pending (see SpecSession::
12723        // pending_tok) enters round 0 directly — the burst boundary becomes a plain round edge.
12724        let mut pending: Option<u32> = carried_pending;
12725        // MEMRA_SPEC_PHASE=1: per-round wall decomposition (draft / verify / accept+commit) —
12726        // no tracing, no extra syncs (each phase is naturally sync-bounded: draft readbacks,
12727        // the verify accept readback). Printed once at loop end via spec-stats.
12728        let anatomy_on = std::env::var("MEMRA_SPEC_PP_ANATOMY").as_deref() == Ok("1");
12729        let phase_on = anatomy_on || std::env::var("MEMRA_SPEC_PHASE").as_deref() == Ok("1");
12730        // MEMRA_SPEC_PHASE_SYNC=1 — reads the phase split correctly, and proves it. `ph_mark` is a
12731        // bare Instant, so `verify-issue` is the host QUEUEING the walk (the GPU is already running
12732        // under it) and `verify-wait` is only the residual drain at the accept readback: one
12733        // overlapped interval cut at the first blocking call, NOT "GPU time" beside "host time".
12734        // Syncing right after the walk is issued moves the whole GPU wall into `verify-issue`. If
12735        // the walk's GPU total is really issue+wait, then with this on verify-issue jumps to that
12736        // sum, verify-wait collapses to the readback alone, and the ROUND WALL DOES NOT MOVE —
12737        // which is what says the queueing time was hidden and is not a target. Diagnostic only.
12738        let phase_sync = std::env::var("MEMRA_SPEC_PHASE_SYNC").as_deref() == Ok("1");
12739        // DRAFT-MASK receipt (lane/draft-mask): speculative-clone wall + rounds, printed with
12740        // spec-stats. The clone is the one cost the design adds per round — measured, not assumed.
12741        let (mut dm_clone_ns, mut dm_rounds) = (0u128, 0usize);
12742        // grammar-truncation counters: how many rounds the verify-side cut fired and how many
12743        // already-verified tokens it threw away. THIS is the quantity draft masking targets.
12744        let (mut dm_cuts, mut dm_cut_tokens) = (0usize, 0usize);
12745        let (mut ph_draft, mut ph_verify, mut ph_rest) = (0f64, 0f64, 0f64);
12746        let mut ph_wait = 0f64;
12747        let mut ph_commit = 0f64;
12748        let mut ph_t = std::time::Instant::now();
12749        let mut ph_mark = |acc: &mut f64, on: bool| {
12750            if on {
12751                let now = std::time::Instant::now();
12752                *acc += (now - ph_t).as_secs_f64();
12753                ph_t = now;
12754            }
12755        };
12756        // MTP-ROUTE VERIFY GRAPHS (`MEMRA_SPEC_VERIFY_GRAPH`, see the flag doc): the
12757        // model-owned capture pool, locked for the whole burst exactly as the dspark serve
12758        // arm holds it — the slab stash is live verify -> commit inside a round, and the
12759        // worker drives rounds from one scheduler thread. PERSISTENT across generations on
12760        // the model (rebuilding per call re-captures the pool per prompt, which is the
12761        // measured way to lose more than the launches cost); the captured bodies are
12762        // cache-independent, every state read going through per-round refreshed pointer
12763        // tables. None = the eager walk, byte-identical.
12764        //
12765        // Never armed together with ROUND-STREAM: the tparallel verify refuses that pair
12766        // loudly, and `stream_active` owns the burst arm above, so the door stays shut
12767        // whenever the stream is live rather than relying on that refusal.
12768        // The lock is taken ONLY when the door is armed: with the flag off this whole block
12769        // is inert, so the default path cannot serialize two spec generations behind a mutex
12770        // it never reads.
12771        let vg_armed =
12772            crate::spec::spec_verify_graph_env().unwrap_or_else(|| self.vgraph_family_default());
12773        let mut vg_guard = if vg_armed && !stream_active {
12774            let mut g = self.dspark_vgraphs.lock().unwrap();
12775            if g.is_none() {
12776                // Size by the WIDEST verify this run can present, which is k+1 and NOT
12777                // k_cap+1: the sampled arm's own window is `t_v_s = k + 1`, so a pool built
12778                // from a smaller adaptive cap gets sliced past its stash rows (a `slice_mut`
12779                // panic in the sampled ON arm, measured before this line said k+1).
12780                let vt_cap = (k.max(k_cap) + 1).max(2);
12781                *g = DsparkVerifyGraphs::new(e, cache, vt_cap, n_embd)?;
12782                if g.is_some() {
12783                    // Engagement receipt (the dead-arm lesson): prove the door is LIVE rather
12784                    // than trusting that a flag set means a pool built.
12785                    eprintln!("[spec-vg] MTP verify-graph pool ENGAGED (vt_cap={vt_cap})");
12786                } else {
12787                    eprintln!(
12788                        "[spec-vg] MTP verify-graph pool declined (no linear layers, \
12789                         non-uniform state, or vt_cap < 2) — eager walk"
12790                    );
12791                }
12792            }
12793            Some(g)
12794        } else {
12795            None
12796        };
12797        // Capacity fail-safe: a round wider than the pool was built for must take the eager
12798        // walk, not slice the stash past its rows. The sizing above already covers every
12799        // round this run can present; this keeps a future caller (or a k that grows behind
12800        // the pool's back) on the byte-identical fallback instead of a panic.
12801        let vg_t_cap = vg_guard
12802            .as_ref()
12803            .and_then(|g| g.as_ref())
12804            .map(|g| g.t_capacity())
12805            .unwrap_or(0);
12806        if let Some(p) = pipe {
12807            p.setup_end();
12808        }
12809        while keep_going && out.len() < max_new {
12810            // MEMRA_SPEC_ROUND_PROF=1: wall of the WHOLE round against the pieces we already
12811            // instrument. Needed because the parts do not add up: the draft step measures 1.27 ms
12812            // ([spec-anatomy] glue 92 / attn 280 / ffn 222 / head 670 us) and the t=2 verify walk
12813            // 25.6 ms ([tcol-prof] attn 10.1 + ffn 15.3), yet a K=1 round takes 177 ms on the
12814            // step37 TP2 stack. This prints where the other ~150 ms lives.
12815            let round_prof = ROUND_PROF
12816                .get_or_init(|| std::env::var("MEMRA_SPEC_ROUND_PROF").as_deref() == Ok("1"));
12817            let round_t0 = round_prof.then(std::time::Instant::now);
12818            // ROUND-STREAM BURST: from round 1 (pending guaranteed by every non-replay arm),
12819            // issue M rounds with zero readbacks, then drain the ring + reconcile mirrors.
12820            if let (true, Some(sg), Some(ptrs)) = (
12821                stream_active && round >= 1 && pending.is_some(),
12822                &stream_graph,
12823                &stream_ptrs,
12824            ) {
12825                if debug_spec {
12826                    static ONCE: std::sync::Once = std::sync::Once::new();
12827                    ONCE.call_once(|| {
12828                        eprintln!("[memra] ROUND-STREAM burst engaged (M={m_rounds} k={k})")
12829                    });
12830                }
12831                e.set_i32_one(&mut pos_ctr, cache.pos as i32)?;
12832                e.set_u32_one(&mut pend_d, pending.unwrap())?;
12833                e.set_u32_one(&mut ring_d, 0)?; // ring count = 0 (writes element 0)
12834                for _mi in 0..m_rounds {
12835                    e.i32_copy_add(&pos_ctr, &mut pos_start_d, 0)?;
12836                    cache.snapshot_into(e, &mut snap)?; // device D2Ds, stream-ordered
12837                    e.i32_copy_add(&pos_ctr, &mut scratch.kv.len_d, 0)?; // draft-KV rollback
12838                    e.i32_copy_add(&pos_ctr, &mut dctx.g_pos, 1)?; // rope pos = pos + base
12839                    e.u32_copy(&pend_d, &mut dctx.g_tok)?;
12840                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
12841                    sg.launch()?;
12842                    e.spec_assemble_verify(
12843                        &g_tokp2k,
12844                        &pend_d,
12845                        d2t_dev.as_ref(),
12846                        &mut vtok_d,
12847                        &mut brk_d,
12848                        p_min,
12849                        k,
12850                        pmin0,
12851                    )?;
12852                    let mut ck = VerifyCkpt::new(self.layers.len());
12853                    let dummy = vec![0u32; t_v_s];
12854                    let (tl_d, vx) = self.decode_step_t_core_stream(
12855                        e,
12856                        &dummy,
12857                        0,
12858                        &mut *cache,
12859                        embd_dev,
12860                        Some(&mut ck),
12861                        Some((&vtok_d, &pos_ctr)),
12862                        None,
12863                        None,
12864                        None,
12865                    )?;
12866                    for j in 0..t_v_s {
12867                        e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
12868                    }
12869                    e.spec_accept_greedy_dc(
12870                        &preds_d,
12871                        &vtok_d,
12872                        &last_pred_d,
12873                        &brk_d,
12874                        &mut stream_acc,
12875                    )?;
12876                    e.spec_seed_gather(&vx, &fill_prev, &stream_acc, &mut h_seed_buf, 1, n_embd)?;
12877                    e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
12878                    self.commit_verified_prefix_stream(
12879                        e,
12880                        &mut *cache,
12881                        &snap,
12882                        &ck,
12883                        &stream_acc,
12884                        1,
12885                        t_v_s,
12886                    )?;
12887                    e.spec_rollback_stream(
12888                        ptrs,
12889                        &pos_start_d,
12890                        &stream_acc,
12891                        1,
12892                        self.layers.len() + 1,
12893                    )?;
12894                    e.spec_ring_commit(&vtok_d, &stream_acc, &brk_d, &mut ring_d, &mut pend_d)?;
12895                }
12896                e.stream().synchronize()?;
12897                let ring_h = e.dtoh_u32(&ring_d)?;
12898                let cnt = ring_h[0] as usize;
12899                for i in 0..cnt {
12900                    if out.len() < max_new {
12901                        out.push(ring_h[1 + i]);
12902                    }
12903                }
12904                let pos_h = e.dtoh_i32(&pos_ctr)?[0] as usize;
12905                for il in 0..self.layers.len() {
12906                    if let Some(kvl) = cache.kv[il].as_mut() {
12907                        kvl.len = pos_h;
12908                    }
12909                }
12910                cache.pos = pos_h;
12911                scratch.kv.len = pos_h;
12912                pending = Some(ring_h[cnt]); // last drained token = the live bonus
12913                last_token = ring_h[cnt];
12914                total_drafted += k * m_rounds; // upper bound (p-min breaks uncounted)
12915                total_accepted += cnt.saturating_sub(m_rounds);
12916                if let Some(t) = sess_telem {
12917                    // totals only — the burst's per-round accept counts stayed on device
12918                    // (that is the point of the round-stream arm). pos_* untouched.
12919                    t.record_totals(m_rounds, k * m_rounds, cnt.saturating_sub(m_rounds));
12920                }
12921                round += m_rounds;
12922                // sse-cadence: the drained ring is committed — flush it at burst-drain cadence.
12923                keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
12924                continue;
12925            }
12926            let pipe_draft = match pipe {
12927                Some(p) => Some(p.draft_begin(round)?),
12928                None => None,
12929            };
12930            let pos = cache.pos; // #tokens committed (EXCLUDES a pending bonus)
12931            let mut current_opti = carried_opti.take();
12932            let mut fork_generation = if current_opti.is_none() && pending.is_some() {
12933                match opti_fork.as_mut() {
12934                    Some(fork) if fork.mode.is_forced() => Some(fork.reserve(&mut snap)?),
12935                    None => None,
12936                    Some(_) => None,
12937                }
12938            } else {
12939                None
12940            };
12941            if current_opti.is_none() {
12942                if let Some(fork) = opti_fork.as_ref() {
12943                    opti_snapshot_stage_owned_into(e, cache, fork.rt, &fork.fence, &mut snap)?;
12944                } else {
12945                    cache.snapshot_into(e, &mut snap)?;
12946                }
12947            } else if snap.pos != pos {
12948                return Err(format!(
12949                    "optipipe carried snapshot pos {} != current pos {pos}",
12950                    snap.pos
12951                )
12952                .into());
12953            } // §C: snapshot BEFORE draft+verify (already retained for a carried successor)
12954            ph_mark(&mut ph_rest, phase_on);
12955
12956            // --- 1. DRAFT k tokens with the NextN head (autoregressive, T=1 each) ---
12957            // p-min semantics (both paths): stop the chain early when the head's confidence in
12958            // its own pick drops below p_min — the just-drafted token is DISCARDED, but its
12959            // scratch append stands (identical to the eager chain's ordering). j==0 always drafts.
12960            let base0 = if pending.is_some() { 1usize } else { 0usize };
12961            // fixed draft length by default; MEMRA_SPEC_ADAPT=1 drafts at last round's
12962            // accepted run + 1 (the gemma law — see the setup block above the loop).
12963            let k_this = if adapt { kc } else { k };
12964            let mut draft: Vec<u32> = Vec::with_capacity(k);
12965            let mut draft_idx: Vec<u32> = Vec::with_capacity(k); // trimmed-vocab ids (== draft when untrimmed)
12966            let mut controller_draft_prob: Option<f32> = None;
12967            let mut controller_eager_state: Option<(u32, CudaSlice<f32>)> = None;
12968            if let Some(ticket) = current_opti.as_mut() {
12969                let carried_pending = pending.ok_or("optipipe carried successor lost pending")?;
12970                if ticket.verify_tokens[0] != carried_pending {
12971                    return Err(format!(
12972                        "optipipe carried pending mismatch: ticket={} live={carried_pending}",
12973                        ticket.verify_tokens[0],
12974                    )
12975                    .into());
12976                }
12977                draft.push(ticket.verify_tokens[1]);
12978                controller_draft_prob = Some(ticket.draft_prob);
12979                controller_eager_state = ticket
12980                    .take_eager_seed()
12981                    .map(|seed| (ticket.verify_tokens[1], seed));
12982            } else {
12983                // Round-start draft-KV sync (BOTH paths). Persistent: truncate/align to the committed
12984                // history — slots 0..P hold entries for the tokens before last_token@P (P = pos +
12985                // base0 - 1); this single set_len IS the draft-side rollback (drops last round's
12986                // rejected drafts and p-min extras via the len mechanism).
12987                scratch.set_len(e, pos + base0 - 1)?;
12988                // dcw door: a captured chain appends k_this device-counter rows (plus the
12989                // pseudo-seed replay) with no host intervention; any ring rebase those appends
12990                // could need happens HERE, host-side, before the replays. The eager arm keeps
12991                // its own per-step prepare, so this is graph-path-only work.
12992                if step35_draft_dcw_on()
12993                    && (dctx.graph.is_some()
12994                        || dctx.graph_s.is_some()
12995                        || dctx.chain.is_some()
12996                        || dctx.chain_s.is_some())
12997                {
12998                    scratch.ensure_dcw_headroom(e, k_this + 2)?;
12999                }
13000                if pen_on {
13001                    // PEN_WINDOW_MAX also bounds the per-round upload and the O(n_hist^2)
13002                    // device dedup: the serve window is already PEN_WINDOW_MAX, and this
13003                    // defensive min also bounds non-server callers.
13004                    let win = sp.penalty_last_n.min(PEN_WINDOW_MAX);
13005                    let w0 = pen_hist.len().saturating_sub(win);
13006                    pen_hist_d = Some(e.htod_u32_v(&pen_hist[w0..])?);
13007                }
13008                if sampled {
13009                    draft_logits.clear();
13010                    draft_stats.clear();
13011                }
13012                // DRAFT-SIDE GRAMMAR MASK: clone the committed grammar state ONCE per round; each
13013                // position's mask is computed on that clone and advanced by the PROPOSED token. The
13014                // real state moves only on emission (verify's job), so the emitted stream is
13015                // unchanged — the mask only removes tokens the verify would have truncated anyway.
13016                let mut dmask_live = dmask_on;
13017                if dmask_live {
13018                    let t_c = std::time::Instant::now();
13019                    constraint
13020                        .as_deref_mut()
13021                        .unwrap()
13022                        .draft_begin()
13023                        .map_err(|e2| format!("constraint: {e2}"))?;
13024                    dm_clone_ns += t_c.elapsed().as_nanos();
13025                    dm_rounds += 1;
13026                }
13027                if let (false, Some(cg)) = (sampled || pen_on, &dctx.chain) {
13028                    // GREEDY CHAIN GRAPH (lane/step37-draft-graph-serving-20260830): the
13029                    // eager multi-head chain's EXACT launch order — step j rewinds head
13030                    // (j % heads)'s plane to the committed length and replays rows 0..=j —
13031                    // with each row's whole head-forward as ONE graph launch. The chain
13032                    // POLICY (head choice, prefix length, stored-seed feed) is host-side,
13033                    // identical to `mtp_chain_forward_dev`, so graph-vs-eager drafts are
13034                    // bit-identical by construction (same launcher, same bucket — the dcw
13035                    // parity contract). Interior rows launch the head-less graph: their
13036                    // logits are dead in the eager chain too, so the consumed bytes match.
13037                    let heads_n = self.mtp_head_count();
13038                    let committed = pos + base0 - 1;
13039                    let mut chain_tokens: Vec<u32> = vec![last_token];
13040                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13041                    for j in 0..k_this {
13042                        let index = mtp_chain_head_index(j, heads_n);
13043                        if debug_spec {
13044                            eprintln!(
13045                                "[mtp-chain-step] round={round} j={j} head={index} \
13046                                 replay_rows={} arm=graph",
13047                                chain_tokens.len(),
13048                            );
13049                        }
13050                        scratch.set_plane_len(e, index, committed)?;
13051                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13052                        for row in 0..=j {
13053                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13054                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13055                            if row < j {
13056                                cg.interior[index].launch()?;
13057                            } else {
13058                                // per-position mask upload before the LAST row only — the
13059                                // eager chain applies the mask on is_last exactly the same.
13060                                if dmask_live
13061                                    && !upload_draft_mask(
13062                                        e,
13063                                        constraint.as_deref_mut().unwrap(),
13064                                        &mut dctx.g_dmask,
13065                                        mtp.d2t.as_ref(),
13066                                        d_vocab,
13067                                        dmask_words,
13068                                    )?
13069                                {
13070                                    e.htod_u32_into(
13071                                        &mut dctx.g_dmask,
13072                                        &vec![u32::MAX; dmask_words],
13073                                    )?;
13074                                    dmask_live = false;
13075                                }
13076                                cg.last[index].launch()?;
13077                            }
13078                            // host mirror (len_d advanced in-graph by the dcw append)
13079                            scratch.plane_mut(index).0.len += 1;
13080                        }
13081                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13082                        // #87 SENTINEL TRAP (see the single-head graph arm below).
13083                        if (idx as usize) >= d_vocab {
13084                            let seed_h = e.dtoh(&dctx.g_seed)?;
13085                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13086                            return Err(format!(
13087                                "draft(chain-graph) argmax sentinel 0x{idx:08x} >= d_vocab \
13088                             {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13089                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference \
13090                             the embed row (#87 trap)"
13091                            )
13092                            .into());
13093                        }
13094                        // multi-head MTP forbids a trimmed head (validated at entry), so the
13095                        // draft index IS the target id; keep the map for uniformity.
13096                        let d = match &mtp.d2t {
13097                            Some(map) => map[idx as usize],
13098                            None => idx,
13099                        };
13100                        let draft_p = if p_min > 0.0 {
13101                            Some(e.dtoh(&dctx.g_p)?[0])
13102                        } else {
13103                            None
13104                        };
13105                        if j == 0 {
13106                            controller_draft_prob = draft_p;
13107                        }
13108                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13109                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13110                                break;
13111                            }
13112                        }
13113                        draft.push(d);
13114                        chain_tokens.push(d);
13115                        // step j's h_nextn: the last-row graph self-fed it into g_seed —
13116                        // snapshot it as the chain history seed for row j+1 (stream-ordered
13117                        // after the launch, exactly the eager chain's chain_seeds push).
13118                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13119                        // speculative grammar advance (see the single-head graph arm).
13120                        if dmask_live
13121                            && !constraint
13122                                .as_deref_mut()
13123                                .unwrap()
13124                                .draft_advance(d)
13125                                .map_err(|e2| format!("constraint: {e2}"))?
13126                        {
13127                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13128                            break;
13129                        }
13130                    }
13131                } else if let (true, Some(cg)) = (
13132                    sampled && s_capturable && dctx.s_key == Some(s_key),
13133                    &dctx.chain_s,
13134                ) {
13135                    if skey_probe() {
13136                        eprintln!(
13137                            "[skey] chain=graph_chain_s round={round} capturable={} top_k={} \
13138                             top_p={} min_p={} s_key_parked={:?}",
13139                            s_capturable as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13140                        );
13141                    }
13142                    // SAMPLED CHAIN GRAPH: the greedy chain arm's launch order with the
13143                    // sampled last-row graphs — in-graph counter bump + (filtered) gumbel
13144                    // draw + argmax; q retained per step into q_slots exactly like the
13145                    // single-head sampled graph arm. Counter continuity: g_ctr host-seeded
13146                    // to sctr-1 once per ROUND; each step's last-row graph bumps it BEFORE
13147                    // the perturb, so step j consumes counter sctr+j — the eager Philox
13148                    // stream (interior rows never draw, never bump).
13149                    let heads_n = self.mtp_head_count();
13150                    let committed = pos + base0 - 1;
13151                    let filtered_stats_in_graph = s_key.filtered();
13152                    let mut chain_tokens: Vec<u32> = vec![last_token];
13153                    let mut chain_seed_bufs: Vec<CudaSlice<f32>> = vec![e.clone_dtod(&h_seed_buf)?];
13154                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13155                    for j in 0..k_this {
13156                        let index = mtp_chain_head_index(j, heads_n);
13157                        if debug_spec {
13158                            eprintln!(
13159                                "[mtp-chain-step] round={round} j={j} head={index} \
13160                                 replay_rows={} arm=graph_s",
13161                                chain_tokens.len(),
13162                            );
13163                        }
13164                        scratch.set_plane_len(e, index, committed)?;
13165                        e.set_i32_one(&mut dctx.g_pos, (committed + 1) as i32)?;
13166                        for row in 0..=j {
13167                            e.set_u32_one(&mut dctx.g_tok, chain_tokens[row])?;
13168                            e.copy_into(&mut dctx.g_seed, 0, &chain_seed_bufs[row], n_embd)?;
13169                            if row < j {
13170                                cg.interior[index].launch()?;
13171                            } else {
13172                                cg.last[index].launch()?;
13173                            }
13174                            scratch.plane_mut(index).0.len += 1;
13175                        }
13176                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13177                        // counts the p-min-discarded token too)
13178                        // q retention: ONE async D2D of the persistent head-logits buffer
13179                        // into this round's slot j (stream-ordered after the replay).
13180                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13181                        // FILTERED capture: read the in-graph filter_stats scalars back per
13182                        // replay instead of a second full-vocab filter_stats per slot post-
13183                        // chain — bit-exact (the values the in-graph perturb consumed) and
13184                        // measured worth ~5% of vendor-default serving tok/s at K=3. Before
13185                        // the p-min break so the discarded slot's stats land too.
13186                        if filtered_stats_in_graph {
13187                            draft_stats.push((
13188                                e.dtoh(&dctx.g_mx)?[0],
13189                                e.dtoh(&dctx.g_th)?[0],
13190                                e.dtoh(&dctx.g_z)?[0],
13191                            ));
13192                        }
13193                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13194                        // #87 SENTINEL TRAP (see the single-head graph arms).
13195                        if (idx as usize) >= d_vocab {
13196                            let seed_h = e.dtoh(&dctx.g_seed)?;
13197                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13198                            return Err(format!(
13199                                "draft(chain-graph-sampled) argmax sentinel 0x{idx:08x} >= \
13200                             d_vocab {d_vocab} at round {round} j={j} head={index} pos={pos}: \
13201                             head-out NaN {seed_nan}/{n_embd} — refusing to dereference the \
13202                             embed row (#87 trap)"
13203                            )
13204                            .into());
13205                        }
13206                        let d = match &mtp.d2t {
13207                            Some(map) => map[idx as usize],
13208                            None => idx,
13209                        };
13210                        draft_idx.push(idx);
13211                        if p_min > 0.0 {
13212                            let p = e.dtoh(&dctx.g_p)?[0];
13213                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13214                                break;
13215                            }
13216                        }
13217                        draft.push(d);
13218                        chain_tokens.push(d);
13219                        chain_seed_bufs.push(e.clone_dtod(&dctx.g_seed)?);
13220                    }
13221                    // PURE-TEMP accept path: stats per used slot recomputed from the RETAINED
13222                    // q with the SAME filter_stats program the eager arm runs (deployment-
13223                    // keyed coop/plain choice, same input bits). The FILTERED graph read its
13224                    // stats back per replay above.
13225                    if !filtered_stats_in_graph {
13226                        for j in 0..draft.len().max(draft_idx.len()) {
13227                            let rows0 = e.htod_i32(&[0])?;
13228                            let (mut th_d, mut z_d, mut mx_d) =
13229                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13230                            e.filter_stats(
13231                                &dctx.q_slots[j],
13232                                d_vocab,
13233                                &rows0,
13234                                &mut th_d,
13235                                &mut z_d,
13236                                &mut mx_d,
13237                                d_vocab,
13238                                1,
13239                                sp_temp,
13240                                sp.top_k,
13241                                sp.top_p,
13242                                sp.min_p,
13243                            )?;
13244                            draft_stats.push((
13245                                e.dtoh(&mx_d)?[0],
13246                                e.dtoh(&th_d)?[0],
13247                                e.dtoh(&z_d)?[0],
13248                            ));
13249                        }
13250                    }
13251                } else if let (false, Some(gr)) = (sampled || pen_on, &dctx.graph) {
13252                    // GRAPH DRAFT: one dispatch per drafted token. The chain feeds itself on-device
13253                    // (in-graph argmax -> tok_d -> next replay's embed; h_nextn -> h_seed_d; pos_d
13254                    // inc'd in-graph); the host only reads 4B token (+4B p) and decides the break.
13255                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13256                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
13257                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13258                    for j in 0..k_this {
13259                        // per-position mask upload (contents only — the graph's baked pointer is
13260                        // dctx.g_dmask). All-ones once masking goes dead mid-chain, so the captured
13261                        // mask node degrades to a no-op ban instead of needing a second graph.
13262                        if dmask_live
13263                            && !upload_draft_mask(
13264                                e,
13265                                constraint.as_deref_mut().unwrap(),
13266                                &mut dctx.g_dmask,
13267                                mtp.d2t.as_ref(),
13268                                d_vocab,
13269                                dmask_words,
13270                            )?
13271                        {
13272                            // no draft-vocab row is grammar-legal here (a trimmed FR-Spec head can
13273                            // genuinely miss the legal set): neutralize the captured mask node and
13274                            // finish the chain UNMASKED — exactly pre-lane behaviour, never worse.
13275                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13276                            dmask_live = false;
13277                        }
13278                        gr.launch()?;
13279                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13280                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13281                        // #87 SENTINEL TRAP: an all-NaN head-logits row leaves the device argmax's
13282                        // init sentinel (0x7FFFFFFF) in g_tok — feeding it onward dereferences
13283                        // embed_row(sentinel) = table + ~4.6TB (never mapped) inside the NEXT graph
13284                        // replay's embed node, and the MMU fault kills the CUDA context for the
13285                        // whole process (research/pp2spec-crash-20260807: 3 coredumps, byte-exact
13286                        // VA arithmetic). Refuse loudly instead; the diagnostics name the first-NaN
13287                        // buffer (g_seed = the verify-side handoff vs head-side compute).
13288                        if (idx as usize) >= d_vocab {
13289                            // g_seed is SELF-FED (the replay writes h_nextn back into it), so it
13290                            // reads as the head's OUTPUT at j; h_seed_buf is the round's INPUT
13291                            // seed, untouched since the round-start copy — the pair discriminates
13292                            // "seed arrived poisoned" from "head forward produced NaN".
13293                            let seed_h = e.dtoh(&dctx.g_seed)?;
13294                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13295                            let in_h = e.dtoh(&h_seed_buf)?;
13296                            let in_nan = in_h.iter().filter(|v| v.is_nan()).count();
13297                            return Err(format!(
13298                                "draft(graph) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
13299                             round {round} j={j} pos={pos}: head-out NaN {seed_nan}/{n_embd}, \
13300                             round-input-seed NaN {in_nan}/{n_embd} — refusing to dereference \
13301                             the embed row (#87 trap)"
13302                            )
13303                            .into());
13304                        }
13305                        // trimmed draft vocab -> target token id (identity when no d2t map)
13306                        let d = match &mtp.d2t {
13307                            Some(map) => map[idx as usize],
13308                            None => idx,
13309                        };
13310                        let draft_p = if p_min > 0.0
13311                            || opti_fork
13312                                .as_ref()
13313                                .is_some_and(|fork| fork.controller.is_some())
13314                        {
13315                            Some(e.dtoh(&dctx.g_p)?[0])
13316                        } else {
13317                            None
13318                        };
13319                        if j == 0 {
13320                            controller_draft_prob = draft_p;
13321                        }
13322                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13323                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13324                                break;
13325                            }
13326                        }
13327                        draft.push(d);
13328                        // with a trimmed head the NEXT embed must read the TARGET id, not the draft
13329                        // index the argmax wrote — patch the persistent token buffer (4B htod).
13330                        if d != idx {
13331                            e.set_u32_one(&mut dctx.g_tok, d)?;
13332                        }
13333                        // advance the SPECULATIVE state with the proposal; a dead chain drops to
13334                        // unmasked drafting for the remaining positions (verify still arbitrates).
13335                        // speculative advance; a chain the grammar can no longer follow (EOS
13336                        // proposed) ends here. The captured mask node always runs, so a dead chain
13337                        // leaves the buffer NEUTRAL (all-ones = ban nothing) before it exits.
13338                        if dmask_live
13339                            && !constraint
13340                                .as_deref_mut()
13341                                .unwrap()
13342                                .draft_advance(d)
13343                                .map_err(|e2| format!("constraint: {e2}"))?
13344                        {
13345                            e.htod_u32_into(&mut dctx.g_dmask, &vec![u32::MAX; dmask_words])?;
13346                            break;
13347                        }
13348                    }
13349                // REGIME RE-TEST (lane/graph-s-key-exactness-20260819, widened by
13350                // lane/step37-draft-graph-serving-20260830): the sampled graph is legal ONLY
13351                // in the regime it was captured in. The condition used to read
13352                // `(sampled, &dctx.graph_s)` and trusted `s_key` to have dropped anything
13353                // else — which it could not, because the key omitted the filters. Both
13354                // halves are enforced: the key drops a stale graph, and this site refuses to
13355                // launch one whose key differs or whose regime is uncapturable (penalties).
13356                } else if let (true, Some(gr)) = (
13357                    sampled && s_capturable && dctx.s_key == Some(s_key),
13358                    &dctx.graph_s,
13359                ) {
13360                    if skey_probe() {
13361                        eprintln!(
13362                            "[skey] chain=graph_s round={round} pure_temp={} capturable={} \
13363                             top_k={} top_p={} min_p={} s_key_parked={:?}",
13364                            pure_temp as u8,
13365                            s_capturable as u8,
13366                            sp.top_k,
13367                            sp.top_p,
13368                            sp.min_p,
13369                            dctx.s_key,
13370                        );
13371                    }
13372                    // SAMPLED GRAPH DRAFT: one replay per drafted token — head forward + gumbel +
13373                    // argmax in ONE dispatch; the host reads 4B token (+4B p), D2Ds q into slot j,
13374                    // and decides the break. Event-counter continuity: g_ctr is host-seeded to
13375                    // sctr-1 ONCE per round (outside the graph); the in-graph bump runs BEFORE the
13376                    // perturb, so replay j consumes counter sctr+j — exactly the eager arm's Philox
13377                    // stream. Host sctr advances in lockstep (computed, no readback needed).
13378                    e.set_i32_one(&mut dctx.g_pos, (pos + base0) as i32)?;
13379                    e.set_u32_one(&mut dctx.g_tok, last_token)?;
13380                    e.copy_into(&mut dctx.g_seed, 0, &h_seed_buf, n_embd)?;
13381                    e.set_u32_one(&mut dctx.g_ctr, sctr.wrapping_sub(1))?;
13382                    let filtered_stats_in_graph = s_key.filtered();
13383                    for j in 0..k_this {
13384                        gr.launch()?;
13385                        scratch.kv.len += 1; // host mirror (len_d advanced in-graph)
13386                        sctr += 1; // mirrors the in-graph g_ctr bump (eager parity:
13387                        // counts the p-min-discarded token too)
13388                        // q retention: ONE async D2D of the persistent head-logits buffer into this
13389                        // round's slot j (stream-ordered after the replay, before the next one).
13390                        e.copy_into(&mut dctx.q_slots[j], 0, &dctx.g_q, d_vocab)?;
13391                        // FILTERED capture: the replay's own filter_stats node already computed
13392                        // (th, z, mx) — read the three scalars back instead of paying a SECOND
13393                        // full-vocab filter_stats per slot post-chain (measured ~5% of vendor-
13394                        // default serving tok/s at K=3). Bit-exact by construction: these are
13395                        // the very values the in-graph perturb consumed. Read BEFORE the p-min
13396                        // break so the discarded slot's stats land too (accept-path indexing).
13397                        if filtered_stats_in_graph {
13398                            draft_stats.push((
13399                                e.dtoh(&dctx.g_mx)?[0],
13400                                e.dtoh(&dctx.g_th)?[0],
13401                                e.dtoh(&dctx.g_z)?[0],
13402                            ));
13403                        }
13404                        let idx = e.dtoh_u32_one(&dctx.g_tok)?;
13405                        // #87 SENTINEL TRAP (see the greedy graph arm above).
13406                        if (idx as usize) >= d_vocab {
13407                            let seed_h = e.dtoh(&dctx.g_seed)?;
13408                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13409                            return Err(format!(
13410                                "draft(graph-sampled) argmax sentinel 0x{idx:08x} >= d_vocab \
13411                             {d_vocab} at round {round} j={j} pos={pos}: round-seed NaN \
13412                             {seed_nan}/{n_embd} — refusing to dereference the embed row \
13413                             (#87 trap)"
13414                            )
13415                            .into());
13416                        }
13417                        let d = match &mtp.d2t {
13418                            Some(map) => map[idx as usize],
13419                            None => idx,
13420                        };
13421                        draft_idx.push(idx);
13422                        if p_min > 0.0 {
13423                            let p = e.dtoh(&dctx.g_p)?[0];
13424                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13425                                break;
13426                            }
13427                        }
13428                        draft.push(d);
13429                        // trimmed head: the NEXT embed must read the TARGET id (see the greedy arm).
13430                        if d != idx {
13431                            e.set_u32_one(&mut dctx.g_tok, d)?;
13432                        }
13433                    }
13434                    // PURE-TEMP accept path: fill draft_stats per used slot post-chain (the
13435                    // stats degenerate to th=0 / full-Z; one filter_stats launch per slot).
13436                    // The FILTERED graph read its stats back per replay above.
13437                    if !filtered_stats_in_graph {
13438                        for j in 0..draft.len().max(draft_idx.len()) {
13439                            let rows0 = e.htod_i32(&[0])?;
13440                            let (mut th_d, mut z_d, mut mx_d) =
13441                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13442                            e.filter_stats(
13443                                &dctx.q_slots[j],
13444                                d_vocab,
13445                                &rows0,
13446                                &mut th_d,
13447                                &mut z_d,
13448                                &mut mx_d,
13449                                d_vocab,
13450                                1,
13451                                sp_temp,
13452                                sp.top_k,
13453                                sp.top_p,
13454                                sp.min_p,
13455                            )?;
13456                            draft_stats.push((
13457                                e.dtoh(&mx_d)?[0],
13458                                e.dtoh(&th_d)?[0],
13459                                e.dtoh(&z_d)?[0],
13460                            ));
13461                        }
13462                    }
13463                } else {
13464                    if skey_probe() && sampled {
13465                        eprintln!(
13466                            "[skey] chain=eager round={round} pure_temp={} top_k={} \
13467                             top_p={} min_p={} s_key_parked={:?}",
13468                            pure_temp as u8, sp.top_k, sp.top_p, sp.min_p, dctx.s_key,
13469                        );
13470                    }
13471                    // EAGER DRAFT (fallback: MoE head/trunk, huge k, MEMRA_SPEC_NOGRAPH, capture fail).
13472                    let chain_heads = !self.mtp_extra.is_empty();
13473                    let mut e_tok = last_token;
13474                    let mut d_seed = e.clone_dtod(&h_seed_buf)?;
13475                    let mut chain_tokens = if chain_heads {
13476                        vec![last_token]
13477                    } else {
13478                        Vec::new()
13479                    };
13480                    let mut chain_seeds = if chain_heads {
13481                        vec![e.clone_dtod(&h_seed_buf)?]
13482                    } else {
13483                        Vec::new()
13484                    };
13485                    for j in 0..k_this {
13486                        // GPU-ARGMAX DRAFT (2026-07-03): device logits + device argmax + 4-byte token
13487                        // read instead of the ~600KB full-vocab dtoh + host argmax per draft token.
13488                        let mtp_pos = pos + base0 + j;
13489                        // draft-side grammar mask (eager twin of the graph arm's in-graph node).
13490                        // A position with no legal draft-vocab row drops to unmasked drafting for
13491                        // the rest of the chain (pre-lane behaviour; verify still arbitrates).
13492                        if dmask_live {
13493                            dmask_live = upload_draft_mask(
13494                                e,
13495                                constraint.as_deref_mut().unwrap(),
13496                                &mut dctx.g_dmask,
13497                                mtp.d2t.as_ref(),
13498                                d_vocab,
13499                                dmask_words,
13500                            )?;
13501                        }
13502                        let mask = if dmask_live {
13503                            Some((&dctx.g_dmask, dmask_words))
13504                        } else {
13505                            None
13506                        };
13507                        let (dl_d, h_nextn) = if chain_heads {
13508                            if debug_spec {
13509                                eprintln!(
13510                                    "[mtp-chain-step] round={round} j={j} head={} replay_rows={}",
13511                                    mtp_chain_head_index(j, self.mtp_head_count()),
13512                                    chain_tokens.len(),
13513                                );
13514                            }
13515                            self.mtp_chain_forward_dev(
13516                                e,
13517                                &chain_tokens,
13518                                &chain_seeds,
13519                                &mut *scratch,
13520                                pos + base0 - 1,
13521                                embd_dev,
13522                                mask,
13523                            )?
13524                        } else {
13525                            self.mtp_head_forward_dev(
13526                                e,
13527                                mtp,
13528                                e_tok,
13529                                &d_seed,
13530                                &mut *scratch,
13531                                mtp_pos,
13532                                embd_dev,
13533                                mask,
13534                            )?
13535                        };
13536                        let tok_d = if sampled {
13537                            // FILTERED Gumbel-max: stats -> masked perturb -> argmax = one draw from
13538                            // the filtered softmax (filters off => th=0, exact v1 semantics).
13539                            if perturb_buf.is_none() {
13540                                perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
13541                            }
13542                            let mut q_row = e.clone_dtod(&dl_d)?; // retained q (penalized when on)
13543                            if pen_on {
13544                                let h = pen_hist_d.as_ref().unwrap();
13545                                let nh = h.len();
13546                                e.penalize_logits(
13547                                    &mut q_row,
13548                                    h,
13549                                    nh,
13550                                    sp.penalty_repeat,
13551                                    sp.penalty_freq,
13552                                    sp.penalty_present,
13553                                    d_vocab,
13554                                )?;
13555                            }
13556                            let rows0 = e.htod_i32(&[0])?;
13557                            let (mut th_d, mut z_d, mut mx_d) =
13558                                (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
13559                            e.filter_stats(
13560                                &q_row, d_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, d_vocab,
13561                                1, sp_temp, sp.top_k, sp.top_p, sp.min_p,
13562                            )?;
13563                            let (th, z, mx) =
13564                                (e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0], e.dtoh(&mx_d)?[0]);
13565                            let pb = perturb_buf.as_mut().unwrap();
13566                            e.gumbel_perturb_filtered(
13567                                &q_row, pb, d_vocab, sp_seed, sctr, sp_temp, mx, th,
13568                            )?;
13569                            sctr += 1;
13570                            draft_logits.push(q_row);
13571                            draft_stats.push((mx, th, z));
13572                            e.argmax_token_device(pb, d_vocab)?
13573                        } else {
13574                            e.argmax_token_device(&dl_d, d_vocab)?
13575                        };
13576                        let idx = e.dtoh_u32_one(&tok_d)?;
13577                        // #87 SENTINEL TRAP (eager twin — see the graph arm). Extra diagnostics
13578                        // here because the eager chain's operands are all readable: dl_d (the head
13579                        // logits row) and d_seed (this step's h_seed) name the first-NaN buffer.
13580                        if (idx as usize) >= d_vocab {
13581                            let dl_h = e.dtoh(&dl_d)?;
13582                            let dl_nan = dl_h.iter().filter(|v| v.is_nan()).count();
13583                            let seed_h = if chain_heads {
13584                                e.dtoh(chain_seeds.last().unwrap())?
13585                            } else {
13586                                e.dtoh(&d_seed)?
13587                            };
13588                            let seed_nan = seed_h.iter().filter(|v| v.is_nan()).count();
13589                            return Err(format!(
13590                                "draft(eager) argmax sentinel 0x{idx:08x} >= d_vocab {d_vocab} at \
13591                             round {round} j={j} pos={pos}: head-logits NaN {dl_nan}/{d_vocab}, \
13592                             step-seed NaN {seed_nan}/{n_embd} — refusing to dereference the \
13593                             embed row (#87 trap)"
13594                            )
13595                            .into());
13596                        }
13597                        let d = match &mtp.d2t {
13598                            Some(map) => map[idx as usize],
13599                            None => idx,
13600                        };
13601                        if sampled {
13602                            draft_idx.push(idx);
13603                        }
13604                        let draft_p = if p_min > 0.0
13605                            || opti_fork
13606                                .as_ref()
13607                                .is_some_and(|fork| fork.controller.is_some())
13608                        {
13609                            let p_d = e.prob_of_token_device(&dl_d, &tok_d, d_vocab)?;
13610                            Some(e.dtoh(&p_d)?[0])
13611                        } else {
13612                            None
13613                        };
13614                        if j == 0 {
13615                            controller_draft_prob = draft_p;
13616                        }
13617                        if let Some(p) = draft_p.filter(|_| p_min > 0.0) {
13618                            if p < p_min && (j > 0 || (pmin0 && base0 == 1)) {
13619                                break;
13620                            }
13621                        }
13622                        draft.push(d);
13623                        if chain_heads {
13624                            chain_tokens.push(d);
13625                            chain_seeds.push(h_nextn);
13626                        } else {
13627                            e_tok = d;
13628                            d_seed = h_nextn;
13629                        }
13630                        // speculative advance; a chain the grammar can no longer follow (EOS
13631                        // proposed) ends here — the prefix already proposed still rides verify.
13632                        if dmask_live
13633                            && !constraint
13634                                .as_deref_mut()
13635                                .unwrap()
13636                                .draft_advance(d)
13637                                .map_err(|e2| format!("constraint: {e2}"))?
13638                        {
13639                            break;
13640                        }
13641                    }
13642                    if !chain_heads
13643                        && opti_fork
13644                            .as_ref()
13645                            .is_some_and(|fork| fork.controller.is_some())
13646                    {
13647                        controller_eager_state = Some((e_tok, d_seed));
13648                    }
13649                }
13650            }
13651            let k_round = draft.len();
13652            if let Some(p) = pipe {
13653                p.draft_end(round);
13654            }
13655            drop(pipe_draft);
13656
13657            ph_mark(&mut ph_draft, phase_on);
13658            // --- 2. VERIFY: one batched target forward. With a pending bonus, it rides as col 0
13659            //         (committing its KV/recur inside the SAME weight read); drafts follow. ---
13660            let verify_tokens: Vec<u32> = match pending {
13661                Some(b) => {
13662                    let mut v = Vec::with_capacity(k_round + 1);
13663                    v.push(b);
13664                    v.extend_from_slice(&draft);
13665                    v
13666                }
13667                None => draft.clone(),
13668            };
13669            let base = if pending.is_some() { 1 } else { 0 };
13670            // ckpt (REPLAY-FREE partial accept): retain per-layer state-rebuild inputs alongside
13671            // the verify. Pure buffer keep-alives + dtod clones — kernel work is unchanged.
13672            let mut ckpt = if let Some(ticket) = current_opti.as_mut() {
13673                Some(ticket.take_ckpt())
13674            } else if spec_replay {
13675                None
13676            } else {
13677                Some(VerifyCkpt::new(self.layers.len()))
13678            };
13679            let controller_can_probe = base == 1
13680                && k_round == 1
13681                && out.len().saturating_add(2) < max_new
13682                && controller_draft_prob.is_some()
13683                && opti_fork
13684                    .as_ref()
13685                    .and_then(|fork| fork.controller.as_ref())
13686                    .is_some_and(|policy| !policy.breaker_tripped);
13687            let mut successor_attempt: Option<OptiControllerTicket> = None;
13688            let mut rejected_probe: Option<(f32, u32)> = None;
13689            let mut controller_prepared: Option<OptiControllerPrepared> = None;
13690            if controller_can_probe {
13691                // Prepare d2/q and, on admission, d3 before either current verify half is
13692                // issued. N stage 0 can then be followed immediately by N+1 stage 0; once N's
13693                // boundary fires, those dev0 launches overlap N stage 1 on dev1. Preparing on
13694                // the primary stream after N stage 1 would serialize the supposed pipeline.
13695                let eager_pos = scratch.kv.len + 1;
13696                let (optimistic_pending, pending_probability) = self.opti_controller_draft_step(
13697                    e,
13698                    mtp,
13699                    &mut dctx,
13700                    &mut *scratch,
13701                    d_vocab,
13702                    &mut controller_eager_state,
13703                    eager_pos,
13704                    embd_dev,
13705                )?;
13706                let first_probability = controller_draft_prob
13707                    .ok_or("optipipe controller probe lost first-token probability")?;
13708                let q_proxy = first_probability * pending_probability;
13709                OPTI_GATE_CHECKS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13710                OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13711                let admitted = opti_fork
13712                    .as_ref()
13713                    .and_then(|fork| fork.controller.as_ref())
13714                    .ok_or("optipipe controller policy disappeared")?
13715                    .admit(q_proxy);
13716                if admitted {
13717                    OPTI_GATE_ADMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13718                    let eager_pos = scratch.kv.len + 1;
13719                    let (optimistic_draft, optimistic_draft_probability) = self
13720                        .opti_controller_draft_step(
13721                            e,
13722                            mtp,
13723                            &mut dctx,
13724                            &mut *scratch,
13725                            d_vocab,
13726                            &mut controller_eager_state,
13727                            eager_pos,
13728                            embd_dev,
13729                        )?;
13730                    OPTI_SHADOW_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13731                    let eager_seed = controller_eager_state.take().map(|(token, seed)| {
13732                        debug_assert_eq!(token, optimistic_draft);
13733                        seed
13734                    });
13735                    controller_prepared = Some(OptiControllerPrepared {
13736                        verify_tokens: [optimistic_pending, optimistic_draft],
13737                        draft_prob: optimistic_draft_probability,
13738                        eager_seed,
13739                        q_proxy,
13740                        scratch_len: scratch.kv.len,
13741                    });
13742                } else {
13743                    OPTI_GATE_REJECTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13744                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13745                    rejected_probe = Some((q_proxy, optimistic_pending));
13746                    eprintln!(
13747                        "[opti-controller] reject q={q_proxy:.6} threshold={:.3}",
13748                        opti_fork
13749                            .as_ref()
13750                            .and_then(|fork| fork.controller.as_ref())
13751                            .expect("controller policy")
13752                            .threshold,
13753                    );
13754                }
13755            }
13756            let fork_attempt = match fork_generation.take() {
13757                Some(generation) if base == 1 && k_round == 1 => Some(generation),
13758                Some(generation) => {
13759                    opti_fork
13760                        .as_mut()
13761                        .expect("fork generation without fork state")
13762                        .retire(generation)?;
13763                    None
13764                }
13765                None => None,
13766            };
13767            let (tlogits_d, vx) = if let Some(p) = pipe {
13768                self.decode_step_t_core_pipelined(
13769                    e,
13770                    &verify_tokens,
13771                    pos,
13772                    &mut *cache,
13773                    embd_dev,
13774                    ckpt.as_mut(),
13775                    p,
13776                    round,
13777                )?
13778            } else if controller_can_probe {
13779                let fence = opti_fork
13780                    .as_ref()
13781                    .ok_or("optipipe controller probe lost fork state")?
13782                    .fence;
13783                let boundary = match current_opti.as_mut() {
13784                    Some(ticket) => ticket.take_boundary(),
13785                    None => self.verify_stage0_issue(
13786                        e,
13787                        &verify_tokens,
13788                        pos,
13789                        &mut *cache,
13790                        embd_dev,
13791                        ckpt.as_mut(),
13792                        None,
13793                        &fence,
13794                        Some(true),
13795                        None,
13796                    )?,
13797                };
13798                if let Some(prepared) = controller_prepared.take() {
13799                    let generation = {
13800                        let fork = opti_fork
13801                            .as_mut()
13802                            .ok_or("optipipe controller admission lost fork state")?;
13803                        let generation = fork.reserve_successor()?;
13804                        let rt = fork.rt;
13805                        let snapshot_fence = fork.fence;
13806                        opti_snapshot_one_stage_owned_into(
13807                            e,
13808                            cache,
13809                            rt,
13810                            &snapshot_fence,
13811                            0,
13812                            fork.successor_snapshot_mut(),
13813                        )?;
13814                        generation
13815                    };
13816                    let mut successor_ckpt = VerifyCkpt::new(self.layers.len());
13817                    let successor_boundary = self.verify_stage0_issue(
13818                        e,
13819                        &prepared.verify_tokens,
13820                        pos + verify_tokens.len(),
13821                        &mut *cache,
13822                        embd_dev,
13823                        Some(&mut successor_ckpt),
13824                        None,
13825                        &fence,
13826                        Some(false),
13827                        None,
13828                    )?;
13829                    OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13830                    let fork = opti_fork
13831                        .as_ref()
13832                        .ok_or("optipipe controller ticket lost fork state")?;
13833                    successor_attempt = Some(fork.controller_ticket(
13834                        generation,
13835                        successor_boundary,
13836                        successor_ckpt,
13837                        prepared.verify_tokens,
13838                        prepared.draft_prob,
13839                        prepared.eager_seed,
13840                        prepared.q_proxy,
13841                        prepared.scratch_len,
13842                    ));
13843                    eprintln!(
13844                        "[opti-controller] issue generation={} q={:.6} threshold={:.3} \
13845                         verify={:?}",
13846                        generation.id,
13847                        prepared.q_proxy,
13848                        fork.controller.expect("controller policy").threshold,
13849                        prepared.verify_tokens,
13850                    );
13851                }
13852                let result = self.verify_stage1_finish(
13853                    e,
13854                    boundary,
13855                    &mut *cache,
13856                    ckpt.as_mut(),
13857                    None,
13858                    &fence,
13859                    successor_attempt.is_none(),
13860                )?;
13861                if let Some(ticket) = current_opti.as_mut() {
13862                    ticket.settle();
13863                }
13864                if successor_attempt.is_some() {
13865                    let fork = opti_fork
13866                        .as_mut()
13867                        .ok_or("optipipe successor snapshot lost fork state")?;
13868                    let rt = fork.rt;
13869                    let snapshot_fence = fork.fence;
13870                    opti_snapshot_one_stage_owned_into(
13871                        e,
13872                        cache,
13873                        rt,
13874                        &snapshot_fence,
13875                        1,
13876                        fork.successor_snapshot_mut(),
13877                    )?;
13878                    // Publish N only after both independent successor-state queues are complete.
13879                    fork.rt.publish_to(1, &e.stream())?;
13880                }
13881                result
13882            } else if let Some(ticket) = current_opti.as_mut() {
13883                let fork = opti_fork
13884                    .as_mut()
13885                    .ok_or("optipipe carried controller ticket lost fork state")?;
13886                let boundary = ticket.take_boundary();
13887                let result = self.verify_stage1_finish(
13888                    e,
13889                    boundary,
13890                    &mut *cache,
13891                    ckpt.as_mut(),
13892                    None,
13893                    &fork.fence,
13894                    true,
13895                )?;
13896                ticket.settle();
13897                result
13898            } else if let Some(generation) = fork_attempt {
13899                let fork = opti_fork
13900                    .as_mut()
13901                    .expect("fork generation without fork state");
13902                fork.capture_seed(e, generation, &h_seed_buf, &fill_prev, scratch.kv.len)?;
13903                let action = fork.mode.action(generation.id);
13904                let boundary = self.verify_stage0_issue(
13905                    e,
13906                    &verify_tokens,
13907                    pos,
13908                    &mut *cache,
13909                    embd_dev,
13910                    ckpt.as_mut(),
13911                    None,
13912                    &fork.fence,
13913                    Some(true),
13914                    None,
13915                )?;
13916                OPTI_FORK_ATTEMPTS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
13917                let mut ticket = fork.ticket(generation, boundary);
13918                if action == OptiForkAction::Abort {
13919                    return Err(format!(
13920                        "optipipe forced abort with generation {} stage0 in flight",
13921                        generation.id,
13922                    )
13923                    .into());
13924                }
13925                fork.reconcile(
13926                    e,
13927                    &mut *cache,
13928                    &mut *scratch,
13929                    &snap,
13930                    &mut h_seed_buf,
13931                    &mut fill_prev,
13932                    generation,
13933                    action,
13934                    verify_tokens[0],
13935                )?;
13936                let result = if action == OptiForkAction::Hit {
13937                    let boundary = ticket.take_boundary();
13938                    self.verify_stage1_finish(
13939                        e,
13940                        boundary,
13941                        &mut *cache,
13942                        ckpt.as_mut(),
13943                        None,
13944                        &fork.fence,
13945                        true,
13946                    )?
13947                } else {
13948                    // The optimistic boundary slot has no reader. Re-run the unchanged serial
13949                    // verify only after E_restart published the restored stage-0 state.
13950                    self.decode_step_t_core(
13951                        e,
13952                        &verify_tokens,
13953                        pos,
13954                        &mut *cache,
13955                        embd_dev,
13956                        ckpt.as_mut(),
13957                    )?
13958                };
13959                ticket.settle();
13960                debug_assert_eq!(ticket.generation, generation);
13961                fork.retire(generation)?;
13962                result
13963            } else {
13964                // The serial verify every non-fork round takes — the MTP route's
13965                // verify-graph door. The pool is None unless MEMRA_SPEC_VERIFY_GRAPH armed
13966                // a pool above, and then the walk replays the captured trunk instead of
13967                // re-issuing it launch by launch.
13968                let vg_round = if verify_tokens.len() <= vg_t_cap {
13969                    vg_guard.as_mut().and_then(|g| g.as_mut())
13970                } else {
13971                    if let Some(g) = vg_guard.as_mut().and_then(|g| g.as_mut()) {
13972                        // The commit reads this flag to pick its arm; a round that declines
13973                        // the pool must not inherit a stale `true` from the round before it.
13974                        g.round_slab = false;
13975                    }
13976                    None
13977                };
13978                self.decode_step_t_core_vg(
13979                    e,
13980                    &verify_tokens,
13981                    pos,
13982                    &mut *cache,
13983                    embd_dev,
13984                    ckpt.as_mut(),
13985                    vg_round,
13986                )?
13987            };
13988            let pipe_accept = match pipe {
13989                Some(p) => Some(p.accept_begin(round)?),
13990                None => None,
13991            };
13992
13993            if phase_sync {
13994                e.stream().synchronize()?;
13995            }
13996            ph_mark(&mut ph_verify, phase_on);
13997            // --- 3. GREEDY ACCEPT (walk prefix, stop at first mismatch) ---
13998            // DEVICE-ARGMAX ACCEPT: argmax every verify column ON DEVICE (same 2-pass kernels +
13999            // smallest-index tie-break as host argmax, argmax_gate-validated) and read back ONE
14000            // [T] u32 — replaces the T x n_vocab f32 dtoh + T host argmaxes per round.
14001            // t_pred[j] = target's greedy prediction for the slot after draft[j-1] (j>=1) or after
14002            // last_token (j==0). With a pending bonus, col 0 IS the prediction after last_token
14003            // (== the bonus), so every index shifts by `base` and last_pred is unused.
14004            let t_v = verify_tokens.len();
14005            let mut preds: Vec<u32> = Vec::new();
14006            if !sampled {
14007                for j in 0..t_v {
14008                    e.argmax_token_device_col(&tlogits_d, j, n_vocab, &mut preds_d, j)?;
14009                }
14010                preds = e.dtoh_u32(&preds_d)?; // <- the verify-GPU wait lands here
14011                // #87 SENTINEL TRAP, verify side: a sentinel pred becomes the round's bonus =
14012                // next round's last_token = the next chain's embed lookup. Catch it at the
14013                // source with the column named — an all-NaN VERIFY column implicates the
14014                // stage-split trunk (decode_step_t_core_ppn), not the draft head.
14015                if let Some(bad) = preds[..t_v].iter().position(|&p| (p as usize) >= n_vocab) {
14016                    let col = &tlogits_d.slice(bad * n_vocab..(bad + 1) * n_vocab);
14017                    let mut probe = e.zeros(n_vocab)?;
14018                    e.copy_view_into(&mut probe, 0, col, n_vocab)?;
14019                    let col_h = e.dtoh(&probe)?;
14020                    let col_nan = col_h.iter().filter(|v| v.is_nan()).count();
14021                    return Err(format!(
14022                        "verify argmax sentinel 0x{:08x} >= n_vocab {n_vocab} at round {round} \
14023                         col {bad}/{t_v} pos={pos}: verify-logits col NaN {col_nan}/{n_vocab} \
14024                         — the verify TRUNK produced a poisoned column (#87 trap). Run \
14025                         MEMRA_SPEC_NAN_SCAN=1 to name the layer that creates it (=2 to split \
14026                         that layer into attention and routed MoE). NOT the draft head, and NOT \
14027                         the PP stage split this message used to name: pp_cuts() returns None \
14028                         without MEMRA_PP_STAGES, so decode_step_t_core_ppn never runs unless \
14029                         that variable is set.",
14030                        preds[bad]
14031                    )
14032                    .into());
14033                }
14034            }
14035            ph_mark(&mut ph_wait, phase_on);
14036            let t_pred = |j: usize| -> u32 {
14037                if j == 0 && base == 0 {
14038                    last_pred
14039                } else {
14040                    // GREEDY-ONLY: `preds` is filled under `if !sampled` above. The debug print
14041                    // used to call this from the sampled arm and panicked the worker; it now goes
14042                    // through `debug_t_pred0`. Keep the strict index here — in the greedy walk an
14043                    // out-of-range pred is a real bug, not something to paper over.
14044                    debug_assert!(
14045                        !sampled,
14046                        "t_pred is greedy-only: `preds` is empty in the sampled arm"
14047                    );
14048                    preds[base + j - 1]
14049                }
14050            };
14051            let mut devacc_seeded = false;
14052            let mut devacc_acc: Option<CudaSlice<u32>> = None;
14053            let (n_acc, bonus) = if !sampled {
14054                // ROUND-STREAM stage (a) (MEMRA_SPEC_DEVACC=1 opt-in): the walk runs ON DEVICE
14055                // (spec_accept_greedy, verbatim rule) and the host reads back 8B (n_acc, bonus)
14056                // instead of the [T] preds. Same sync count — machinery for stages (b)/(c),
14057                // gated on token identity vs the host walk (the arms below are bit-equal rules).
14058                if crate::spec::spec_devacc() && k_round > 0 && !spec_replay && constraint.is_none()
14059                {
14060                    let draft_d = e.htod_u32_v(&draft)?;
14061                    let mut acc_out = e.alloc_u32_zeroed(2)?;
14062                    e.spec_accept_greedy(
14063                        &preds_d,
14064                        &draft_d,
14065                        last_pred,
14066                        base,
14067                        k_round,
14068                        &mut acc_out,
14069                    )?;
14070                    devacc_acc = Some(acc_out.clone());
14071                    // stage (b): next-round seed gathered ON DEVICE from acc_out before the host
14072                    // ever reads n_acc (j=base+n_acc -> vx col j-1; j==0 -> fill_prev). The three
14073                    // non-replay commit arms skip their host-offset seed copies (guarded below);
14074                    // the legacy spec_replay arm keeps its own rx-based seeding (excluded here).
14075                    // NOTE: fill_prev is NOT updated here — the commit arms' TRUE-HIDDEN
14076                    // REFRESH reads the OLD fill_prev (predecessor of this round's verify batch);
14077                    // the update lands after the arms (devacc_seeded guard below).
14078                    e.spec_seed_gather(&vx, &fill_prev, &acc_out, &mut h_seed_buf, base, n_embd)?;
14079                    // 3a: KV lens roll back on device (len = saved + base + n_acc, all arms'
14080                    // unified rule; full accept rewrites the verify-left value). Host mirrors
14081                    // update after the readback; commit_verified_prefix skips its len_d writes.
14082                    if let Some(successor) = successor_attempt.as_ref() {
14083                        opti_fork
14084                            .as_mut()
14085                            .ok_or("optipipe successor reconcile lost fork state")?
14086                            .queue_actual_reconcile(
14087                                e,
14088                                &snap,
14089                                &acc_out,
14090                                successor.verify_tokens[0],
14091                                base,
14092                            )?;
14093                    } else if let Some(ptrs) = &kv_len_ptrs {
14094                        let saved: Vec<i32> = (0..self.layers.len())
14095                            .map(|il| snap.kv_len[il].map(|v| v as i32).unwrap_or(0))
14096                            .collect();
14097                        let saved_d = e.htod_i32(&saved)?;
14098                        e.spec_rollback_kv(ptrs, &saved_d, &acc_out, base, self.layers.len())?;
14099                    }
14100                    devacc_seeded = true;
14101                    let ab = e.dtoh_u32(&acc_out)?;
14102                    (ab[0] as usize, ab[1])
14103                } else {
14104                    let mut n_acc = 0usize;
14105                    for j in 0..k_round {
14106                        if t_pred(j) == draft[j] {
14107                            n_acc += 1;
14108                        } else {
14109                            break;
14110                        }
14111                    }
14112                    // bonus = target's own token at the first non-accepted slot. n_acc in 0..=k; t_pred
14113                    // is defined for j in 0..=k (j==0 -> last_logits, j>=1 -> col j-1, last col = k-1).
14114                    (n_acc, t_pred(n_acc))
14115                }
14116            } else {
14117                // --- SAMPLED ACCEPT (rejection sampling): u_j < p_j(x_j)/q_j(x_j) walk ---
14118                if col_buf.is_none() {
14119                    col_buf = Some(e.zeros(n_vocab)?);
14120                }
14121                // FILTERED p_j: per-verify-col stats (one batched filter_stats call), then the
14122                // filtered gather. j==0&&base==0 reads last_col (its own stats row appended).
14123                let mut pj = vec![0f32; k_round.max(1)];
14124                let mut col_stats: Vec<(f32, f32, f32)> = Vec::new(); // (max, th, z) per verify col used
14125                if k_round > 0 {
14126                    let mut ids: Vec<u32> = Vec::new();
14127                    let mut rows: Vec<i32> = Vec::new();
14128                    for j in 0..k_round {
14129                        if j > 0 || base == 1 {
14130                            ids.push(draft[j]);
14131                            rows.push((base + j) as i32 - 1);
14132                        }
14133                    }
14134                    if !ids.is_empty() {
14135                        let nr = rows.len();
14136                        // penalties: materialize the used columns into one contiguous penalized
14137                        // buffer (rows remapped 0..nr) so stats+gathers see the penalized p.
14138                        // penalties: materialize used columns contiguously, penalize all rows in
14139                        // one launch, and point stats+gathers at the penalized buffer (rows 0..nr).
14140                        let p_rows: Vec<i32> = if pen_on {
14141                            (0..nr as i32).collect()
14142                        } else {
14143                            rows.clone()
14144                        };
14145                        if pen_on {
14146                            if pcol_buf.as_ref().map(|b| b.len()).unwrap_or(0) < nr * n_vocab {
14147                                pcol_buf = Some(e.zeros(nr * n_vocab)?);
14148                            }
14149                            let pc = pcol_buf.as_mut().unwrap();
14150                            for (i2, &r) in rows.iter().enumerate() {
14151                                let c = r as usize;
14152                                e.copy_view_into(
14153                                    pc,
14154                                    i2 * n_vocab,
14155                                    &tlogits_d.slice(c * n_vocab..(c + 1) * n_vocab),
14156                                    n_vocab,
14157                                )?;
14158                            }
14159                            let h = pen_hist_d.as_ref().unwrap();
14160                            let nh = h.len();
14161                            e.penalize_logits_rows(
14162                                pc,
14163                                h,
14164                                nh,
14165                                sp.penalty_repeat,
14166                                sp.penalty_freq,
14167                                sp.penalty_present,
14168                                n_vocab,
14169                                nr,
14170                            )?;
14171                        }
14172                        let p_src: &CudaSlice<f32> = if pen_on {
14173                            pcol_buf.as_ref().unwrap()
14174                        } else {
14175                            &tlogits_d
14176                        };
14177                        let rowsd = e.htod_i32(&p_rows)?;
14178                        let (mut th_d, mut z_d, mut mx_d) =
14179                            (e.zeros(nr)?, e.zeros(nr)?, e.zeros(nr)?);
14180                        e.filter_stats(
14181                            p_src, n_vocab, &rowsd, &mut th_d, &mut z_d, &mut mx_d, n_vocab, nr,
14182                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14183                        )?;
14184                        let idsd = e.htod_u32_v(&ids)?;
14185                        let mut outd = e.zeros(nr)?;
14186                        e.softmax_gather_filtered(
14187                            p_src, n_vocab, &idsd, &rowsd, &th_d, &z_d, &mut outd, n_vocab, nr,
14188                            sp_temp,
14189                        )?;
14190                        let outv = e.dtoh(&outd)?;
14191                        let (thv, zv, mxv) = (e.dtoh(&th_d)?, e.dtoh(&z_d)?, e.dtoh(&mx_d)?);
14192                        let mut oi = 0usize;
14193                        for j in 0..k_round {
14194                            if j > 0 || base == 1 {
14195                                pj[j] = outv[oi];
14196                                oi += 1;
14197                            }
14198                        }
14199                        col_stats = (0..nr).map(|i| (mxv[i], thv[i], zv[i])).collect();
14200                    }
14201                    if base == 0 {
14202                        let lc: &CudaSlice<f32> = if pen_on {
14203                            if col_buf.is_none() {
14204                                col_buf = Some(e.zeros(n_vocab)?);
14205                            }
14206                            let cb = col_buf.as_mut().unwrap();
14207                            e.copy_into(
14208                                cb,
14209                                0,
14210                                last_col_logits
14211                                    .as_ref()
14212                                    .expect("sampled: last_col_logits unset"),
14213                                n_vocab,
14214                            )?;
14215                            let h = pen_hist_d.as_ref().unwrap();
14216                            let nh = h.len();
14217                            e.penalize_logits(
14218                                cb,
14219                                h,
14220                                nh,
14221                                sp.penalty_repeat,
14222                                sp.penalty_freq,
14223                                sp.penalty_present,
14224                                n_vocab,
14225                            )?;
14226                            col_buf.as_ref().unwrap()
14227                        } else {
14228                            last_col_logits
14229                                .as_ref()
14230                                .expect("sampled: last_col_logits unset")
14231                        };
14232                        let rows0 = e.htod_i32(&[0])?;
14233                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14234                        e.filter_stats(
14235                            lc, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14236                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14237                        )?;
14238                        let idsd = e.htod_u32_v(&[draft[0]])?;
14239                        let mut outd = e.zeros(1)?;
14240                        e.softmax_gather_filtered(
14241                            lc, n_vocab, &idsd, &rows0, &th_d, &z_d, &mut outd, n_vocab, 1, sp_temp,
14242                        )?;
14243                        pj[0] = e.dtoh(&outd)?[0];
14244                        last_col_stats =
14245                            Some((e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0], e.dtoh(&z_d)?[0]));
14246                    }
14247                }
14248                // q source: the graph arms (single-head AND chain) retained the head logits
14249                // in the persistent q_slots; the eager arm in per-round draft_logits clones.
14250                // Same raw-logit values either way. FILTERED q_j: stats from draft_stats
14251                // (eager pushes in-chain; the graph arms compute them post-replay from the
14252                // retained q with the same filter_stats program — bit-identical to the
14253                // in-graph stats that shaped the draw, keeping ONE accept path).
14254                let q_bufs: &[CudaSlice<f32>] = if dctx.graph_s.is_some() || dctx.chain_s.is_some()
14255                {
14256                    &dctx.q_slots
14257                } else {
14258                    &draft_logits
14259                };
14260                let mut n_acc = 0usize;
14261                for j in 0..k_round {
14262                    let (qmx, qth, qz) = draft_stats[j];
14263                    let idsd = e.htod_u32_v(&[draft_idx[j]])?;
14264                    let rowsd = e.htod_i32(&[0])?;
14265                    let thd = e.htod(&[qth])?;
14266                    let zd = e.htod(&[qz])?;
14267                    let _ = qmx;
14268                    let mut outd = e.zeros(1)?;
14269                    e.softmax_gather_filtered(
14270                        &q_bufs[j], d_vocab, &idsd, &rowsd, &thd, &zd, &mut outd, d_vocab, 1,
14271                        sp_temp,
14272                    )?;
14273                    let qj = e.dtoh(&outd)?[0];
14274                    let u = host_u01(sp_seed, uctr);
14275                    uctr += 1;
14276                    let accept = (u as f64) * (qj as f64) < pj[j] as f64;
14277                    // SKEY PROBE: q == 0 for the token the draft actually proposed is the
14278                    // exactness signature (see `skey_probe`). Impossible when the draft was
14279                    // drawn from the same filtered distribution the verify reconstructs here;
14280                    // `u * 0 < p` makes it an UNCONDITIONAL accept whenever p > 0.
14281                    if skey_probe() && qj == 0.0 {
14282                        eprintln!(
14283                            "[skey] EXACTNESS q=0 round={round} j={j} draft_tok={} \
14284                             draft_idx={} p={:e} u={u} accepted={} th_z={:?}",
14285                            draft[j], draft_idx[j], pj[j], accept as u8, draft_stats[j],
14286                        );
14287                    }
14288                    if accept {
14289                        n_acc += 1;
14290                    } else {
14291                        break;
14292                    }
14293                }
14294                let bonus = if n_acc == k_round {
14295                    // FULL ACCEPT: bonus ~ FILTERED softmax at the last verify column.
14296                    let col = base + k_round - 1;
14297                    let cb = col_buf.as_mut().unwrap();
14298                    e.copy_view_into(
14299                        cb,
14300                        0,
14301                        &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14302                        n_vocab,
14303                    )?;
14304                    if pen_on {
14305                        let h = pen_hist_d.as_ref().unwrap();
14306                        let nh = h.len();
14307                        e.penalize_logits(
14308                            cb,
14309                            h,
14310                            nh,
14311                            sp.penalty_repeat,
14312                            sp.penalty_freq,
14313                            sp.penalty_present,
14314                            n_vocab,
14315                        )?;
14316                    }
14317                    if perturb_buf.is_none() {
14318                        perturb_buf = Some(e.zeros(d_vocab.max(n_vocab))?);
14319                    }
14320                    // STATS MUST COME FROM THIS COLUMN (bug fix 2026-08-05, lane/sampler-
14321                    // truncation-fix; receipts research/sampfix-20260805/). The old code reused
14322                    // `col_stats.last()` here, which is ALWAYS the wrong row: the gathered set
14323                    // covers verify columns 0..=(base+k_round-2) (rows pushed as base+j-1), while
14324                    // the full-accept bonus samples column base+k_round-1 — exactly ONE PAST the
14325                    // last gathered column, in both base arms. `th` is a threshold in e-units of
14326                    // its OWN row's max, so feeding a neighbour's (row_max, th) into
14327                    // gumbel_perturb_filtered mis-scales every e0 = exp((x-row_max)/T). When the
14328                    // donor column's peak is higher by more than T*ln(1/th), EVERY id fails
14329                    // `e0 >= th`, the whole perturbed row becomes -3.4e38, and the 2-pass argmax
14330                    // falls through to its smallest-index tie-break => token id 0 ("!") spliced
14331                    // mid-word. Fragility is ordered by how large th is: min_p pins th = min_p
14332                    // (0.05 => trigger at delta > 2.4 at T=0.8, fires constantly), top_p's
14333                    // mass-boundary th is smaller, top_k's k-th-largest th smaller still — which
14334                    // is why the head-to-head matrix saw min_p and top_p corrupt while top_k-only
14335                    // stayed clean. The pure-temp default regime is immune (th == 0 masks nothing,
14336                    // and row_max is unused once nothing is masked), so this fix is a byte-level
14337                    // no-op for the untruncated serve default. One extra one-block filter_stats
14338                    // per full-accept round is the whole cost.
14339                    let (mx, th) = {
14340                        let rows0 = e.htod_i32(&[0])?;
14341                        let (mut th_d, mut z_d, mut mx_d) = (e.zeros(1)?, e.zeros(1)?, e.zeros(1)?);
14342                        let cb0 = col_buf.as_ref().unwrap();
14343                        e.filter_stats(
14344                            cb0, n_vocab, &rows0, &mut th_d, &mut z_d, &mut mx_d, n_vocab, 1,
14345                            sp_temp, sp.top_k, sp.top_p, sp.min_p,
14346                        )?;
14347                        (e.dtoh(&mx_d)?[0], e.dtoh(&th_d)?[0])
14348                    };
14349                    let pb = perturb_buf.as_mut().unwrap();
14350                    let cb2 = col_buf.as_ref().unwrap();
14351                    e.gumbel_perturb_filtered(cb2, pb, n_vocab, sp_seed, sctr, sp_temp, mx, th)?;
14352                    sctr += 1;
14353                    let td = e.argmax_token_device(pb, n_vocab)?;
14354                    e.dtoh_u32_one(&td)?
14355                } else {
14356                    // REJECT at n_acc: bonus ~ norm(max(0, softmax_T(p) - softmax_T(q))).
14357                    let cb = col_buf.as_mut().unwrap();
14358                    if n_acc > 0 || base == 1 {
14359                        let col = base + n_acc - 1;
14360                        e.copy_view_into(
14361                            cb,
14362                            0,
14363                            &tlogits_d.slice(col * n_vocab..(col + 1) * n_vocab),
14364                            n_vocab,
14365                        )?;
14366                    } else {
14367                        let lc = last_col_logits.as_ref().unwrap();
14368                        e.copy_into(cb, 0, lc, n_vocab)?;
14369                    }
14370                    if pen_on {
14371                        let h = pen_hist_d.as_ref().unwrap();
14372                        let nh = h.len();
14373                        e.penalize_logits(
14374                            cb,
14375                            h,
14376                            nh,
14377                            sp.penalty_repeat,
14378                            sp.penalty_freq,
14379                            sp.penalty_present,
14380                            n_vocab,
14381                        )?;
14382                    }
14383                    let cb2 = col_buf.as_ref().unwrap();
14384                    let sc = sctr;
14385                    sctr += 1;
14386                    // p-stats for the reject column: from col_stats when the col was gathered,
14387                    // else (j==0&&base==0) from last_col_stats.
14388                    let p_stats = if n_acc > 0 || base == 1 {
14389                        // col index within the gathered set == number of gathered cols before n_acc
14390                        let gi = if base == 1 { n_acc } else { n_acc - 1 };
14391                        col_stats.get(gi).copied().unwrap_or_else(|| {
14392                            (0.0, 0.0, 1.0) // unreachable: gathered cols always cover the reject slot
14393                        })
14394                    } else {
14395                        last_col_stats.expect("sampled: last_col_stats unset at reject")
14396                    };
14397                    let q_stats = draft_stats[n_acc];
14398                    if let Some(map) = &d2t_dev {
14399                        if q_full_buf.is_none() {
14400                            q_full_buf = Some(e.zeros(n_vocab)?);
14401                        }
14402                        let qf = q_full_buf.as_mut().unwrap();
14403                        e.scatter_trim_logits(&q_bufs[n_acc], map, qf, d_vocab, n_vocab)?;
14404                        let qf2 = q_full_buf.as_ref().unwrap();
14405                        e.residual_sample_filtered(
14406                            cb2,
14407                            Some(qf2),
14408                            n_vocab,
14409                            sp_temp,
14410                            sp_seed,
14411                            sc,
14412                            p_stats,
14413                            q_stats,
14414                            &mut sample_tok,
14415                        )?;
14416                    } else {
14417                        e.residual_sample_filtered(
14418                            cb2,
14419                            Some(&q_bufs[n_acc]),
14420                            n_vocab,
14421                            sp_temp,
14422                            sp_seed,
14423                            sc,
14424                            p_stats,
14425                            q_stats,
14426                            &mut sample_tok,
14427                        )?;
14428                    }
14429                    e.dtoh_u32(&sample_tok)?[0]
14430                };
14431                (
14432                    n_acc,
14433                    guard_vocab_token(
14434                        bonus,
14435                        n_vocab,
14436                        &format!("sampled verify bonus at round {round} pos={pos} n_acc={n_acc}"),
14437                    )?,
14438                )
14439            };
14440            // --- 3b. GRAMMAR TRUNCATION (constrained spec, 2026-08-03): the grammar is
14441            // an extra rejection rule AFTER the exactness verify (the batched-verify-twins
14442            // ordering). Walk the accepted drafts through the grammar in commit order; the
14443            // first illegal token truncates acceptance at its slot, and that slot's emission
14444            // is recomputed as the MASKED argmax of the target's own verify column — token-
14445            // identical to constrained plain greedy decode (an unmasked argmax that is
14446            // grammar-legal IS the masked argmax: masking only removes competitors). The
14447            // column D2H (~1MB) is paid only when a cut fires — the tight-grammar cost,
14448            // measured in acceptance numbers, never hidden.
14449            let (n_acc, bonus) = match constraint.as_deref_mut() {
14450                None => (n_acc, bonus),
14451                Some(c) => {
14452                    fn ce(e2: String) -> Box<dyn std::error::Error> {
14453                        format!("constraint: {e2}").into()
14454                    }
14455                    let mut na = n_acc;
14456                    let mut cut = false;
14457                    for (j, &d) in draft.iter().enumerate().take(n_acc) {
14458                        if c.is_allowed(d).map_err(ce)? {
14459                            c.consume(d).map_err(ce)?;
14460                        } else {
14461                            na = j;
14462                            cut = true;
14463                            dm_cut_tokens += n_acc - j;
14464                            break;
14465                        }
14466                    }
14467                    if cut {
14468                        dm_cuts += 1;
14469                    }
14470                    let mut bo = bonus;
14471                    if cut || !c.is_allowed(bo).map_err(ce)? {
14472                        let mut row = if na == 0 && base == 0 {
14473                            init_logits_host
14474                                .clone()
14475                                .ok_or("constraint: init logits missing (round-0 cut)")?
14476                        } else {
14477                            e.dtoh_view(
14478                                &tlogits_d.slice((base + na - 1) * n_vocab..(base + na) * n_vocab),
14479                            )?
14480                        };
14481                        c.mask_logits(&mut row).map_err(ce)?;
14482                        bo = argmax(&row) as u32;
14483                    }
14484                    c.consume(bo).map_err(ce)?;
14485                    (na, bo)
14486                }
14487            };
14488            let mut successor_valid = false;
14489            if let Some((q_proxy, expected_d2)) = rejected_probe {
14490                let v_n = n_acc == 1 && bonus == expected_d2;
14491                eprintln!(
14492                    "[opti-controller] shadow q={q_proxy:.6} admitted=false v_n={v_n} \
14493                     expected_d2={expected_d2} n_acc={n_acc} bonus={bonus}",
14494                );
14495            }
14496            if let Some(successor) = successor_attempt.as_ref() {
14497                successor_valid = n_acc == 1 && bonus == successor.verify_tokens[0];
14498                let generation = successor.generation;
14499                let q_proxy = successor.q_proxy;
14500                let expected_pending = successor.verify_tokens[0];
14501                let resolution_ms = successor.issued_at.elapsed().as_secs_f64() * 1e3;
14502                let fork = opti_fork
14503                    .as_mut()
14504                    .ok_or("optipipe successor resolution lost fork state")?;
14505                fork.finish_actual_reconcile(e, &mut *cache, &snap, n_acc, base, successor_valid)?;
14506                if successor_valid {
14507                    OPTI_FORK_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14508                } else {
14509                    OPTI_FORK_MISSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14510                    OPTI_RECONCILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14511                    OPTI_WASTED_DRAFT_TOKENS.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
14512                }
14513                let breaker_tripped = fork
14514                    .controller
14515                    .as_mut()
14516                    .expect("controller policy")
14517                    .resolve(successor_valid);
14518                if breaker_tripped {
14519                    OPTI_BREAKER_TRIPS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
14520                }
14521                eprintln!(
14522                    "[opti-controller] resolve generation={} hit={} q={q_proxy:.6} \
14523                     expected_pending={expected_pending} n_acc={n_acc} bonus={bonus} \
14524                     resolution_ms={resolution_ms:.3} reconcile={} breaker={}",
14525                    generation.id, successor_valid, !successor_valid, breaker_tripped,
14526                );
14527                if !successor_valid {
14528                    let mut successor = successor_attempt
14529                        .take()
14530                        .expect("controller successor disappeared on miss");
14531                    successor.settle();
14532                    fork.retire(generation)?;
14533                }
14534            }
14535            total_drafted += k_round;
14536            total_accepted += n_acc;
14537            if let Some(t) = sess_telem {
14538                // Greedy, rejection-sampling, and grammar truncation all converge here after
14539                // the accept decision is already on host. Fixed-size relaxed atomics only.
14540                t.record_round(k_round, n_acc);
14541            }
14542            if spec_stats {
14543                st_len_hist[k_round] += 1;
14544                for j in 0..k_round {
14545                    st_drafted[j] += 1;
14546                }
14547                for j in 0..n_acc {
14548                    st_accepted[j] += 1;
14549                }
14550                if n_acc == k_round {
14551                    st_full += 1;
14552                }
14553            }
14554
14555            if debug_spec {
14556                eprintln!(
14557                    "[R{round}] pos={pos} out_len={} last_tok={last_token} draft={draft:?} n_acc={n_acc} bonus={bonus} t_pred0={}",
14558                    out.len(),
14559                    // NOT `t_pred(0)`: `preds` is filled only under `if !sampled` above, so on a
14560                    // sampled request round >= 1 (base == 1) indexed an EMPTY vector and PANICKED
14561                    // the GPU worker thread — a debug flag that killed the exact regime you would
14562                    // set it to investigate. See `debug_t_pred0`.
14563                    debug_t_pred0(sampled, base, last_pred, &preds)
14564                );
14565            }
14566
14567            // --- 4. COMMIT: draft[0..n_acc] then bonus (n_acc + 1 tokens) ---
14568            let commit_started = std::time::Instant::now();
14569            // SESSION MODE: every accepted column is already in the CACHE — `out` must carry all
14570            // of them (overshoot past max_new included) or `committed` under-counts the cache rows
14571            // and the next turn's continuation seeds one token off (gate-caught 2026-07-05). The
14572            // single-shot path keeps the cap (its caller truncates + drops the cache anyway).
14573            for j in 0..n_acc {
14574                if !session_mode && out.len() >= max_new {
14575                    break;
14576                }
14577                out.push(draft[j]);
14578            }
14579            if pen_on {
14580                pen_hist.extend_from_slice(&draft[0..n_acc]);
14581                pen_hist.push(bonus);
14582            }
14583            let bonus_emitted = session_mode || out.len() < max_new;
14584            if bonus_emitted {
14585                out.push(bonus);
14586            }
14587            last_token = bonus;
14588
14589            // --- 5. ROLLBACK + advance (§C) ---
14590            if n_acc == k_round && !spec_replay {
14591                // FULL ACCEPT, BONUS FOLD: all verify columns (pending? + drafts) are committed in
14592                // cache; the NEW bonus stays PENDING for the next round's verify batch — NO extra
14593                // T=1 trunk pass. The next draft chain seeds from the MTP block's h_nextn at the
14594                // bonus position: one MTP-block pass (~1/33 trunk cost) replaces the trunk read.
14595                // last_pred is dead in the pending path (t_pred reads verify col 0).
14596                //
14597                // PERSISTENT DRAFT KV, full-accept fill: the chain covered last_token +
14598                // draft[0..k_round-2] as INPUTS (slots P..P'-2); draft[k_round-1] (slot P'-1) was
14599                // only ever an output, so its entry is MISSING. Fill it from vh_seed — its EXACT
14600                // trunk hidden (the last verify column). set_len first: a p-min break may have
14601                // left one extra chain append at that slot. Partial accepts need NO fill (the
14602                // chain already covered every accepted position; round-start set_len truncates).
14603                let mut vh_seed = e.zeros(n_embd)?;
14604                e.copy_view_into(
14605                    &mut vh_seed,
14606                    0,
14607                    &vx.slice((t_v - 1) * n_embd..t_v * n_embd),
14608                    n_embd,
14609                )?;
14610                if refresh {
14611                    // TRUE-HIDDEN REFRESH (2026-07-03, the HANDOVER-listed acceptance lever):
14612                    // overwrite ALL committed positions' scratch entries with K/V from their EXACT
14613                    // verify hiddens — the reference engine's mtp_update fills from true hiddens;
14614                    // the full stack (vx) is already resident from the verify. Replaces both the
14615                    // chain-approximate entries AND the old last-token-only fill. Acceptance-only
14616                    // (draft attention quality); exactness stays the verify's job.
14617                    scratch.set_len(e, pos)?;
14618                    // PREDECESSOR pairing: row i gets vx[i-1]; row 0 the carried fill_prev
14619                    // (hidden of the last committed row before this verify batch).
14620                    let mut vxs = e.zeros(t_v * n_embd)?;
14621                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14622                    if t_v > 1 {
14623                        e.copy_view_into(
14624                            &mut vxs,
14625                            n_embd,
14626                            &vx.slice(0..(t_v - 1) * n_embd),
14627                            (t_v - 1) * n_embd,
14628                        )?;
14629                    }
14630                    self.mtp_kv_fill_all(e, &verify_tokens, &vxs, pos, &mut *scratch, embd_dev)?;
14631                } else {
14632                    scratch.set_len(e, pos + base + k_round - 1)?;
14633                    // predecessor of the last draft = verify col t_v-2 (or fill_prev at t_v==1)
14634                    let mut hp = e.zeros(n_embd)?;
14635                    if t_v >= 2 {
14636                        e.copy_view_into(
14637                            &mut hp,
14638                            0,
14639                            &vx.slice((t_v - 2) * n_embd..(t_v - 1) * n_embd),
14640                            n_embd,
14641                        )?;
14642                    } else {
14643                        e.copy_into(&mut hp, 0, &fill_prev, n_embd)?;
14644                    }
14645                    self.mtp_kv_fill_all(
14646                        e,
14647                        &[draft[k_round - 1]],
14648                        &hp,
14649                        pos + base + k_round - 1,
14650                        &mut *scratch,
14651                        embd_dev,
14652                    )?;
14653                }
14654                // REFERENCE SEEDING: no pseudo pass — the next chain's step 0 IS the
14655                // reference's (id_last, h_prev) draft row; it appends the bonus's scratch
14656                // entry itself. Seed = TRUE hidden of the bonus's predecessor (last verify
14657                // col). Saves one MTP-block pass per round on top of the pairing fix.
14658                if !devacc_seeded {
14659                    e.copy_into(&mut h_seed_buf, 0, &vh_seed, n_embd)?;
14660                    e.copy_into(&mut fill_prev, 0, &vh_seed, n_embd)?;
14661                }
14662                pending = Some(bonus);
14663                if debug_spec {
14664                    eprintln!("  -> FULL ACCEPT (bonus pending, prev-h seed)");
14665                }
14666            } else if !spec_replay && base + n_acc >= 1 {
14667                // PARTIAL ACCEPT, REPLAY-FREE (2026-07-03 — the profiled #1 long-ctx spec cost):
14668                // the verify's first j = base+n_acc columns ARE the committed sequence, computed
14669                // bit-identically to eager (decode-exact contract) — so KEEP them: KV truncates to
14670                // pos+j, recurrent state rebuilds from the VerifyCkpt (same-kernel gdn prefix
14671                // re-run / pure state-clone restore), and the bonus stays PENDING exactly like the
14672                // full-accept path — the legacy duplicate trunk replay is gone. The next chain
14673                // seeds from the MTP pseudo-hidden of the bonus, whose seed = the TRUE verify
14674                // hidden of its predecessor (col j-1) — same one-hop pseudo structure as full
14675                // accept (never compounds: the next verify recomputes true hiddens for all
14676                // committed columns).
14677                let j = base + n_acc;
14678                // VERIFY-GRAPH SLAB COMMIT: when the captured trunk ran, the linear layers'
14679                // column stash was written into the graphs ctx's persistent slabs as in-graph
14680                // memcpy nodes, NOT into the per-column VerifyCkpt the cols arm reads — so the
14681                // commit must take the slab twin (same semantics, slab-addressed sources). The
14682                // ctx states which of the two this round produced via `round_slab`; trusting the
14683                // flag rather than the env keeps a round that fell back to the eager walk (a
14684                // capture that declined, a t the pool never captured) on the cols arm.
14685                let slab_commit = vg_guard
14686                    .as_ref()
14687                    .and_then(|g| g.as_ref())
14688                    .map(|g| g.round_slab)
14689                    .unwrap_or(false);
14690                if slab_commit {
14691                    self.dspark_commit_prefix_slab(
14692                        e,
14693                        &mut *cache,
14694                        &snap,
14695                        vg_guard
14696                            .as_ref()
14697                            .and_then(|g| g.as_ref())
14698                            .expect("slab_commit implies a graphs ctx"),
14699                        j,
14700                    )?;
14701                } else {
14702                    self.commit_verified_prefix(
14703                        e,
14704                        &mut *cache,
14705                        &snap,
14706                        ckpt.as_ref().unwrap(),
14707                        j,
14708                        devacc_seeded,
14709                        if devacc_seeded {
14710                            devacc_acc.as_ref().map(|a| (a, base, t_v))
14711                        } else {
14712                            None
14713                        },
14714                    )?;
14715                }
14716                let mut seed = e.zeros(n_embd)?;
14717                e.copy_view_into(
14718                    &mut seed,
14719                    0,
14720                    &vx.slice((j - 1) * n_embd..j * n_embd),
14721                    n_embd,
14722                )?;
14723                // Draft scratch: TRUE-HIDDEN REFRESH of the committed prefix (see the full-accept
14724                // branch); without it the chain entries stand and only the tail truncates. Either
14725                // way len ends at pos+j so the pseudo append lands at the bonus's slot pos+j
14726                // (persistent mode), rope pos+j+1 (chain convention).
14727                if refresh {
14728                    scratch.set_len(e, pos)?;
14729                    let mut vxs = e.zeros(j * n_embd)?;
14730                    e.copy_into(&mut vxs, 0, &fill_prev, n_embd)?;
14731                    if j > 1 {
14732                        e.copy_view_into(
14733                            &mut vxs,
14734                            n_embd,
14735                            &vx.slice(0..(j - 1) * n_embd),
14736                            (j - 1) * n_embd,
14737                        )?;
14738                    }
14739                    self.mtp_kv_fill_all(
14740                        e,
14741                        &verify_tokens[0..j],
14742                        &vxs,
14743                        pos,
14744                        &mut *scratch,
14745                        embd_dev,
14746                    )?;
14747                } else {
14748                    scratch.set_len(e, pos + j)?;
14749                }
14750                // REFERENCE SEEDING (see the full-accept branch): seed = TRUE hidden of the
14751                // bonus's predecessor (verify col j-1); no pseudo pass.
14752                if !devacc_seeded {
14753                    e.copy_into(&mut h_seed_buf, 0, &seed, n_embd)?;
14754                    e.copy_into(&mut fill_prev, 0, &seed, n_embd)?;
14755                }
14756                pending = Some(bonus);
14757                if debug_spec {
14758                    eprintln!("  -> PARTIAL(replay-free j={j}, bonus pending, prev-h seed)");
14759                }
14760            } else if !spec_replay {
14761                // ZERO ROUND FOLD (2026-07-10, verify-cost target #3): base+n_acc == 0 — a
14762                // pending-less round where nothing was accepted (PMIN0 zero-draft chains after a
14763                // replay/commit, or plain 0-accept rounds at round 0). The old path replayed
14764                // [bonus] through a FULL m=1 trunk+head forward (the 489us full-vocab head pass
14765                // measured at ~0.75/round on PMIN0 configs). Instead: restore the pre-round
14766                // snapshot and let the bonus ride the NEXT round's verify as col 0 — the existing
14767                // base=1 pending machinery, bit-identical by the decode-exact verify contract.
14768                // Seed: the bonus's predecessor is the last COMMITTED token, whose hidden
14769                // fill_prev already carries (same seeding as the 1-token-replay case it replaces).
14770                cache.rollback(e, &snap, 0)?;
14771                scratch.set_len(e, pos)?;
14772                e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14773                pending = Some(bonus);
14774                if debug_spec {
14775                    eprintln!("  -> ZERO-ROUND FOLD (bonus pending, fill_prev seed)");
14776                }
14777            } else {
14778                // PARTIAL ACCEPT, LEGACY REPLAY (seam MEMRA_SPEC_REPLAY=1 — or j==0: nothing of
14779                // this round survives, only possible before the first pending exists, ~round 0):
14780                // restore EVERYTHING to the pre-round snapshot (KV truncate to pos + recur
14781                // restore), then replay the committed prefix pending? ++ draft[0..n_acc] ++
14782                // [bonus] as ONE batched T forward — single weight read, bit-identical to greedy
14783                // (the verify-all-columns path is the same math). Commits the bonus with a TRUE
14784                // trunk hidden.
14785                cache.rollback(e, &snap, 0)?; // accept_len=0: KV len = pos, recur = snapshot
14786                let mut replay: Vec<u32> = Vec::with_capacity(base + n_acc + 1);
14787                if let Some(b) = pending.take() {
14788                    replay.push(b);
14789                }
14790                replay.extend_from_slice(&draft[0..n_acc]);
14791                replay.push(bonus);
14792                // Full-stack forward (decode_step_t_core = decode_step_t_h_emb_dev's body):
14793                // Predecessor pairing seeds from the PREDECESSOR row (col len-2) — the same-row path takes the
14794                // last col exactly as before (byte-identical to the old _h_emb_dev call).
14795                let (rl_d, rx) = if self.batched_serving_numeric_class() {
14796                    let mut logits = Vec::with_capacity(replay.len() * n_vocab);
14797                    let mut hidden = e.uninit(replay.len() * n_embd)?;
14798                    for (row, &token) in replay.iter().enumerate() {
14799                        let (row_logits, row_hidden) =
14800                            self.spec_target_step_h(e, token, &mut *cache)?;
14801                        logits.extend_from_slice(&row_logits);
14802                        e.dtod_copy_into(&row_hidden, &mut hidden, row * n_embd)?;
14803                    }
14804                    (e.htod(&logits)?, hidden)
14805                } else {
14806                    self.decode_step_t_core(e, &replay, pos, &mut *cache, embd_dev, None)?
14807                };
14808                // last_pred = argmax of the LAST column's logits (predicts the token after `bonus`)
14809                // — device argmax + one 4-byte read instead of the full-vocab column dtoh.
14810                e.argmax_token_device_col(&rl_d, replay.len() - 1, n_vocab, &mut preds_d, 0)?;
14811                last_pred = guard_vocab_token(
14812                    e.dtoh_u32(&preds_d)?[0],
14813                    n_vocab,
14814                    &format!("replay last_pred at round {round} pos={pos}"),
14815                )?;
14816                if sampled {
14817                    let lr0 = replay.len();
14818                    let lc = last_col_logits
14819                        .as_mut()
14820                        .expect("sampled: last_col_logits unset");
14821                    e.copy_view_into(
14822                        lc,
14823                        0,
14824                        &rl_d.slice((lr0 - 1) * n_vocab..lr0 * n_vocab),
14825                        n_vocab,
14826                    )?;
14827                }
14828                let lr = replay.len();
14829                if lr >= 2 {
14830                    e.copy_view_into(
14831                        &mut h_seed_buf,
14832                        0,
14833                        &rx.slice((lr - 2) * n_embd..(lr - 1) * n_embd),
14834                        n_embd,
14835                    )?;
14836                } else {
14837                    // 1-token replay (round-0 miss): the bonus's predecessor is the OLD
14838                    // last_token, whose own-row hidden fill_prev still holds.
14839                    e.copy_into(&mut h_seed_buf, 0, &fill_prev, n_embd)?;
14840                }
14841                // the bonus is COMMITTED here — it becomes the last committed row.
14842                let mut rh_last = e.zeros(n_embd)?;
14843                e.copy_view_into(
14844                    &mut rh_last,
14845                    0,
14846                    &rx.slice((lr - 1) * n_embd..lr * n_embd),
14847                    n_embd,
14848                )?;
14849                e.copy_into(&mut fill_prev, 0, &rh_last, n_embd)?;
14850                if debug_spec {
14851                    eprintln!("  -> PARTIAL(replay={replay:?}), next_pred={last_pred}");
14852                }
14853            }
14854            if devacc_seeded {
14855                // stage (b) epilogue: fill_prev takes the gathered seed AFTER the refresh fills
14856                // consumed the old value (both slots carry the same value in every non-replay arm).
14857                e.copy_into(&mut fill_prev, 0, &h_seed_buf, n_embd)?;
14858            }
14859            if successor_valid {
14860                let optimistic_scratch_len = successor_attempt
14861                    .as_ref()
14862                    .expect("valid controller successor disappeared")
14863                    .scratch_len;
14864                // The normal current-round commit refreshed/truncated the logical scratch tail.
14865                // Its optimistic successor row was already written physically, so restoring only
14866                // the retained logical length makes that row live for the carried round.
14867                scratch.set_len(e, optimistic_scratch_len)?;
14868            }
14869            if let Some(current) = current_opti.take() {
14870                opti_fork
14871                    .as_mut()
14872                    .ok_or("optipipe current retirement lost fork state")?
14873                    .retire(current.generation)?;
14874            }
14875            if successor_valid {
14876                let successor = successor_attempt
14877                    .take()
14878                    .expect("valid controller successor disappeared before promotion");
14879                let generation = successor.generation;
14880                opti_fork
14881                    .as_mut()
14882                    .ok_or("optipipe successor promotion lost fork state")?
14883                    .promote_successor_snapshot(&mut snap, generation);
14884                carried_opti = Some(successor);
14885            }
14886            if anatomy_on {
14887                // Commit/rollback is normally asynchronous on the primary/head stream. Bound it
14888                // only for this diagnostic so it does not disappear into the following draft's
14889                // first token readback.
14890                e.stream().synchronize()?;
14891                ph_commit += commit_started.elapsed().as_secs_f64();
14892            }
14893            // adaptive-K update (host math, zero syncs): next round drafts accepted-run + 1,
14894            // clamped to [floor(pos), k_cap]. cache.pos is post-rollback here (the round's
14895            // final position — the floor's position key reads the committed depth). Burst
14896            // rounds (`continue` above) draft the captured fixed depth and skip this, exactly
14897            // like gemma's burst arm.
14898            if adapt {
14899                let fl_now = floor_at(cache.pos);
14900                kc = (n_acc + 1).clamp(fl_now.min(k_cap), k_cap);
14901            }
14902            ph_mark(&mut ph_rest, phase_on);
14903            if let Some(p) = pipe {
14904                p.accept_end(round);
14905            }
14906            drop(pipe_accept);
14907            if let Some(t0) = round_t0 {
14908                let ms = t0.elapsed().as_secs_f64() * 1e3;
14909                ROUND_MS.fetch_add((ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed);
14910                let n = ROUND_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
14911                if n % 32 == 0 {
14912                    eprintln!(
14913                        "[spec-round] rounds={n} avg round wall={:.2} ms (emitted={} drafted so far)",
14914                        ROUND_MS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e3 / n as f64,
14915                        out.len()
14916                    );
14917                }
14918            }
14919            round += 1;
14920            // sse-cadence: this round's accepted drafts + bonus are committed (out is
14921            // append-only past step 4) — flush at round cadence.
14922            keep_going = flush_commit(&mut on_commit, &out, &mut flushed);
14923        }
14924        if let Some(mut ticket) = carried_opti.take() {
14925            opti_fork
14926                .as_mut()
14927                .ok_or("optipipe tail drain lost fork state")?
14928                .cancel_controller_ticket(e, &mut *cache, &mut *scratch, &snap, &mut ticket)?;
14929        }
14930        // sse-cadence: nothing below appends to `out`; flush any remainder (defensive).
14931        // (verdict ignored — the burst is over either way; the session tail runs unchanged.)
14932        let _ = flush_commit(&mut on_commit, &out, &mut flushed);
14933
14934        if spec_stats {
14935            let per_slot: Vec<String> = (0..k)
14936                .map(|j| {
14937                    if st_drafted[j] > 0 {
14938                        format!(
14939                            "{}/{}={:.3}",
14940                            st_accepted[j],
14941                            st_drafted[j],
14942                            st_accepted[j] as f64 / st_drafted[j] as f64
14943                        )
14944                    } else {
14945                        "0/0".into()
14946                    }
14947                })
14948                .collect();
14949            let acc = if total_drafted > 0 {
14950                total_accepted as f64 / total_drafted as f64
14951            } else {
14952                0.0
14953            };
14954            eprintln!(
14955                "[spec-stats] rounds={round} full_accept={st_full} len_hist={st_len_hist:?} \
14956                       per_slot=[{}] total={total_accepted}/{total_drafted}={acc:.3} \
14957                       tok_per_round={:.3}",
14958                per_slot.join(" "),
14959                (total_accepted + round) as f64 / round.max(1) as f64
14960            );
14961        }
14962        if constraint.is_some() {
14963            eprintln!(
14964                "[draft-mask] mask_rounds={dm_rounds} clone_total={:.3}ms \
14965                 clone_per_round={:.4}ms gram_cuts={dm_cuts}/{round} cut_tokens={dm_cut_tokens}",
14966                dm_clone_ns as f64 / 1e6,
14967                dm_clone_ns as f64 / 1e6 / dm_rounds.max(1) as f64
14968            );
14969        }
14970        if phase_on {
14971            let tot = ph_draft + ph_verify + ph_wait + ph_rest;
14972            eprintln!(
14973                "[spec-phase] draft={:.1}ms ({:.1}%) verify-issue={:.1}ms ({:.1}%) verify-wait={:.1}ms ({:.1}%) commit-host={:.1}ms ({:.1}%) rounds={round}",
14974                ph_draft * 1e3,
14975                ph_draft / tot * 100.0,
14976                ph_verify * 1e3,
14977                ph_verify / tot * 100.0,
14978                ph_wait * 1e3,
14979                ph_wait / tot * 100.0,
14980                ph_rest * 1e3,
14981                ph_rest / tot * 100.0
14982            );
14983        }
14984        if anatomy_on {
14985            let rounds_f = round.max(1) as f64;
14986            let other = (ph_rest - ph_commit).max(0.0);
14987            eprintln!(
14988                "[spec-anatomy] per-round draft={:.3}ms pp-verify={:.3}ms \
14989                 verify-accept={:.3}ms commit-rollback={:.3}ms other={:.3}ms rounds={round}",
14990                ph_draft * 1e3 / rounds_f,
14991                ph_verify * 1e3 / rounds_f,
14992                ph_wait * 1e3 / rounds_f,
14993                ph_commit * 1e3 / rounds_f,
14994                other * 1e3 / rounds_f,
14995            );
14996        }
14997        let _pipe_tail = pipe.map(|p| p.primary());
14998        // SESSION TAIL: leave the session in the exact invariant the next turn's suffix prime
14999        // expects — every row in `committed` has trunk KV/recur state AND an exact draft-KV row.
15000        // Park the draft-graph ctx back on the session (the serve-burst fixed-cost fix): the next
15001        // burst replays instead of recapturing. Error paths (`?` above) drop it — recaptured then.
15002        if let Some(slot) = sess_draft_slot.take() {
15003            *slot = Some(dctx);
15004        }
15005        let t_rounds = t_ent.elapsed();
15006        if let Some((committed, last_h, next_pred_slot, sctr_slot, uctr_slot)) = sess_tail.take() {
15007            // NEXT BURST'S BOUNDARY TOKEN (lane/sampled-spec-quality, Item 1). Greedy stashes
15008            // the argmax `last_pred` exactly as before (byte contract). SAMPLED draws the token
15009            // HERE, where the sampler, the session Philox counters and the penalty window are
15010            // all live and the boundary logits row still exists — that is the "make the state
15011            // available" half of the fix; the consuming burst then just emits it. `sctr` is
15012            // written to the session BELOW the draws so the advance is never lost.
15013            *next_pred_slot = Some(last_pred);
15014            let sample_boundary = sampled && constraint.is_none() && spec_sampled_boundary_on();
15015            let mut stashed_pending = false;
15016            if let Some(b) = pending.take() {
15017                if !sampled {
15018                    // PENDING-CARRY (2026-08-01): stash the bonus on the session instead of
15019                    // committing it with a solo T=1 pass — the next empty-suffix greedy burst
15020                    // consumes it as round-0 verify col 0 (a plain round edge; the old tail
15021                    // commit + next burst's init feed were 11.6+11.5ms solo trunk passes per
15022                    // burst on H100 q27, [spec-setup] trace). b stays in `out` (emitted) but
15023                    // OUT of `committed` (cache rows == committed); the consuming call
15024                    // prepends it once its verify commits the row. next_pred is unknowable
15025                    // without the commit pass — None; callers gate on pending_tok too.
15026                    debug_assert_eq!(out.last(), Some(&b), "pending must be the last emitted");
15027                    if let Some(slot) = sess_pending_slot.take() {
15028                        *slot = Some(b);
15029                    }
15030                    *next_pred_slot = None;
15031                    // fill_prev = hidden of the last COMMITTED row (b's predecessor) — the
15032                    // exact chain-seed/fill anchor the consuming burst (or a flush) needs.
15033                    *last_h = Some(e.clone_dtod(&fill_prev)?);
15034                    stashed_pending = true;
15035                } else {
15036                    // SAMPLED tail (unchanged): commit the bonus (one T=1 pass) + draft fill —
15037                    // the sampled round-0 accept needs this pass's logits (last_col_logits).
15038                    let pos_b = cache.pos;
15039                    scratch.set_len(e, pos_b)?;
15040                    let (lg_b, hb) = self.spec_target_step_h(e, b, &mut *cache)?;
15041                    // after a FULL-accept exit `last_pred` is STALE (it predicted the bonus
15042                    // itself — the prediction AFTER the bonus never materialized; it would have
15043                    // been the next round's verify col 0). The commit's logits ARE that
15044                    // prediction — so they are also the row the next burst's boundary token
15045                    // comes off, and (lane/sampled-spec-quality) it is DRAWN from them here.
15046                    *next_pred_slot = Some(if sample_boundary {
15047                        sample_boundary_token(
15048                            e,
15049                            &lg_b,
15050                            &sp,
15051                            &pen_hist,
15052                            &mut sctr,
15053                            "burst-tail-commit",
15054                        )?
15055                    } else {
15056                        argmax(&lg_b) as u32
15057                    });
15058                    self.mtp_kv_fill_all(e, &[b], &fill_prev, pos_b, &mut *scratch, embd_dev)?;
15059                    *last_h = Some(hb);
15060                }
15061            } else {
15062                // fill_prev tracks the hidden of the last COMMITTED row throughout the loop.
15063                *last_h = Some(e.clone_dtod(&fill_prev)?);
15064                if sample_boundary {
15065                    // No pending to commit, so the boundary row is the one `last_pred` was
15066                    // argmaxed from and the sampled path keeps it on device: the init feed's
15067                    // logits when the burst ran zero rounds, else the legacy-replay path's
15068                    // last verify column (both predict the token AFTER the last committed
15069                    // row). It is retained precisely because round 0's accept test needs it,
15070                    // so the draw costs no extra D2H of the [n_vocab] row.
15071                    match last_col_logits.as_ref() {
15072                        Some(lc) => {
15073                            *next_pred_slot = Some(sample_boundary_token_dev(
15074                                e,
15075                                lc,
15076                                n_vocab,
15077                                &sp,
15078                                &pen_hist,
15079                                &mut sctr,
15080                                "burst-tail-nopending",
15081                            )?);
15082                        }
15083                        // NAME THE FALLBACK (house standard): unreachable today — a sampled
15084                        // burst always feeds or replays, so the row exists — but if it ever
15085                        // is, the stream takes a greedy token and SAYS so rather than
15086                        // silently regressing to the pre-lane behaviour.
15087                        None => eprintln!(
15088                            "[spec-boundary] sampled tail kept the ARGMAX boundary token \
15089                             (reason: no retained boundary logits row)"
15090                        ),
15091                    }
15092                }
15093            }
15094            *sctr_slot = sctr;
15095            *uctr_slot = uctr;
15096            committed.extend_from_slice(prompt);
15097            if let Some(cb) = carried_pending {
15098                // the consumed carry's cache row landed in round 0's verify (every pending
15099                // round commits col 0) — it joins `committed` here, in sequence order.
15100                committed.push(cb);
15101            }
15102            if stashed_pending {
15103                // ZERO-EMIT BURST (2026-08-06 c=8 serve panic, pre-existing since b4aea184):
15104                // `out.len() - 1` underflowed on an EMPTY `out` — "range end index
15105                // 18446744073709551615 out of range for slice of length 0", killing the
15106                // memra-gpu-worker and failing 31 of 32 concurrent requests with "worker closed
15107                // stream". Reachable because `pending` starts as `carried_pending` (a bonus
15108                // stashed by the PREVIOUS burst) while `out` starts empty, and the carry is
15109                // deliberately NOT pushed to `out` (line ~3239: the burst that emitted it already
15110                // did). So a burst that stashes a pending without emitting anything of its own —
15111                // the round loop exits before a push, e.g. the ring drain's `out.len() < max_new`
15112                // guard skipping every token under a tight budget — arrives here with
15113                // out.len() == 0 and stashed_pending == true.
15114                //
15115                // The invariant is unchanged: `committed` gets every emitted token EXCEPT the
15116                // stashed bonus. With nothing emitted, that is nothing — and the carry pushed
15117                // just above is already accounted. Saturating, not a min/assert: an empty `out`
15118                // here is a legitimate burst shape, not a corrupt state.
15119                let emitted = out.len().saturating_sub(1);
15120                committed.extend_from_slice(&out[..emitted]);
15121            } else {
15122                committed.extend_from_slice(&out); // FULL out incl. overshoot — all committed
15123            }
15124            debug_assert_eq!(
15125                cache.pos,
15126                committed.len(),
15127                "session invariant: cache rows == committed tokens"
15128            );
15129            if setup_trace {
15130                e.stream().synchronize()?; // bound the async tail fill in the trace
15131                let t_tail = t_ent.elapsed();
15132                eprintln!(
15133                    "[spec-setup] init={:.2}ms cap={:.2}ms fill={:.2}ms rounds={:.2}ms tail={:.2}ms total={:.2}ms out={} cont={}",
15134                    t_init.as_secs_f64() * 1e3,
15135                    (t_cap - t_init).as_secs_f64() * 1e3,
15136                    (t_fill - t_cap).as_secs_f64() * 1e3,
15137                    (t_rounds - t_fill).as_secs_f64() * 1e3,
15138                    (t_tail - t_rounds).as_secs_f64() * 1e3,
15139                    t_tail.as_secs_f64() * 1e3,
15140                    out.len(),
15141                    continuation
15142                );
15143            }
15144            return Ok((out, total_drafted, total_accepted));
15145        }
15146        out.truncate(max_new);
15147        Ok((out, total_drafted, total_accepted))
15148    }
15149
15150    /// Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape;
15151    /// only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs
15152    /// with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.
15153    pub fn extract_dspark_anchors(
15154        &self,
15155        e: &Engine,
15156        tokens: &[u32],
15157        anchor_positions: &[usize],
15158        gamma: usize,
15159        top_k: usize,
15160        chunk: usize,
15161        temperature: f32,
15162    ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn std::error::Error>> {
15163        if tokens.len() < gamma + 2 || gamma == 0 || chunk < 2 {
15164            return Err("DSpark extraction token tape/gamma/chunk is invalid".into());
15165        }
15166        if anchor_positions.windows(2).any(|pair| pair[0] >= pair[1]) {
15167            return Err("DSpark anchor positions must be sorted and unique".into());
15168        }
15169        for &position in anchor_positions {
15170            if position == 0 || position + gamma >= tokens.len() {
15171                return Err(format!(
15172                    "DSpark anchor {position} has no predecessor or cannot cover gamma={gamma} in {} tokens",
15173                    tokens.len()
15174                )
15175                .into());
15176            }
15177        }
15178
15179        let n_vocab = self.output.out_features();
15180        let n_embd = self.cfg.n_embd as usize;
15181        let mut cache =
15182            crate::pp::new_cache_planned(e, &self.cfg, &self.plan, tokens.len() + gamma + 8)?;
15183        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15184        let embd_gpu = if spec_host_embd() {
15185            None
15186        } else {
15187            Some(
15188                self.embd_gpu
15189                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15190            )
15191        };
15192        let embd_dev = embd_gpu.map(|gpu| (gpu, embd_qt, embd_rb));
15193
15194        struct PendingRecord {
15195            position: usize,
15196            hidden: Option<Vec<f32>>,
15197            tokens: Vec<u32>,
15198            target_top_ids: Vec<Option<Vec<u32>>>,
15199            target_top_logits: Vec<Option<Vec<f32>>>,
15200            target_top_probs: Vec<Option<Vec<f32>>>,
15201            target_tail_probs: Vec<Option<f32>>,
15202        }
15203
15204        let mut pending: Vec<PendingRecord> = anchor_positions
15205            .iter()
15206            .map(|&position| PendingRecord {
15207                position,
15208                hidden: None,
15209                tokens: tokens[position..=position + gamma].to_vec(),
15210                target_top_ids: vec![None; gamma],
15211                target_top_logits: vec![None; gamma],
15212                target_top_probs: vec![None; gamma],
15213                target_tail_probs: vec![None; gamma],
15214            })
15215            .collect();
15216
15217        let mut start = 0usize;
15218        while start < tokens.len() {
15219            let end = (start + chunk).min(tokens.len());
15220            let chunk_tokens = &tokens[start..end];
15221            let (target_logits, hidden_rows) =
15222                self.decode_step_t_core(e, chunk_tokens, start, &mut cache, embd_dev, None)?;
15223            for record in &mut pending {
15224                let hidden_position = record.position - 1;
15225                if hidden_position >= start && hidden_position < end {
15226                    let local = hidden_position - start;
15227                    record.hidden = Some(
15228                        e.dtoh_view(&hidden_rows.slice(local * n_embd..(local + 1) * n_embd))?,
15229                    );
15230                }
15231                for slot in 0..gamma {
15232                    let target_row = record.position + slot;
15233                    if target_row < start || target_row >= end {
15234                        continue;
15235                    }
15236                    let local = target_row - start;
15237                    let logits =
15238                        e.dtoh_view(&target_logits.slice(local * n_vocab..(local + 1) * n_vocab))?;
15239                    let (ids, top_logits, probs, tail) =
15240                        dspark_sparse_softmax_topk(&logits, top_k, temperature)?;
15241                    record.target_top_ids[slot] = Some(ids);
15242                    record.target_top_logits[slot] = Some(top_logits);
15243                    record.target_top_probs[slot] = Some(probs);
15244                    record.target_tail_probs[slot] = Some(tail);
15245                }
15246            }
15247            start = end;
15248        }
15249
15250        pending
15251            .into_iter()
15252            .map(|record| {
15253                let hidden = record
15254                    .hidden
15255                    .ok_or_else(|| format!("missing DSpark hidden at {}", record.position))?;
15256                let target_top_ids =
15257                    flatten_dspark_rows(record.target_top_ids, record.position, "target ids")?;
15258                let target_top_logits = flatten_dspark_rows(
15259                    record.target_top_logits,
15260                    record.position,
15261                    "target logits",
15262                )?;
15263                let target_top_probs =
15264                    flatten_dspark_rows(record.target_top_probs, record.position, "target probs")?;
15265                let target_tail_probs = record
15266                    .target_tail_probs
15267                    .into_iter()
15268                    .enumerate()
15269                    .map(|(slot, value)| {
15270                        value.ok_or_else(|| {
15271                            format!("missing DSpark tail at {} slot {slot}", record.position)
15272                        })
15273                    })
15274                    .collect::<Result<Vec<_>, _>>()?;
15275                Ok(DsparkAnchorRecord {
15276                    position: record.position,
15277                    hidden,
15278                    tokens: record.tokens,
15279                    target_top_ids,
15280                    target_top_logits,
15281                    target_top_probs,
15282                    target_tail_probs,
15283                })
15284            })
15285            .collect()
15286    }
15287
15288    /// TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token
15289    /// sequence and, at sampled positions, compare the MTP head's K-token draft chain against
15290    /// the trunk's own teacher-forced greedy predictions. Nothing is generated — the context is
15291    /// the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance
15292    /// and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the
15293    /// quant-induced head/hidden-state mismatch from text drift.
15294    ///
15295    /// Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode):
15296    ///   draft_j  = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact
15297    ///              eager spec-decode chain (same mtp_head_forward_dev, same rope positions).
15298    ///   target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits
15299    ///              at forced context tokens[0..p+j]). For j==0 this equals live spec
15300    ///              acceptance; for j>=1 live verify would condition on the drafts, here it
15301    ///              conditions on the corpus — deterministic and arm-comparable by design.
15302    ///
15303    /// Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p),
15304    /// plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1)
15305    /// so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).
15306    ///
15307    /// `hdump`: when Some, every position's pre-output_norm trunk hidden (the exact rows the
15308    /// draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] —
15309    /// the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk
15310    /// hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy
15311    /// agreement vs this path — not usable as a training-data source).
15312    pub fn replay_acceptance(
15313        &self,
15314        e: &Engine,
15315        tokens: &[u32],
15316        k: usize,
15317        stride: usize,
15318        chunk: usize,
15319        mut hdump: Option<&mut std::fs::File>,
15320    ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn std::error::Error>> {
15321        assert!(k >= 1 && stride >= 1 && chunk >= 2);
15322        let mtp = self
15323            .mtp
15324            .as_ref()
15325            .expect("replay_acceptance requires an MTP head");
15326        let n_vocab = self.output.out_features();
15327        let d_vocab = mtp
15328            .shared_head_head
15329            .as_ref()
15330            .unwrap_or(&self.output)
15331            .out_features();
15332        let n_embd = self.cfg.n_embd as usize;
15333        let t_total = tokens.len();
15334        assert!(t_total >= 8, "corpus too short ({t_total} tokens)");
15335        // STAGE-OWNED KV (lane/pp2-spec 2026-08-06) — see `new_session`. Door shut = `Cache::new`.
15336        let mut cache = crate::pp::new_cache_planned(e, &self.cfg, &self.plan, t_total + k + 8)?;
15337        let mut scratch = self.new_mtp_scratch(e, t_total + k + 8)?;
15338        let (embd_qt, embd_rb) = self.embd.qt_and_row_bytes(n_embd);
15339        let embd_gpu = if spec_host_embd() {
15340            None
15341        } else {
15342            Some(
15343                self.embd_gpu
15344                    .get_or_init(|| e.upload_u8(&self.embd.raw).expect("embed table upload")),
15345            )
15346        };
15347        let embd_dev = embd_gpu.map(|g| (g, embd_qt, embd_rb));
15348
15349        // bg[i] = the trunk's greedy pick for position i under the forced context (i >= 1).
15350        let mut bg: Vec<u32> = vec![0; t_total + 1];
15351        let mut rows: Vec<(usize, Vec<u32>, Vec<u32>)> = Vec::new();
15352        let mut prev_last_h = e.zeros(n_embd)?; // predecessor hidden entering the chunk
15353        let mut seed_buf = e.zeros(n_embd)?;
15354        let mut preds_d = e.alloc_u32_zeroed(chunk)?;
15355        let nll_on = std::env::var("MEMRA_REPLAY_NLL").as_deref() == Ok("1");
15356        let (mut nll_sum, mut nll_cnt) = (0f64, 0u64);
15357        let mut s = 0usize;
15358        while s < t_total {
15359            let cend = (s + chunk).min(t_total);
15360            let tc = cend - s;
15361            let ch = &tokens[s..cend];
15362            // 1. forced trunk pass — verify path (decode-exact contract): all-column logits +
15363            //    the chunk's true hiddens.
15364            let (tl_d, vx) = self.decode_step_t_core(e, ch, s, &mut cache, embd_dev, None)?;
15365            for j in 0..tc {
15366                e.argmax_token_device_col(&tl_d, j, n_vocab, &mut preds_d, j)?;
15367            }
15368            let preds = e.dtoh_u32(&preds_d)?;
15369            for j in 0..tc {
15370                bg[s + j + 1] = preds[j];
15371            }
15372            // MEMRA_REPLAY_NLL=1: teacher-forced NLL/perplexity over the same forced pass — the
15373            // checkpoint-quality metric (position j's logits score the GOLD next token).
15374            if nll_on {
15375                let jmax = if cend < t_total { tc } else { tc - 1 }; // last pos has no gold next
15376                if jmax > 0 {
15377                    let ids: Vec<u32> = (0..jmax).map(|j| tokens[s + j + 1]).collect();
15378                    let rows: Vec<i32> = (0..jmax as i32).collect();
15379                    let idsd = e.htod_u32_v(&ids)?;
15380                    let rowsd = e.htod_i32(&rows)?;
15381                    let mut outd = e.zeros(jmax)?;
15382                    e.softmax_gather(&tl_d, n_vocab, &idsd, &rowsd, &mut outd, n_vocab, jmax, 1.0)?;
15383                    for pr in e.dtoh(&outd)? {
15384                        nll_sum += -((pr.max(1e-30)) as f64).ln();
15385                        nll_cnt += 1;
15386                    }
15387                }
15388            }
15389            if let Some(f) = hdump.as_deref_mut() {
15390                use std::io::Write;
15391                let host: Vec<f32> = e.dtoh(&vx)?;
15392                // bf16 round-to-nearest-even — f32 doubled the disk bill at bulk
15393                // extraction scale (20M tokens x 4096 = 320GB f32 vs 160GB bf16).
15394                let mut bytes = Vec::with_capacity(tc * n_embd * 2);
15395                for v in &host[..tc * n_embd] {
15396                    let b = v.to_bits();
15397                    let r = b.wrapping_add(0x7FFF + ((b >> 16) & 1));
15398                    bytes.extend_from_slice(&((r >> 16) as u16).to_le_bytes());
15399                }
15400                f.write_all(&bytes)?;
15401            }
15402            // CHAINLESS extraction (stride > corpus, the bulk-hdump mode): no chunk ever
15403            // drafts, so the draft-KV fills are pure waste — skip them (2 MTP-block passes
15404            // per token saved; the forced trunk pass + hdump is all the mode needs).
15405            let chainless = stride > t_total;
15406            if chainless {
15407                e.copy_view_into(
15408                    &mut prev_last_h,
15409                    0,
15410                    &vx.slice((tc - 1) * n_embd..tc * n_embd),
15411                    n_embd,
15412                )?;
15413                s = cend;
15414                continue;
15415            }
15416            // 2. TRUE predecessor-paired draft-KV fill for the chunk (row i carries h_{i-1};
15417            //    row s reads the previous chunk's last true hidden, zeros at corpus start).
15418            let mut vxs = e.zeros(tc * n_embd)?;
15419            e.copy_into(&mut vxs, 0, &prev_last_h, n_embd)?;
15420            if tc > 1 {
15421                e.copy_view_into(
15422                    &mut vxs,
15423                    n_embd,
15424                    &vx.slice(0..(tc - 1) * n_embd),
15425                    (tc - 1) * n_embd,
15426                )?;
15427            }
15428            scratch.set_len(e, s)?;
15429            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15430            // 3. draft chains at sampled positions, DESCENDING: a chain reads only slots
15431            //    [0..p) (true fills) and appends at >= p; the next (smaller-p) chain's set_len
15432            //    truncates those approximate appends before they can ever be read.
15433            let ps: Vec<usize> = (s..cend)
15434                .filter(|p| *p >= 1 && *p % stride == 0 && *p + k <= t_total)
15435                .collect();
15436            for &p in ps.iter().rev() {
15437                scratch.set_len(e, p)?;
15438                if p == s {
15439                    e.copy_into(&mut seed_buf, 0, &prev_last_h, n_embd)?;
15440                } else {
15441                    e.copy_view_into(
15442                        &mut seed_buf,
15443                        0,
15444                        &vx.slice((p - 1 - s) * n_embd..(p - s) * n_embd),
15445                        n_embd,
15446                    )?;
15447                }
15448                let mut e_tok = tokens[p];
15449                let mut d_seed = e.clone_dtod(&seed_buf)?;
15450                let chain_heads = !self.mtp_extra.is_empty();
15451                let mut chain_tokens = if chain_heads {
15452                    vec![tokens[p]]
15453                } else {
15454                    Vec::new()
15455                };
15456                let mut chain_seeds = if chain_heads {
15457                    vec![e.clone_dtod(&seed_buf)?]
15458                } else {
15459                    Vec::new()
15460                };
15461                let mut drafts: Vec<u32> = Vec::with_capacity(k);
15462                for j in 0..k {
15463                    let (dl_d, h_nextn) = if chain_heads {
15464                        self.mtp_chain_forward_dev(
15465                            e,
15466                            &chain_tokens,
15467                            &chain_seeds,
15468                            &mut scratch,
15469                            p,
15470                            embd_dev,
15471                            None,
15472                        )?
15473                    } else {
15474                        self.mtp_head_forward_dev(
15475                            e,
15476                            mtp,
15477                            e_tok,
15478                            &d_seed,
15479                            &mut scratch,
15480                            p + 1 + j,
15481                            embd_dev,
15482                            None,
15483                        )?
15484                    };
15485                    let tok_d = e.argmax_token_device(&dl_d, d_vocab)?;
15486                    let idx = e.dtoh_u32_one(&tok_d)?;
15487                    let d = match &mtp.d2t {
15488                        Some(map) => map[idx as usize],
15489                        None => idx,
15490                    };
15491                    drafts.push(d);
15492                    if chain_heads {
15493                        chain_tokens.push(d);
15494                        chain_seeds.push(h_nextn);
15495                    } else {
15496                        e_tok = d;
15497                        d_seed = h_nextn;
15498                    }
15499                }
15500                // targets may live in a LATER chunk's bg — resolved after the walk.
15501                rows.push((p, drafts, Vec::new()));
15502            }
15503            // 4. restore TRUE entries for the whole chunk (the next chunk's chains and fills
15504            //    expect scratch.len == cend with exact rows).
15505            scratch.set_len(e, s)?;
15506            self.mtp_kv_fill_all(e, ch, &vxs, s, &mut scratch, embd_dev)?;
15507            e.copy_view_into(
15508                &mut prev_last_h,
15509                0,
15510                &vx.slice((tc - 1) * n_embd..tc * n_embd),
15511                n_embd,
15512            )?;
15513            s = cend;
15514        }
15515        for (p, drafts, targets) in rows.iter_mut() {
15516            for j in 0..drafts.len() {
15517                targets.push(bg[*p + 1 + j]);
15518            }
15519        }
15520        rows.sort_by_key(|r| r.0);
15521        if nll_cnt > 0 {
15522            let mean = nll_sum / nll_cnt as f64;
15523            println!(
15524                "[replay-nll] tokens={nll_cnt} nll/token={mean:.5} ppl={:.4}",
15525                mean.exp()
15526            );
15527        }
15528        Ok((rows, bg))
15529    }
15530}
15531
15532#[cfg(test)]
15533mod vg_debt_tests {
15534    use super::dspark_vg_debt_projection;
15535
15536    /// TOOTH for the verify-graph admission accounting: the pool's projected remaining
15537    /// growth must be charged (pre-fix, admission charged 0 for a pool measured at
15538    /// 8,852 MiB), the projection must price the MARGINAL cost of one more key rather than
15539    /// extrapolating the pool's one-time shared allocation, and the doors that make growth
15540    /// impossible must zero the debt.
15541    #[test]
15542    fn vg_debt_projects_remaining_growth_and_respects_the_freeze_valves() {
15543        const MIB: usize = 1 << 20;
15544        let d = dspark_vg_debt_projection;
15545        // cold pool: nothing observed, one capture fits inside SPEC_SHRINK_RESERVE.
15546        assert_eq!(d(0, 256, 0, None), 0);
15547        // freeze valve MEMRA_DSPARK_VG_MAX=0: the pool cannot grow.
15548        assert_eq!(d(10, 0, 500 * MIB, None), 0);
15549        // saturated pool: at/past the cap the pool FREEZES, nothing left to reserve.
15550        assert_eq!(d(256, 256, 8852 * MIB, None), 0);
15551        assert_eq!(d(300, 256, 8852 * MIB, None), 0);
15552
15553        // BOOTSTRAP (one observation, growth unmeasurable): at most one more pool's worth.
15554        // The pre-fix mean rule extrapolated 255x here — the measured 8.5 GB phantom.
15555        assert_eq!(d(1, 256, 33 * MIB, None), 33 * MIB);
15556
15557        // MARGINAL, flat pool (the box9 receipt: reserved stayed ~33.6 MiB across captures
15558        // 1..3, so an additional key costs ~nothing and the debt must collapse to ~0 —
15559        // NOT the 8,556/4,261/2,830 MB the mean rule printed).
15560        assert_eq!(d(3, 256, 33 * MIB, Some((1, 33 * MIB))), 0);
15561
15562        // MARGINAL, genuinely growing pool: 40 MiB per new key over 2 keys, 250 slots left.
15563        let debt = d(6, 256, 273 * MIB, Some((4, 193 * MIB)));
15564        assert_eq!(debt, 250 * (40 * MIB));
15565        assert!(
15566            debt > 3 * (1536 * MIB),
15567            "real growth must dwarf SPEC_SHRINK_RESERVE"
15568        );
15569
15570        // a shrinking/recycled reading never becomes a negative charge.
15571        assert_eq!(d(6, 256, 10 * MIB, Some((4, 99 * MIB))), 0);
15572        // a stale observation at the same capture count falls back to bootstrap.
15573        assert_eq!(d(4, 256, 80 * MIB, Some((4, 80 * MIB))), 80 * MIB);
15574    }
15575}
15576
15577#[cfg(test)]
15578mod mtp_chain_tests {
15579    use super::mtp_chain_head_index;
15580
15581    #[test]
15582    fn embedded_step_heads_cycle_in_declared_order() {
15583        let actual: Vec<usize> = (0..8).map(|step| mtp_chain_head_index(step, 3)).collect();
15584        assert_eq!(actual, [0, 1, 2, 0, 1, 2, 0, 1]);
15585    }
15586
15587    #[test]
15588    fn standalone_draft_remains_single_head() {
15589        assert!((0..8).all(|step| mtp_chain_head_index(step, 1) == 0));
15590    }
15591}
15592
15593#[cfg(test)]
15594mod tp_verified_prefix_tests {
15595    use super::rewind_tp_kv_verified_prefix;
15596    use crate::tp::ResidentTpKvCache;
15597
15598    fn cache_with_committed_len(committed: usize) -> ResidentTpKvCache {
15599        let mut cache = ResidentTpKvCache::new(Vec::new(), 1, 1, 1, 1, 8);
15600        let transaction = cache.begin_transaction().unwrap();
15601        let target = cache.append_target(transaction, committed).unwrap();
15602        cache.publish_append(transaction, target).unwrap();
15603        let target = cache.commit_target(transaction, committed).unwrap();
15604        cache.publish_finalize(transaction, target).unwrap();
15605        cache
15606    }
15607
15608    #[test]
15609    fn replay_free_prefix_rewinds_tp_visibility_to_snapshot_plus_accepts() {
15610        let mut layers = vec![Some(cache_with_committed_len(5)), None];
15611        rewind_tp_kv_verified_prefix(&mut layers, &[Some(2), None], 1).unwrap();
15612        let cache = layers[0].as_ref().unwrap();
15613        assert_eq!(cache.committed_len(), 3);
15614        assert_eq!(cache.staged_len(), 3);
15615    }
15616
15617    #[test]
15618    fn replay_free_prefix_rejects_a_changed_tp_cache_shape() {
15619        let mut layers = vec![Some(cache_with_committed_len(1))];
15620        let error = rewind_tp_kv_verified_prefix(&mut layers, &[None], 1)
15621            .unwrap_err()
15622            .to_string();
15623        assert!(error.contains("changed shape"), "unexpected error: {error}");
15624    }
15625}
15626
15627#[cfg(test)]
15628mod dspark_sparse_tests {
15629    use super::dspark_sparse_softmax_topk;
15630
15631    #[test]
15632    fn topk_keeps_full_softmax_mass_and_stable_ties() {
15633        let logits = [1.0f32, 3.0, 3.0, -2.0];
15634        let (ids, top_logits, probs, tail) = dspark_sparse_softmax_topk(&logits, 2, 1.0).unwrap();
15635        assert_eq!(ids, vec![1, 2]);
15636        assert_eq!(top_logits, vec![3.0, 3.0]);
15637        let denominator = logits.iter().map(|value| (value - 3.0).exp()).sum::<f32>();
15638        let expected = 1.0 / denominator;
15639        assert!((probs[0] - expected).abs() < 1.0e-6);
15640        assert!((probs[1] - expected).abs() < 1.0e-6);
15641        assert!((tail - (1.0 - 2.0 * expected)).abs() < 1.0e-6);
15642        assert!((probs.iter().sum::<f32>() + tail - 1.0).abs() < 1.0e-6);
15643    }
15644}
15645
15646#[cfg(test)]
15647mod spec_replay_env_tests {
15648    use super::spec_replay_env_on;
15649
15650    #[test]
15651    fn replay_requires_literal_one() {
15652        assert!(!spec_replay_env_on(None));
15653        assert!(!spec_replay_env_on(Some("")));
15654        assert!(!spec_replay_env_on(Some("0")));
15655        assert!(!spec_replay_env_on(Some("true")));
15656        assert!(!spec_replay_env_on(Some("2")));
15657        assert!(spec_replay_env_on(Some("1")));
15658    }
15659}
15660
15661#[cfg(test)]
15662mod telem_tests {
15663    use super::{SPEC_TELEM_POS, SpecTelemetry, SpecTelemetryCounters};
15664
15665    #[test]
15666    fn synthetic_accept_masks_produce_tau_and_position_histogram() {
15667        let counters = SpecTelemetryCounters::default();
15668        for mask in [
15669            [true, true, true],
15670            [true, true, false],
15671            [true, false, false],
15672            [false, false, false],
15673        ] {
15674            let accepted = mask.iter().take_while(|&&value| value).count();
15675            counters.record_round(mask.len(), accepted);
15676        }
15677
15678        let snapshot = counters.snapshot();
15679        assert_eq!(
15680            (snapshot.rounds, snapshot.drafted, snapshot.accepted),
15681            (4, 12, 6)
15682        );
15683        assert_eq!(&snapshot.pos_drafted[..3], &[4, 4, 4]);
15684        assert_eq!(&snapshot.pos_accepted[..3], &[3, 2, 1]);
15685        assert_eq!(snapshot.tau(), 1.5);
15686        assert_eq!(snapshot.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15687        assert_eq!(snapshot.pos_accepted[3..], [0; SPEC_TELEM_POS - 3]);
15688    }
15689
15690    /// The worker's per-burst pattern: stash, accumulate, diff — the delta must isolate
15691    /// exactly the burst's contribution (pool-resumed sessions carry prior requests' counts).
15692    #[test]
15693    fn delta_isolates_burst_contribution() {
15694        let mut t = SpecTelemetry::default();
15695        // "previous request": 2 rounds of k=3, accepts 3 then 1.
15696        for (kr, na) in [(3usize, 3usize), (3, 1)] {
15697            t.rounds += 1;
15698            t.drafted += kr as u64;
15699            t.accepted += na as u64;
15700            for j in 0..kr {
15701                t.pos_drafted[j] += 1;
15702            }
15703            for j in 0..na {
15704                t.pos_accepted[j] += 1;
15705            }
15706        }
15707        let before = t;
15708        // "this burst": 1 round k=3, accepts 2.
15709        t.rounds += 1;
15710        t.drafted += 3;
15711        t.accepted += 2;
15712        for j in 0..3 {
15713            t.pos_drafted[j] += 1;
15714        }
15715        for j in 0..2 {
15716            t.pos_accepted[j] += 1;
15717        }
15718        let d = t.delta_since(&before);
15719        assert_eq!((d.rounds, d.drafted, d.accepted), (1, 3, 2));
15720        assert_eq!(&d.pos_drafted[..3], &[1, 1, 1]);
15721        assert_eq!(&d.pos_accepted[..3], &[1, 1, 0]);
15722        assert_eq!(d.pos_drafted[3..], [0; SPEC_TELEM_POS - 3]);
15723    }
15724
15725    /// merge(delta) then merge(delta2) equals accumulating both — the per-model /metrics
15726    /// aggregation invariant.
15727    #[test]
15728    fn merge_accumulates_fieldwise() {
15729        let mut agg = SpecTelemetry::default();
15730        let mut d1 = SpecTelemetry {
15731            rounds: 2,
15732            drafted: 6,
15733            accepted: 4,
15734            ..Default::default()
15735        };
15736        d1.pos_drafted[0] = 2;
15737        d1.pos_accepted[0] = 2;
15738        let mut d2 = SpecTelemetry {
15739            rounds: 1,
15740            drafted: 3,
15741            accepted: 1,
15742            ..Default::default()
15743        };
15744        d2.pos_drafted[0] = 1;
15745        d2.pos_accepted[0] = 1;
15746        d2.pos_drafted[1] = 1;
15747        agg.merge(&d1);
15748        agg.merge(&d2);
15749        assert_eq!((agg.rounds, agg.drafted, agg.accepted), (3, 9, 5));
15750        assert_eq!(agg.pos_drafted[0], 3);
15751        assert_eq!(agg.pos_accepted[0], 3);
15752        assert_eq!(agg.pos_drafted[1], 1);
15753        assert_eq!(agg.pos_accepted[1], 0);
15754    }
15755
15756    /// Wrong-snapshot diff saturates to zero instead of wrapping — the counters feed a
15757    /// public metrics surface and must never publish a u64-wrapped garbage value.
15758    #[test]
15759    fn delta_saturates_never_wraps() {
15760        let small = SpecTelemetry {
15761            rounds: 1,
15762            drafted: 2,
15763            accepted: 1,
15764            ..Default::default()
15765        };
15766        let big = SpecTelemetry {
15767            rounds: 5,
15768            drafted: 15,
15769            accepted: 9,
15770            ..Default::default()
15771        };
15772        let d = small.delta_since(&big);
15773        assert_eq!((d.rounds, d.drafted, d.accepted), (0, 0, 0));
15774    }
15775}
15776
15777#[cfg(test)]
15778mod opti_fork_tests {
15779    use super::{
15780        OptiControllerPolicy, OptiForkAction, OptiForkGateMode, OptiForkGenerationTracker,
15781    };
15782
15783    #[test]
15784    fn controller_threshold_and_three_miss_breaker_are_exact() {
15785        let mut policy = OptiControllerPolicy {
15786            threshold: 0.7,
15787            consecutive_misses: 0,
15788            breaker_tripped: false,
15789        };
15790        assert!(!policy.admit(0.699_999));
15791        assert!(policy.admit(0.7));
15792        assert!(!policy.resolve(false));
15793        assert!(!policy.resolve(false));
15794        assert!(policy.resolve(false));
15795        assert!(policy.breaker_tripped);
15796        assert!(!policy.admit(1.0));
15797        assert!(
15798            !policy.resolve(true),
15799            "a resolved hit cannot re-arm a tripped request"
15800        );
15801        assert!(policy.breaker_tripped);
15802    }
15803
15804    #[test]
15805    fn zero_threshold_is_the_true_unconditional_measurement_arm() {
15806        let mut policy = OptiControllerPolicy {
15807            threshold: 0.0,
15808            consecutive_misses: 0,
15809            breaker_tripped: false,
15810        };
15811        for _ in 0..16 {
15812            assert!(policy.admit(0.0));
15813            assert!(!policy.resolve(false));
15814        }
15815        for invalid in [f32::NAN, f32::INFINITY, -0.01, 1.01] {
15816            assert!(
15817                !policy.admit(invalid),
15818                "invalid q proxy must fail closed: {invalid}"
15819            );
15820        }
15821        assert!(!policy.breaker_tripped);
15822        assert_eq!(policy.consecutive_misses, 0);
15823    }
15824
15825    #[test]
15826    fn alternating_mode_flips_by_generation_not_round_parity() {
15827        assert_eq!(OptiForkGateMode::Alternate.action(0), OptiForkAction::Hit);
15828        assert_eq!(OptiForkGateMode::Alternate.action(1), OptiForkAction::Miss);
15829        assert_eq!(OptiForkGateMode::Alternate.action(8), OptiForkAction::Hit);
15830        assert_eq!(OptiForkGateMode::Alternate.action(9), OptiForkAction::Miss);
15831    }
15832
15833    #[test]
15834    fn live_generation_cannot_be_overwritten() {
15835        let mut tracker = OptiForkGenerationTracker::default();
15836        let g0 = tracker.reserve().unwrap();
15837        let g1 = tracker.reserve().unwrap();
15838        let err = tracker.reserve().unwrap_err().to_string();
15839        assert!(
15840            err.contains("still owns generation 0"),
15841            "unexpected error: {err}"
15842        );
15843        tracker.retire(g0).unwrap();
15844        let g2 = tracker.reserve().unwrap();
15845        assert_eq!((g2.id, g2.slot), (2, 0));
15846        tracker.retire(g1).unwrap();
15847        tracker.retire(g2).unwrap();
15848    }
15849
15850    #[test]
15851    fn teardown_rejects_a_stale_generation_tag() {
15852        let mut tracker = OptiForkGenerationTracker::default();
15853        let g0 = tracker.reserve().unwrap();
15854        tracker.retire(g0).unwrap();
15855        let err = tracker.retire(g0).unwrap_err().to_string();
15856        assert!(err.contains("teardown mismatch"), "unexpected error: {err}");
15857    }
15858}
15859
15860#[cfg(test)]
15861mod draft_graph_fallback_tests {
15862    use super::DraftGraphFallback;
15863
15864    /// The Q2 contract, part (a): a fallback flip is LOUD — exactly once per flip.
15865    #[test]
15866    fn flip_is_loud_once_and_memoized_after() {
15867        let mut f = DraftGraphFallback::default();
15868        let line = f
15869            .mark_greedy("out of memory")
15870            .expect("first flip must return the warn line");
15871        assert!(
15872            line.contains("WARN"),
15873            "flip line must be warn-level: {line}"
15874        );
15875        assert!(
15876            line.contains("out of memory"),
15877            "flip line must carry the reason: {line}"
15878        );
15879        assert!(f.greedy_failed());
15880        // re-marking an already-failed graph is the memoization: quiet, still failed.
15881        assert!(f.mark_greedy("out of memory").is_none());
15882        assert!(f.greedy_failed());
15883        // the two graphs' flags are independent (greedy flip leaves sampled capturable).
15884        assert!(!f.sampled_failed());
15885        let line_s = f
15886            .mark_sampled("capture unsupported")
15887            .expect("sampled flip is its own flip");
15888        assert!(
15889            line_s.contains("sampled"),
15890            "sampled flip names itself: {line_s}"
15891        );
15892        assert!(f.mark_sampled("capture unsupported").is_none());
15893    }
15894
15895    /// The Q2 contract, part (b): resume-from-pool RESETS both flags (fresh capture chance),
15896    /// and says so exactly when there was something to reset.
15897    #[test]
15898    fn reset_on_resume_clears_flags_and_logs_once() {
15899        let mut f = DraftGraphFallback::default();
15900        // clean session: resume is silent, nothing to reset.
15901        assert!(f.reset_on_resume().is_none());
15902        f.mark_greedy("oom").unwrap();
15903        f.mark_sampled("oom").unwrap();
15904        let note = f
15905            .reset_on_resume()
15906            .expect("a set flag must produce the reset note");
15907        assert!(
15908            note.contains("greedy+sampled"),
15909            "note names what was reset: {note}"
15910        );
15911        assert!(
15912            !f.greedy_failed() && !f.sampled_failed(),
15913            "both flags cleared"
15914        );
15915        // and the NEXT failure after a reset is a fresh flip — loud again.
15916        assert!(f.mark_greedy("oom again").is_some());
15917        let note2 = f.reset_on_resume().expect("greedy-only reset");
15918        assert!(note2.contains("(greedy)"), "single-flag note: {note2}");
15919    }
15920
15921    /// Shape-change clears (dmask realloc / mask-shape mismatch / s_key change) stay silent —
15922    /// they precede a fresh capture attempt whose own failure re-flips loudly.
15923    #[test]
15924    fn shape_change_clears_are_silent() {
15925        let mut f = DraftGraphFallback::default();
15926        f.mark_greedy("oom").unwrap();
15927        f.clear_greedy();
15928        assert!(!f.greedy_failed());
15929        f.mark_sampled("oom").unwrap();
15930        f.clear_sampled();
15931        assert!(!f.sampled_failed());
15932        // after a silent clear there is nothing left for resume to report.
15933        assert!(f.reset_on_resume().is_none());
15934    }
15935}
15936
15937/// SAMPLED DRAFT-GRAPH KEY (lane/graph-s-key-exactness-20260819).
15938///
15939/// These are the CPU teeth for an exactness bug whose live reproduction needs a GPU, a trunk, a
15940/// drafter and a two-turn session: the key itself. Every test below fails against the pre-fix key
15941/// `(seed, temp.to_bits(), k)` — `legacy_key` restates it so the collision is explicit rather
15942/// than remembered.
15943#[cfg(test)]
15944mod sampled_graph_key_tests {
15945    use super::{SampledGraphKey, debug_t_pred0};
15946
15947    /// The pre-fix key, verbatim: `let s_key = (sp_seed, sp_temp.to_bits(), k);`
15948    fn legacy_key(k: &SampledGraphKey) -> (u64, u32, usize) {
15949        (k.seed, k.temp_bits, k.k)
15950    }
15951
15952    fn pure_temp_key() -> SampledGraphKey {
15953        // temperature 1.0, filters off — today's serve default, the shape that parks a graph.
15954        SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, false)
15955    }
15956
15957    /// THE COLLISION. Two requests that differ ONLY in the truncation filters shared one key, so
15958    /// a parked pure-temp graph survived into a filtered request and the launch site launched it.
15959    #[test]
15960    fn vendor_filters_change_the_key() {
15961        let parked = pure_temp_key();
15962        // qwen3.8 generation_config.json — what the vendor-default flip makes the default shape.
15963        let vendor = SampledGraphKey::new(12345, 1.0, 3, 20, 0.95, 0.0, false);
15964        assert_eq!(
15965            legacy_key(&parked),
15966            legacy_key(&vendor),
15967            "pre-fix key collided: this is the bug, and the reason a test asserts on it",
15968        );
15969        assert_ne!(parked, vendor, "post-fix key must separate the two regimes");
15970        assert!(parked.pure_temp());
15971        assert!(!vendor.pure_temp());
15972    }
15973
15974    /// Each distribution-shaping field alone is enough to drop the parked graph.
15975    #[test]
15976    fn every_filter_field_is_keyed() {
15977        let base = pure_temp_key();
15978        for (what, other) in [
15979            (
15980                "top_k",
15981                SampledGraphKey::new(12345, 1.0, 3, 20, 1.0, 0.0, false),
15982            ),
15983            (
15984                "top_p",
15985                SampledGraphKey::new(12345, 1.0, 3, 0, 0.95, 0.0, false),
15986            ),
15987            (
15988                "min_p",
15989                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.05, false),
15990            ),
15991            (
15992                "penalties",
15993                SampledGraphKey::new(12345, 1.0, 3, 0, 1.0, 0.0, true),
15994            ),
15995        ] {
15996            assert_ne!(base, other, "{what} must be part of the key");
15997            assert!(!other.pure_temp(), "{what} leaves the pure-temp regime");
15998            assert_eq!(
15999                legacy_key(&base),
16000                legacy_key(&other),
16001                "{what} was invisible to the pre-fix key",
16002            );
16003        }
16004    }
16005
16006    /// The baked constants stay keyed (this half was always right — regression cover for it).
16007    #[test]
16008    fn baked_constants_stay_keyed() {
16009        let base = pure_temp_key();
16010        assert_ne!(
16011            base,
16012            SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false),
16013            "seed"
16014        );
16015        assert_ne!(
16016            base,
16017            SampledGraphKey::new(12345, 0.7, 3, 0, 1.0, 0.0, false),
16018            "temp"
16019        );
16020        assert_ne!(
16021            base,
16022            SampledGraphKey::new(12345, 1.0, 4, 0, 1.0, 0.0, false),
16023            "k"
16024        );
16025        // bitwise on temperature: 0.7f32 vs the same value re-derived must NOT differ.
16026        assert_eq!(
16027            SampledGraphKey::new(1, 0.7, 3, 0, 1.0, 0.0, false),
16028            SampledGraphKey::new(1, 7.0 / 10.0, 3, 0, 1.0, 0.0, false),
16029        );
16030    }
16031
16032    /// THE LOAD-BEARING HALF OF THE SEED DECISION (lane/session-resume-sampler-predicate-
16033    /// 20260820). The whole-session resume predicate deliberately does NOT compare `seed`: an
16034    /// omitted serve `seed` draws fresh per-request entropy, so comparing it would refuse every
16035    /// seed-omitting sampled conversation. That is only sound because the one piece of parked state
16036    /// that BAKES the seed — this graph — is re-keyed on it, so a seed change drops and recaptures.
16037    ///
16038    /// This test is the other end of that argument, asserted here rather than remembered in a
16039    /// comment: if a future change dropped `seed` from the key, the resume predicate's exclusion
16040    /// would silently become the unsound thing it is documented not to be.
16041    /// (Paired with `seed_alone_does_not_refuse` in `memra-sampling`.)
16042    #[test]
16043    fn seed_alone_still_rekeys_the_draft_graph() {
16044        let parked = pure_temp_key();
16045        let reseeded = SampledGraphKey::new(999, 1.0, 3, 0, 1.0, 0.0, false);
16046        assert_ne!(
16047            parked, reseeded,
16048            "a seed-only change MUST drop the parked sampled graph — the resume predicate's \
16049             decision not to compare seed rests on exactly this",
16050        );
16051        // Same regime on both sides: the drop is a recapture, not a fall to the eager chain
16052        // because of a filter difference.
16053        assert!(parked.pure_temp() && reseeded.pure_temp());
16054    }
16055
16056    /// `pure_temp()` is the capture guard's predicate, computed from the key so the two cannot
16057    /// drift. The equality below is the invariant the launch-site guard asserts: identical keys
16058    /// agree on the regime, so a graph that survives the drop is legal to launch.
16059    #[test]
16060    fn equal_keys_agree_on_the_regime() {
16061        let a = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16062        let b = SampledGraphKey::new(7, 0.8, 3, 20, 0.95, 0.0, false);
16063        assert_eq!(a, b);
16064        assert_eq!(a.pure_temp(), b.pure_temp());
16065        // top_p slightly above 1.0 (a client sending 1.0 exactly, or an operator default) is
16066        // still the unfiltered regime, matching the original `sp.top_p >= 1.0` test.
16067        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.0, 0.0, false).pure_temp());
16068        assert!(SampledGraphKey::new(7, 0.8, 3, 0, 1.5, -1.0, false).pure_temp());
16069    }
16070
16071    /// The WIDENED capture regime (lane/step37-draft-graph-serving-20260830): truncation-
16072    /// filtered shapes are capturable — the filter runs IN-GRAPH (`filter_stats` +
16073    /// `gumbel_perturb_filtered_ctr`), so the draft draws from the same filtered
16074    /// distribution the accept test reconstructs. Penalties never are: the per-round
16075    /// history cannot be baked. The step37 vendor-default shape (temp 0.5 / top_p 0.9) is
16076    /// exactly the previously-excluded regime this lane exists to capture.
16077    #[test]
16078    fn filtered_regimes_are_capturable_penalties_never() {
16079        let vendor = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, false);
16080        assert!(!vendor.pure_temp());
16081        assert!(vendor.filtered());
16082        assert!(
16083            vendor.graph_capturable(),
16084            "the vendor-default filtered shape must be capturable (default door state)",
16085        );
16086        assert!(pure_temp_key().graph_capturable());
16087        assert!(
16088            !pure_temp_key().filtered(),
16089            "pure-temp takes the legacy (filterless) capture body",
16090        );
16091        let pen = SampledGraphKey::new(12345, 0.5, 3, 0, 0.9, 0.0, true);
16092        assert!(
16093            !pen.graph_capturable(),
16094            "penalty history varies per round and can never be baked into a graph",
16095        );
16096    }
16097
16098    /// MEMRA_DEBUG_SPEC on a SAMPLED spec request past round 0: the print must render without
16099    /// indexing the empty greedy `preds` vector (it panicked the GPU worker before this lane).
16100    #[test]
16101    fn debug_print_survives_the_sampled_arm() {
16102        // round >= 1 with a pending bonus == base 1, sampled == `preds` empty.
16103        assert_eq!(debug_t_pred0(true, 1, 4242, &[]), "n/a");
16104        assert_eq!(debug_t_pred0(true, 2, 4242, &[]), "n/a");
16105        // round 0 without a pending bonus still reports last_pred, in both arms.
16106        assert_eq!(debug_t_pred0(true, 0, 4242, &[]), "4242");
16107        assert_eq!(debug_t_pred0(false, 0, 4242, &[7, 8]), "4242");
16108        // greedy keeps the real prediction it always printed.
16109        assert_eq!(debug_t_pred0(false, 1, 4242, &[7, 8]), "7");
16110        assert_eq!(debug_t_pred0(false, 2, 4242, &[7, 8]), "8");
16111    }
16112}